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..4f3de8f 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,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 diff --git a/internal/codec/lz4_test.go b/internal/codec/lz4_test.go index 417891f..e912f02 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,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() 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..947bea1 100644 --- a/server.go +++ b/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 { diff --git a/server_security_test.go b/server_security_test.go new file mode 100644 index 0000000..758ea47 --- /dev/null +++ b/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") + } +}