Add some decoding size limits

go/otel/otelstef/testdata/fuzz/FuzzMetricsReader/e2f5b27e0c475b18

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("STEF\x02\x00\x00\x00\x02\x00\x00\x009\x0f\xca\xca\xca\xca\xca\n@\x00\x00\x00\xec\xeb\rK\x8a\xb0\x00\x00\x80\x00\x00\x00\x00\x04\x04\xa6\x00\x00\xc0\x00\x80\x90\x020\x03\x03@\x020\x000\xcb\xe1\x8dѺ\xfd\xea\u008b\x01\xf9\x87\xba\x9a\xebDž\xaa\x95\x01")

go/otel/otelstef/testdata/fuzz/FuzzMetricsReader/edd03aaec4027c8b

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("STEF\x02\x00\x01\x00\x02\v(\xb5/\xfd\x04h\x10\x00\x00\x00\x00\x00\x1c\x1f\xe0\x00\x00\x0f\x04,\x80\xff\xff\xff\x00\x00\x00\x00\x00\x00\x80\x00\x00 \x00\xc0\xe3\xec\u0099\xf2\xe1\xa7\xe8\x1e")

go/otel/otelstef/testdata/fuzz/FuzzMetricsReader/fbd78644626e4442

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("STEF\x0201\x00\xee\xee\xee\xee\xee0\xee\xee\xee\xee\xee\xee\xee\xee0")

go/pkg/basereader.go

@@ -52,7 +52,7 @@ func (r *BaseReader) ReadFixedHeader() error {
 		return err
 	}
 
-	if contentSize < 2 || contentSize > HdrContentSizeLimit {
+	if contentSize < 2 || contentSize > FixedHdrContentSizeLimit {
 		return ErrInvalidHeader
 	}
 
@@ -85,7 +85,12 @@ func (r *BaseReader) ReadVarHeader(ownSchema schema.WireSchema) error {
 		return err
 	}
 
-	hdrBytes := make([]byte, r.FrameDecoder.RemainingSize())
+	hdrSize := r.FrameDecoder.RemainingSize()
+	if hdrSize > VarHdrContentSizeLimit {
+		return ErrInvalidVarHeader
+	}
+
+	hdrBytes := make([]byte, hdrSize)
 	n, err := r.FrameDecoder.Read(hdrBytes)
 	if err != nil {
 		return err
@@ -124,7 +129,7 @@ func (r *BaseReader) NextFrame() (FrameFlags, error) {
 		return 0, err
 	}
 
-	if err := r.ReadBufs.ReadFrom(&r.FrameDecoder); err != nil {
+	if err := r.ReadBufs.ReadFrom(&r.FrameDecoder, r.FrameDecoder.RemainingSize()); err != nil {
 		return 0, err
 	}
 

go/pkg/errors.go

@@ -1,13 +1,28 @@
 package pkg
 
-import "errors"
+var ErrMultimap = NewDecodeError("invalid multimap")
+var ErrMultimapCountLimit = NewDecodeError("too many elements in the multimap")
+var ErrInvalidRefNum = NewDecodeError("invalid refNum")
+var ErrInvalidOneOfType = NewDecodeError("invalid oneof type")
 
-var ErrMultimap = errors.New("invalid multimap")
-var ErrMultimapCountLimit = errors.New("too many elements in the multimap")
-var ErrInvalidRefNum = errors.New("invalid refNum")
-var ErrInvalidOneOfType = errors.New("invalid oneof type")
+var ErrInvalidHeader = NewDecodeError("invalid FixedHeader")
+var ErrInvalidHeaderSignature = NewDecodeError("invalid FixedHeader signature")
+var ErrInvalidFormatVersion = NewDecodeError("invalid format version in the FixedHeader")
+var ErrInvalidCompression = NewDecodeError("invalid compression method")
 
-var ErrInvalidHeader = errors.New("invalid FixedHeader")
-var ErrInvalidHeaderSignature = errors.New("invalid FixedHeader signature")
-var ErrInvalidFormatVersion = errors.New("invalid format version in the FixedHeader")
-var ErrInvalidCompression = errors.New("invalid compression method")
+var ErrInvalidVarHeader = NewDecodeError("invalid VarHeader")
+
+var ErrColumnSizeLimitExceeded = NewDecodeError("column size limit exceeded")
+var ErrTotalColumnSizeLimitExceeded = NewDecodeError("total column size limit exceeded")
+
+type DecodeError struct {
+	msg string
+}
+
+func (e *DecodeError) Error() string {
+	return e.msg
+}
+
+func NewDecodeError(msg string) error {
+	return &DecodeError{msg: msg}
+}

go/pkg/limits.go

@@ -2,4 +2,5 @@ package pkg
 
 const MultimapElemCountLimit = 1024
 
-const HdrContentSizeLimit = 1 << 20
+const FixedHdrContentSizeLimit = 1 << 20
+const VarHdrContentSizeLimit = 1 << 20

go/pkg/recordbuf.go

@@ -155,9 +155,17 @@ func (s *ReadColumnSet) SubColumnLen() int {
 	return len(s.subColumns)
 }
 
-func (s *ReadColumnSet) ReadSizesFrom(buf *BitsReader) error {
+// ReadSizesFrom reads sizes of the column and its subcolumns from buf.
+// It will honor the readLimit to avoid reading too much data and will
+// decrease the readLimit by the size of data that is read.
+func (s *ReadColumnSet) ReadSizesFrom(buf *BitsReader, readLimit *uint64) error {
 	// Read data size
 	dataSize := buf.ReadUvarintCompact()
+	if dataSize > *readLimit {
+		return ErrColumnSizeLimitExceeded
+	}
+	*readLimit -= dataSize
+
 	s.column.data = EnsureLen(s.column.data, int(dataSize))
 
 	if dataSize == 0 {
@@ -170,7 +178,7 @@ func (s *ReadColumnSet) ReadSizesFrom(buf *BitsReader) error {
 
 	// Recursively read subcolumns
 	for i := 0; i < len(s.subColumns); i++ {
-		if err := s.subColumns[i].ReadSizesFrom(buf); err != nil {
+		if err := s.subColumns[i].ReadSizesFrom(buf, readLimit); err != nil {
 			return err
 		}
 	}
@@ -191,18 +199,9 @@ func (s *ReadColumnSet) ReadDataFrom(buf ByteAndBlockReader) error {
 		}
 	}
 
-	//s.readIndex = 0
-
 	return nil
 }
 
-func (s *ReadColumnSet) PrintSchema(indent int) {
-	//fmt.Printf("%s%d\n", strings.Repeat("-", indent), len(s.subColumns))
-	//for _, subColumn := range s.subColumns {
-	//	subColumn.PrintSchema(indent + 1)
-	//}
-}
-
 func (s *ReadColumnSet) ResetData() {
 	s.column.data = nil
 	for i := range s.subColumns {
@@ -214,20 +213,28 @@ type ReadBufs struct {
 	Columns      ReadColumnSet
 	tempBuf      BitsReader
 	tempBufBytes []byte
+	readLimit    uint64
 }
 
-func (s *ReadBufs) ReadFrom(buf ByteAndBlockReader) error {
+func (s *ReadBufs) ReadFrom(buf ByteAndBlockReader, readLimit uint64) error {
 	bufSize, err := binary.ReadUvarint(buf)
 	if err != nil {
 		return err
 	}
+
+	if bufSize > readLimit {
+		return ErrTotalColumnSizeLimitExceeded
+	}
+
 	s.tempBufBytes = EnsureLen(s.tempBufBytes, int(bufSize))
 	if _, err := io.ReadFull(buf, s.tempBufBytes); err != nil {
 		return err
 	}
 	s.tempBuf.Reset(s.tempBufBytes)
 
-	if err := s.Columns.ReadSizesFrom(&s.tempBuf); err != nil {
+	// Keep track of remaining read limit for column sizes and data
+	s.readLimit = readLimit - bufSize
+	if err := s.Columns.ReadSizesFrom(&s.tempBuf, &s.readLimit); err != nil {
 		return err
 	}