feat: validate agent protocol boundaries

internal/agentproto/compression.go

@@ -0,0 +1,46 @@
+package agentproto
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"io"
+
+	"github.com/klauspost/compress/zstd"
+	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
+)
+
+var ErrInvalidOutputChunk = errors.New("invalid output chunk")
+
+// DecodeOutputChunk validates declared sizes before bounded decompression.
+func DecodeOutputChunk(chunk *rvboxv1.OutputChunk, maxRawBytes uint64) ([]byte, error) {
+	if chunk == nil || (chunk.Stream != rvboxv1.StreamKind_STREAM_STDOUT && chunk.Stream != rvboxv1.StreamKind_STREAM_STDERR) || chunk.CompressedSize != uint64(len(chunk.Data)) || chunk.UncompressedSize > maxRawBytes {
+		return nil, ErrInvalidOutputChunk
+	}
+	switch chunk.Compression {
+	case rvboxv1.Compression_COMPRESSION_NONE:
+		if chunk.UncompressedSize != uint64(len(chunk.Data)) {
+			return nil, ErrInvalidOutputChunk
+		}
+		return bytes.Clone(chunk.Data), nil
+	case rvboxv1.Compression_COMPRESSION_ZSTD:
+		if chunk.UncompressedSize == 0 && len(chunk.Data) != 0 {
+			return nil, ErrInvalidOutputChunk
+		}
+		decoder, err := zstd.NewReader(bytes.NewReader(chunk.Data), zstd.WithDecoderConcurrency(1), zstd.WithDecoderMaxMemory(maxRawBytes+1))
+		if err != nil {
+			return nil, fmt.Errorf("%w: zstd header: %v", ErrInvalidOutputChunk, err)
+		}
+		defer decoder.Close()
+		decoded, err := io.ReadAll(io.LimitReader(decoder, int64(maxRawBytes)+1))
+		if err != nil {
+			return nil, fmt.Errorf("%w: zstd decode: %v", ErrInvalidOutputChunk, err)
+		}
+		if uint64(len(decoded)) != chunk.UncompressedSize || uint64(len(decoded)) > maxRawBytes {
+			return nil, ErrInvalidOutputChunk
+		}
+		return decoded, nil
+	default:
+		return nil, ErrInvalidOutputChunk
+	}
+}