fix: bound frame payload allocations

docs/ARCHITECTURE.md

@@ -96,6 +96,11 @@ client.PublishFrame(compose.Frame{ Pixels: rgba, Width: 400, Height: 120 })
 | 56 | PayloadSize | 4B | Compressed payload bytes |
 | 60 | UncompressedSize | 4B | Original pixel bytes |
 
+The wire fields remain `uint32` for protocol compatibility. The socket
+transport and decompression boundary enforce `protocol.MaxPayloadSize` (64
+MiB) before converting a declared size or allocating a payload buffer. Header
+decoding itself remains permissive and accepts the full `uint32` range.
+
 ## Frame Delivery (ADR-002)
 
 Both push and pull delivery coexist. No mode negotiation — inferred from behavior (Chromium pattern).
@@ -163,7 +168,7 @@ Module connects → Handshake (name, size, fps)
 ```
 compose (root) ──→ internal/protocol (leaf, no deps)
     ├──→ internal/transport/socket ──→ internal/protocol
-    ├──→ internal/codec (standalone)
+    ├──→ internal/codec ──→ internal/protocol
     ├──→ internal/flow (standalone)
     └──→ internal/conn (standalone)
 ```

internal/codec/doc.go

@@ -10,6 +10,9 @@
 //   - Raw (ID 0x00): Pass-through copy, no compression.
 //   - LZ4 (ID 0x01): LZ4 block compression via github.com/pierrec/lz4/v4.
 //
+// LZ4's allocation fallback is bounded by protocol.MaxPayloadSize (64 MiB),
+// matching the transport and server decompression boundary.
+//
 // Registration happens automatically via init() in each codec's source file.
 // Use Get(id) to retrieve a codec by its protocol identifier.
 package codec

internal/codec/lz4.go

@@ -5,6 +5,8 @@ import (
 	"sync"
 
 	"github.com/pierrec/lz4/v4"
+
+	"github.com/gogpu/compose/internal/protocol"
 )
 
 func init() {
@@ -32,6 +34,17 @@ type lz4Codec struct {
 	pool sync.Pool
 }
 
+// maxDecodeBuf preserves the codec's existing internal name while sourcing
+// the bound from the shared protocol allocation policy.
+const maxDecodeBuf = protocol.MaxPayloadSize
+
+// maxDecodeRatio bounds the nil-destination fallback for a compressed block.
+// LZ4's block format encodes matches with a 16-bit offset and bounded length
+// extensions, so a few hundred times expansion is the practical upper range.
+// Keeping the fallback proportional to the input also prevents tiny malformed
+// blocks from forcing the global 64 MiB ceiling.
+const maxDecodeRatio = 256
+
 // Encode compresses src using LZ4 block compression. Returns a sub-slice of
 // dst containing the compressed data. If dst is nil or too small, allocates a
 // new buffer.
@@ -81,22 +94,25 @@ func (c *lz4Codec) Decode(dst, src []byte) ([]byte, error) {
 	if cap(dst) == 0 {
 		// Caller didn't provide a buffer. Start with 10x compressed size
 		// as initial guess. LZ4 GUI data often compresses 100:1 or better,
-		// but 10x covers most cases in one attempt.
-		dst = make([]byte, len(src)*10)
+		// but 10x covers most cases in one attempt. Clamp before multiplying:
+		// malformed input must not overflow int or allocate beyond the growth
+		// limit before the limit below gets a chance to run.
+		dst = make([]byte, initialDecodeSize(len(src)))
 	} else {
 		dst = dst[:cap(dst)]
 	}
-
-	// maxDecodeBuf caps the growth strategy to prevent runaway allocation
-	// on corrupt or adversarial input (64 MB covers 4K RGBA frames).
-	const maxDecodeBuf = 64 * 1024 * 1024
+	maxSize := maxDecodeSize(len(src))
 
 	for {
 		n, err := lz4.UncompressBlock(src, dst)
 		if err != nil {
 			// If buffer might be too small, double and retry.
-			if len(dst) < maxDecodeBuf {
-				dst = make([]byte, len(dst)*2)
+			if len(dst) < maxSize {
+				next := len(dst) * 2
+				if next <= len(dst) || next > maxSize {
+					next = maxSize
+				}
+				dst = make([]byte, next)
 				continue
 			}
 			return nil, fmt.Errorf("codec: lz4 decode: %w", err)
@@ -105,6 +121,26 @@ func (c *lz4Codec) Decode(dst, src []byte) ([]byte, error) {
 	}
 }
 
+// initialDecodeSize returns the first allocation used by the nil-destination
+// fallback without allowing len(src)*10 to overflow or exceed the shared
+// payload bound.
+func initialDecodeSize(srcLen int) int {
+	if srcLen > maxDecodeBuf/10 {
+		return maxDecodeSize(srcLen)
+	}
+	return srcLen * 10
+}
+
+// maxDecodeSize returns the largest fallback destination permitted for a
+// compressed block of srcLen bytes. It is overflow-safe and never exceeds the
+// shared protocol allocation bound.
+func maxDecodeSize(srcLen int) int {
+	if srcLen >= maxDecodeBuf/maxDecodeRatio {
+		return maxDecodeBuf
+	}
+	return srcLen * maxDecodeRatio
+}
+
 // ID returns the LZ4 codec protocol identifier (0x01).
 func (c *lz4Codec) ID() byte {
 	return IDLZ4

internal/codec/lz4_test.go

@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"crypto/rand"
 	"testing"
+
+	"github.com/gogpu/compose/internal/protocol"
 )
 
 func TestLZ4RoundTrip(t *testing.T) {
@@ -85,6 +87,35 @@ func TestLZ4EmptyInput(t *testing.T) {
 	}
 }
 
+func TestLZ4InitialDecodeSizeIsBounded(t *testing.T) {
+	for _, srcLen := range []int{
+		protocol.MaxPayloadSize/10 + 1,
+		int(^uint(0) >> 1),
+	} {
+		if got := initialDecodeSize(srcLen); got != protocol.MaxPayloadSize {
+			t.Errorf("initialDecodeSize(%d) = %d, want %d", srcLen, got, protocol.MaxPayloadSize)
+		}
+	}
+
+	if got := initialDecodeSize(1024); got != 10240 {
+		t.Errorf("initialDecodeSize(1024) = %d, want 10240", got)
+	}
+}
+
+func TestLZ4DecodeFallbackScalesWithInput(t *testing.T) {
+	if got := maxDecodeSize(1); got != maxDecodeRatio {
+		t.Errorf("maxDecodeSize(1) = %d, want %d", got, maxDecodeRatio)
+	}
+	if got := maxDecodeSize(maxDecodeBuf / maxDecodeRatio); got != maxDecodeBuf {
+		t.Errorf("maxDecodeSize(threshold) = %d, want %d", got, maxDecodeBuf)
+	}
+
+	// A malformed one-byte block must fail without walking the global cap.
+	if _, err := LZ4().Decode(nil, []byte{0xFF}); err == nil {
+		t.Fatal("Decode accepted malformed one-byte block")
+	}
+}
+
 func TestLZ4CompressionRatio(t *testing.T) {
 	c := LZ4()
 

internal/protocol/header_test.go

@@ -306,6 +306,29 @@ func TestDecode_ExactBuffer(t *testing.T) {
 	}
 }
 
+func TestDecode_AcceptsMaximumWirePayloadSizes(t *testing.T) {
+	h := Header{
+		Magic:            Magic,
+		Version:          ProtocolVersion,
+		MsgType:          MsgFrame,
+		PayloadSize:      math.MaxUint32,
+		UncompressedSize: math.MaxUint32,
+	}
+	buf := make([]byte, HeaderSize)
+	if err := Encode(&h, buf); err != nil {
+		t.Fatalf("Encode: %v", err)
+	}
+
+	got, err := Decode(buf)
+	if err != nil {
+		t.Fatalf("Decode: %v", err)
+	}
+	if got.PayloadSize != math.MaxUint32 || got.UncompressedSize != math.MaxUint32 {
+		t.Fatalf("decoded sizes = (%d, %d), want (%d, %d)",
+			got.PayloadSize, got.UncompressedSize, uint32(math.MaxUint32), uint32(math.MaxUint32))
+	}
+}
+
 func TestHeader_ReservedZeroed(t *testing.T) {
 	// After encoding, reserved bytes should be zero.
 	h := Header{

internal/protocol/limits.go

@@ -0,0 +1,10 @@
+package protocol
+
+// MaxPayloadSize is the maximum payload size accepted by the socket
+// transport and decompression boundary.
+//
+// Header payload fields remain uint32 so Decode continues to parse every
+// representable wire value. The limit is enforced only before a transport or
+// codec allocation, keeping malformed frames from forcing unbounded memory
+// use while allowing the protocol representation to remain forward-compatible.
+const MaxPayloadSize = 64 * 1024 * 1024

internal/transport/socket/conn.go

@@ -85,6 +85,14 @@ func (c *Conn) ReadFrameInto(buf []byte) (protocol.Header, []byte, error) {
 		return protocol.Header{}, nil, fmt.Errorf("socket: decode header: %w", err)
 	}
 
+	// Validate the wire-sized field before converting it to int or allocating.
+	// Header.Decode intentionally accepts the complete uint32 field range for
+	// protocol compatibility; this boundary is where memory use is bounded.
+	if hdr.PayloadSize > protocol.MaxPayloadSize {
+		return protocol.Header{}, nil, fmt.Errorf("%w: declared %d bytes (limit %d)",
+			ErrPayloadTooLarge, hdr.PayloadSize, protocol.MaxPayloadSize)
+	}
+
 	// Read payload.
 	size := int(hdr.PayloadSize)
 	if size == 0 {

internal/transport/socket/conn_test.go

@@ -100,6 +100,35 @@ func TestWriteReadFrame_RoundTrip(t *testing.T) {
 	}
 }
 
+func TestReadFrame_OversizedPayloadHeader(t *testing.T) {
+	client, server := net.Pipe()
+	defer client.Close()
+	defer server.Close()
+
+	hdr := makeHeader(protocol.MsgFrame, uint32(protocol.MaxPayloadSize+1))
+	var encoded [protocol.HeaderSize]byte
+	if err := protocol.Encode(&hdr, encoded[:]); err != nil {
+		t.Fatalf("Encode: %v", err)
+	}
+
+	writeErr := make(chan error, 1)
+	go func() {
+		_, err := client.Write(encoded[:])
+		writeErr <- err
+	}()
+
+	_, _, err := NewConn(server).ReadFrame()
+	if !errors.Is(err, ErrPayloadTooLarge) {
+		t.Fatalf("ReadFrame error = %v, want ErrPayloadTooLarge", err)
+	}
+
+	// The reader must reject the header before attempting to read the
+	// declared payload. A header-only peer therefore completes immediately.
+	if err := <-writeErr; err != nil {
+		t.Fatalf("header write: %v", err)
+	}
+}
+
 func TestWriteReadFrame_LargePayload(t *testing.T) {
 	client, server := newPipeConns(t)
 

internal/transport/socket/doc.go

@@ -7,8 +7,9 @@
 //
 // [Conn] wraps a [net.Conn] with framed I/O. Each frame on the wire is a
 // 64-byte header (see [protocol.HeaderSize]) followed by PayloadSize bytes
-// of payload. Reads and writes are independently locked, so concurrent
-// producers and consumers are safe on the same connection.
+// of payload. Incoming payloads larger than [protocol.MaxPayloadSize] are
+// rejected before allocation. Reads and writes are independently locked, so
+// concurrent producers and consumers are safe on the same connection.
 //
 // # Listener
 //

internal/transport/socket/errors.go

@@ -0,0 +1,7 @@
+package socket
+
+import "errors"
+
+// ErrPayloadTooLarge reports a frame whose declared payload exceeds the
+// transport's allocation limit.
+var ErrPayloadTooLarge = errors.New("socket: payload too large")

server.go

@@ -424,6 +424,18 @@ func (s *Server) decodePayload(hdr protocol.Header, payload []byte) ([]byte, err
 		return payload, nil
 	}
 
+	// Header.Decode keeps the uint32 wire representation permissive, but do
+	// not turn an untrusted declared size into an allocation at this boundary.
+	// A non-empty compressed frame must declare its original size; otherwise
+	// the codec would have to use its allocation fallback.
+	if hdr.UncompressedSize > protocol.MaxPayloadSize {
+		return nil, fmt.Errorf("compose: uncompressed payload size %d exceeds limit %d",
+			hdr.UncompressedSize, protocol.MaxPayloadSize)
+	}
+	if len(payload) > 0 && hdr.UncompressedSize == 0 {
+		return nil, fmt.Errorf("compose: compressed payload has zero uncompressed size")
+	}
+
 	// Look up the codec by compression ID.
 	c := codec.Get(byte(hdr.Compression))
 	if c == nil {

server_security_test.go

@@ -0,0 +1,54 @@
+package compose
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/gogpu/compose/internal/codec"
+	"github.com/gogpu/compose/internal/protocol"
+)
+
+func TestDecodePayloadRejectsOversizedUncompressedSize(t *testing.T) {
+	hdr := protocol.Header{
+		Flags:            protocol.FlagCompressed,
+		Compression:      protocol.CompressionLZ4,
+		UncompressedSize: ^uint32(0),
+	}
+
+	_, err := (&Server{}).decodePayload(hdr, []byte("compressed"))
+	if err == nil {
+		t.Fatal("decodePayload accepted an oversized uncompressed size")
+	}
+}
+
+func TestDecodePayloadLZ4RoundTrip(t *testing.T) {
+	src := bytes.Repeat([]byte{0xA5}, 4096)
+	c := codec.LZ4()
+	encoded, err := c.Encode(nil, src)
+	if err != nil {
+		t.Fatalf("Encode: %v", err)
+	}
+
+	hdr := protocol.Header{
+		Flags:            protocol.FlagCompressed,
+		Compression:      protocol.CompressionLZ4,
+		UncompressedSize: uint32(len(src)),
+	}
+	decoded, err := (&Server{}).decodePayload(hdr, encoded)
+	if err != nil {
+		t.Fatalf("decodePayload: %v", err)
+	}
+	if !bytes.Equal(decoded, src) {
+		t.Fatalf("decoded payload differs: got %d bytes, want %d", len(decoded), len(src))
+	}
+}
+
+func TestDecodePayloadRejectsZeroDeclaredSize(t *testing.T) {
+	hdr := protocol.Header{
+		Flags:       protocol.FlagCompressed,
+		Compression: protocol.CompressionLZ4,
+	}
+	if _, err := (&Server{}).decodePayload(hdr, []byte{0xFF}); err == nil {
+		t.Fatal("decodePayload accepted a non-empty compressed payload with zero declared size")
+	}
+}