From 06bfcdb6b23855575dc5394213027cad6c6dcfa3 Mon Sep 17 00:00:00 2001 From: Mark Date: Mon, 10 Aug 2026 16:25:22 +0300 Subject: [PATCH 1/4] fix: bound frame payload allocations --- docs/ARCHITECTURE.md | 7 ++- internal/codec/doc.go | 3 ++ internal/codec/lz4.go | 32 +++++++++--- internal/codec/lz4_test.go | 17 +++++++ internal/protocol/header_test.go | 23 +++++++++ internal/protocol/limits.go | 10 ++++ internal/transport/socket/conn.go | 8 +++ internal/transport/socket/conn_test.go | 29 +++++++++++ internal/transport/socket/doc.go | 5 +- internal/transport/socket/errors.go | 7 +++ server.go | 11 +++++ server_security_test.go | 67 ++++++++++++++++++++++++++ 12 files changed, 209 insertions(+), 10 deletions(-) create mode 100644 internal/protocol/limits.go create mode 100644 internal/transport/socket/errors.go create mode 100644 server_security_test.go diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 1e44448..dc1805f 100644 --- a/docs/ARCHITECTURE.md +++ b/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) ``` diff --git a/internal/codec/doc.go b/internal/codec/doc.go index 4d47bc2..007ec49 100644 --- a/internal/codec/doc.go +++ b/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 diff --git a/internal/codec/lz4.go b/internal/codec/lz4.go index 5125732..b6f7879 100644 --- a/internal/codec/lz4.go +++ b/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,10 @@ 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 + // 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 +87,24 @@ 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 - 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) + next := len(dst) * 2 + if next <= len(dst) || next > maxDecodeBuf { + next = maxDecodeBuf + } + dst = make([]byte, next) continue } return nil, fmt.Errorf("codec: lz4 decode: %w", err) @@ -105,6 +113,16 @@ 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 maxDecodeBuf + } + return srcLen * 10 +} + // ID returns the LZ4 codec protocol identifier (0x01). func (c *lz4Codec) ID() byte { return IDLZ4 diff --git a/internal/codec/lz4_test.go b/internal/codec/lz4_test.go index 417891f..b6d43d9 100644 --- a/internal/codec/lz4_test.go +++ b/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,21 @@ 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 TestLZ4CompressionRatio(t *testing.T) { c := LZ4() diff --git a/internal/protocol/header_test.go b/internal/protocol/header_test.go index 5a029e7..78a976e 100644 --- a/internal/protocol/header_test.go +++ b/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{ diff --git a/internal/protocol/limits.go b/internal/protocol/limits.go new file mode 100644 index 0000000..88bd8cf --- /dev/null +++ b/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 diff --git a/internal/transport/socket/conn.go b/internal/transport/socket/conn.go index f0445e8..72e3670 100644 --- a/internal/transport/socket/conn.go +++ b/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 { diff --git a/internal/transport/socket/conn_test.go b/internal/transport/socket/conn_test.go index 0843b7a..fa8650e 100644 --- a/internal/transport/socket/conn_test.go +++ b/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) diff --git a/internal/transport/socket/doc.go b/internal/transport/socket/doc.go index 1b563a6..e327ea6 100644 --- a/internal/transport/socket/doc.go +++ b/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 // diff --git a/internal/transport/socket/errors.go b/internal/transport/socket/errors.go new file mode 100644 index 0000000..42731a0 --- /dev/null +++ b/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") diff --git a/server.go b/server.go index 29efca4..f444828 100644 --- a/server.go +++ b/server.go @@ -424,6 +424,17 @@ 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. + if hdr.UncompressedSize > protocol.MaxPayloadSize { + return nil, fmt.Errorf("compose: uncompressed payload size %d exceeds limit %d", + hdr.UncompressedSize, protocol.MaxPayloadSize) + } + if len(payload) > protocol.MaxPayloadSize { + return nil, fmt.Errorf("compose: compressed payload size %d exceeds limit %d", + len(payload), protocol.MaxPayloadSize) + } + // Look up the codec by compression ID. c := codec.Get(byte(hdr.Compression)) if c == nil { diff --git a/server_security_test.go b/server_security_test.go new file mode 100644 index 0000000..bc3482b --- /dev/null +++ b/server_security_test.go @@ -0,0 +1,67 @@ +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 TestDecodePayloadLZ4ZeroDeclaredSize(t *testing.T) { + src := bytes.Repeat([]byte{0x5A}, 4096) + c := codec.LZ4() + encoded, err := c.Encode(nil, src) + if err != nil { + t.Fatalf("Encode: %v", err) + } + + // A zero declaration exercises the codec's nil-destination fallback. It + // remains valid for callers that do not know the decoded size in advance. + hdr := protocol.Header{ + Flags: protocol.FlagCompressed, + Compression: protocol.CompressionLZ4, + } + decoded, err := (&Server{}).decodePayload(hdr, encoded) + if err != nil { + t.Fatalf("decodePayload with zero declared size: %v", err) + } + if !bytes.Equal(decoded, src) { + t.Fatalf("decoded payload differs: got %d bytes, want %d", len(decoded), len(src)) + } +} From 8d05bf1a0f4c75e1a1bb33985e10c938bb0f426a Mon Sep 17 00:00:00 2001 From: Mark Date: Mon, 10 Aug 2026 16:41:15 +0300 Subject: [PATCH 2/4] fix: keep decompression growth proportional --- internal/codec/lz4.go | 31 +++++++++++++++++++++++++++---- internal/codec/lz4_test.go | 14 ++++++++++++++ server.go | 5 +++++ server_security_test.go | 19 +++---------------- 4 files changed, 49 insertions(+), 20 deletions(-) diff --git a/internal/codec/lz4.go b/internal/codec/lz4.go index b6f7879..d6f90bf 100644 --- a/internal/codec/lz4.go +++ b/internal/codec/lz4.go @@ -38,6 +38,13 @@ type lz4Codec struct { // 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. @@ -94,15 +101,16 @@ func (c *lz4Codec) Decode(dst, src []byte) ([]byte, error) { } else { dst = dst[:cap(dst)] } + 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 { + if len(dst) < maxSize { next := len(dst) * 2 - if next <= len(dst) || next > maxDecodeBuf { - next = maxDecodeBuf + if next <= len(dst) || next > maxSize { + next = maxSize } dst = make([]byte, next) continue @@ -117,10 +125,25 @@ func (c *lz4Codec) Decode(dst, src []byte) ([]byte, error) { // fallback without allowing len(src)*10 to overflow or exceed the shared // payload bound. func initialDecodeSize(srcLen int) int { + maxSize := maxDecodeSize(srcLen) if srcLen > maxDecodeBuf/10 { + return maxSize + } + initial := srcLen * 10 + if initial > maxSize { + return maxSize + } + return initial +} + +// 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 * 10 + return srcLen * maxDecodeRatio } // ID returns the LZ4 codec protocol identifier (0x01). diff --git a/internal/codec/lz4_test.go b/internal/codec/lz4_test.go index b6d43d9..e912f02 100644 --- a/internal/codec/lz4_test.go +++ b/internal/codec/lz4_test.go @@ -102,6 +102,20 @@ func TestLZ4InitialDecodeSizeIsBounded(t *testing.T) { } } +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() diff --git a/server.go b/server.go index f444828..886aa57 100644 --- a/server.go +++ b/server.go @@ -426,10 +426,15 @@ func (s *Server) decodePayload(hdr protocol.Header, payload []byte) ([]byte, err // 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") + } if len(payload) > protocol.MaxPayloadSize { return nil, fmt.Errorf("compose: compressed payload size %d exceeds limit %d", len(payload), protocol.MaxPayloadSize) diff --git a/server_security_test.go b/server_security_test.go index bc3482b..758ea47 100644 --- a/server_security_test.go +++ b/server_security_test.go @@ -43,25 +43,12 @@ func TestDecodePayloadLZ4RoundTrip(t *testing.T) { } } -func TestDecodePayloadLZ4ZeroDeclaredSize(t *testing.T) { - src := bytes.Repeat([]byte{0x5A}, 4096) - c := codec.LZ4() - encoded, err := c.Encode(nil, src) - if err != nil { - t.Fatalf("Encode: %v", err) - } - - // A zero declaration exercises the codec's nil-destination fallback. It - // remains valid for callers that do not know the decoded size in advance. +func TestDecodePayloadRejectsZeroDeclaredSize(t *testing.T) { hdr := protocol.Header{ Flags: protocol.FlagCompressed, Compression: protocol.CompressionLZ4, } - decoded, err := (&Server{}).decodePayload(hdr, encoded) - if err != nil { - t.Fatalf("decodePayload with zero declared size: %v", err) - } - if !bytes.Equal(decoded, src) { - t.Fatalf("decoded payload differs: got %d bytes, want %d", len(decoded), len(src)) + if _, err := (&Server{}).decodePayload(hdr, []byte{0xFF}); err == nil { + t.Fatal("decodePayload accepted a non-empty compressed payload with zero declared size") } } From 5e13b7588b28dbb3ba0a0b7f8dd3e7466346dedb Mon Sep 17 00:00:00 2001 From: Mark Date: Tue, 11 Aug 2026 00:24:59 +0300 Subject: [PATCH 3/4] ci: close coverage and lint gaps --- .golangci.yml | 1 + internal/codec/lz4.go | 9 ++------- server.go | 4 ---- 3 files changed, 3 insertions(+), 11 deletions(-) diff --git a/.golangci.yml b/.golangci.yml index 11b9746..1c1edf5 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -161,5 +161,6 @@ issues: new: false output: + sort-results: true sort-order: - file diff --git a/internal/codec/lz4.go b/internal/codec/lz4.go index d6f90bf..4f3de8f 100644 --- a/internal/codec/lz4.go +++ b/internal/codec/lz4.go @@ -125,15 +125,10 @@ func (c *lz4Codec) Decode(dst, src []byte) ([]byte, error) { // fallback without allowing len(src)*10 to overflow or exceed the shared // payload bound. func initialDecodeSize(srcLen int) int { - maxSize := maxDecodeSize(srcLen) if srcLen > maxDecodeBuf/10 { - return maxSize + return maxDecodeSize(srcLen) } - initial := srcLen * 10 - if initial > maxSize { - return maxSize - } - return initial + return srcLen * 10 } // maxDecodeSize returns the largest fallback destination permitted for a diff --git a/server.go b/server.go index 886aa57..947bea1 100644 --- a/server.go +++ b/server.go @@ -435,10 +435,6 @@ func (s *Server) decodePayload(hdr protocol.Header, payload []byte) ([]byte, err if len(payload) > 0 && hdr.UncompressedSize == 0 { return nil, fmt.Errorf("compose: compressed payload has zero uncompressed size") } - if len(payload) > protocol.MaxPayloadSize { - return nil, fmt.Errorf("compose: compressed payload size %d exceeds limit %d", - len(payload), protocol.MaxPayloadSize) - } // Look up the codec by compression ID. c := codec.Get(byte(hdr.Compression)) From eb9d885f7cdf1f7647565746e083a34b8a1a4b94 Mon Sep 17 00:00:00 2001 From: Mark Date: Wed, 12 Aug 2026 09:52:18 +0300 Subject: [PATCH 4/4] chore: update golangci-lint v2 output config --- .golangci.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.golangci.yml b/.golangci.yml index 1c1edf5..11b9746 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -161,6 +161,5 @@ issues: new: false output: - sort-results: true sort-order: - file