rtmp: add limit on message body size (#2252)

internal/rtmp/message/reader.go

@@ -145,7 +145,10 @@ func (r *Reader) Read() (Message, error) {
 
 	switch tmsg := msg.(type) {
 	case *SetChunkSize:
-		r.r.SetChunkSize(tmsg.Value)
+		err := r.r.SetChunkSize(tmsg.Value)
+		if err != nil {
+			return nil, err
+		}
 
 	case *SetWindowAckSize:
 		r.r.SetWindowAckSize(tmsg.Value)

internal/rtmp/message/reader_test.go

@@ -290,3 +290,16 @@ func TestReader(t *testing.T) {
 		})
 	}
 }
+
+func FuzzReader(f *testing.F) {
+	f.Add([]byte{
+		0x04, 0x00, 0x3a, 0xfc, 0x00, 0x00, 0x08, 0x09,
+		0x01, 0x00, 0x00, 0x00, 0x88, 0x68, 0x76, 0x63,
+		0x31, 0x01, 0x02, 0x03,
+	})
+	f.Fuzz(func(t *testing.T, b []byte) {
+		bc := bytecounter.NewReader(bytes.NewReader(b))
+		r := NewReader(bc, bc, nil)
+		r.Read() //nolint:errcheck
+	})
+}

internal/rtmp/message/testdata/fuzz/FuzzReader/05d2521061b772dd

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000\x00\x00000000")

internal/rtmp/message/testdata/fuzz/FuzzReader/06f5bdb4e0ba6885

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000\x00\x00\f\x120000\f\x00\x00\x00\x010000000")

internal/rtmp/message/testdata/fuzz/FuzzReader/2fb5da434799f2aa

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000\x00\x00\x00\t0000")

internal/rtmp/message/testdata/fuzz/FuzzReader/420fac969d79c3d0

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000\x00\x00\f\x040000000000000000")

internal/rtmp/message/testdata/fuzz/FuzzReader/6b1d357b508b38a4

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000\x00\x00\v\t0000\xe50000000000")

internal/rtmp/message/testdata/fuzz/FuzzReader/b18a01392c7c91e9

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000\x00\x00\x0000000")

internal/rtmp/message/testdata/fuzz/FuzzReader/f244ff2f55d1103f

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000\x00\x00\x00\x040000")

internal/rtmp/rawmessage/reader.go

@@ -13,14 +13,30 @@ import (
 
 var errMoreChunksNeeded = errors.New("more chunks are needed")
 
+const (
+	maxBodySize = 10 * 1024 * 1024
+)
+
+func joinFragments(fragments [][]byte, size uint32) []byte {
+	ret := make([]byte, size)
+	n := 0
+	for _, p := range fragments {
+		n += copy(ret[n:], p)
+	}
+	return ret
+}
+
 type readerChunkStream struct {
-	mr                 *Reader
-	curTimestamp       *uint32
-	curType            *uint8
-	curMessageStreamID *uint32
-	curBodyLen         *uint32
-	curBody            []byte
-	curTimestampDelta  *uint32
+	mr                         *Reader
+	curTimestamp               uint32
+	curTimestampAvailable      bool
+	curType                    uint8
+	curMessageStreamID         uint32
+	curBodyLen                 uint32
+	curBodyFragments           [][]byte
+	curBodyRecv                uint32
+	curTimestampDelta          uint32
+	curTimestampDeltaAvailable bool
 }
 
 func (rc *readerChunkStream) readChunk(c chunk.Chunk, chunkBodySize uint32) error {
@@ -50,7 +66,7 @@ func (rc *readerChunkStream) readChunk(c chunk.Chunk, chunkBodySize uint32) erro
 func (rc *readerChunkStream) readMessage(typ byte) (*Message, error) {
 	switch typ {
 	case 0:
-		if rc.curBody != nil {
+		if rc.curBodyRecv != 0 {
 			return nil, fmt.Errorf("received type 0 chunk but expected type 3 chunk")
 		}
 
@@ -59,18 +75,22 @@ func (rc *readerChunkStream) readMessage(typ byte) (*Message, error) {
 			return nil, err
 		}
 
-		v1 := rc.mr.c0.MessageStreamID
-		rc.curMessageStreamID = &v1
-		v2 := rc.mr.c0.Type
-		rc.curType = &v2
-		v3 := rc.mr.c0.Timestamp
-		rc.curTimestamp = &v3
-		v4 := rc.mr.c0.BodyLen
-		rc.curBodyLen = &v4
-		rc.curTimestampDelta = nil
-
-		if rc.mr.c0.BodyLen != uint32(len(rc.mr.c0.Body)) {
-			rc.curBody = rc.mr.c0.Body
+		rc.curMessageStreamID = rc.mr.c0.MessageStreamID
+		rc.curType = rc.mr.c0.Type
+		rc.curTimestamp = rc.mr.c0.Timestamp
+		rc.curTimestampAvailable = true
+		rc.curTimestampDeltaAvailable = false
+		rc.curBodyLen = rc.mr.c0.BodyLen
+
+		if rc.curBodyLen > maxBodySize {
+			return nil, fmt.Errorf("body size (%d) exceeds maximum (%d)", rc.curBodyLen, maxBodySize)
+		}
+
+		le := uint32(len(rc.mr.c0.Body))
+
+		if rc.mr.c0.BodyLen != le {
+			rc.curBodyFragments = append(rc.curBodyFragments, rc.mr.c0.Body)
+			rc.curBodyRecv = le
 			return nil, errMoreChunksNeeded
 		}
 
@@ -81,11 +101,11 @@ func (rc *readerChunkStream) readMessage(typ byte) (*Message, error) {
 		return &rc.mr.msg, nil
 
 	case 1:
-		if rc.curTimestamp == nil {
+		if !rc.curTimestampAvailable {
 			return nil, fmt.Errorf("received type 1 chunk without previous chunk")
 		}
 
-		if rc.curBody != nil {
+		if rc.curBodyRecv != 0 {
 			return nil, fmt.Errorf("received type 1 chunk but expected type 3 chunk")
 		}
 
@@ -94,36 +114,40 @@ func (rc *readerChunkStream) readMessage(typ byte) (*Message, error) {
 			return nil, err
 		}
 
-		v2 := rc.mr.c1.Type
-		rc.curType = &v2
-		v3 := *rc.curTimestamp + rc.mr.c1.TimestampDelta
-		rc.curTimestamp = &v3
-		v4 := rc.mr.c1.BodyLen
-		rc.curBodyLen = &v4
-		v5 := rc.mr.c1.TimestampDelta
-		rc.curTimestampDelta = &v5
-
-		if rc.mr.c1.BodyLen != uint32(len(rc.mr.c1.Body)) {
-			rc.curBody = rc.mr.c1.Body
+		rc.curType = rc.mr.c1.Type
+		rc.curTimestamp += rc.mr.c1.TimestampDelta
+		rc.curTimestampDelta = rc.mr.c1.TimestampDelta
+		rc.curTimestampDeltaAvailable = true
+		rc.curBodyLen = rc.mr.c1.BodyLen
+
+		if rc.curBodyLen > maxBodySize {
+			return nil, fmt.Errorf("body size (%d) exceeds maximum (%d)", rc.curBodyLen, maxBodySize)
+		}
+
+		le := uint32(len(rc.mr.c1.Body))
+
+		if rc.mr.c1.BodyLen != le {
+			rc.curBodyFragments = append(rc.curBodyFragments, rc.mr.c1.Body)
+			rc.curBodyRecv = le
 			return nil, errMoreChunksNeeded
 		}
 
-		rc.mr.msg.Timestamp = time.Duration(*rc.curTimestamp) * time.Millisecond
+		rc.mr.msg.Timestamp = time.Duration(rc.curTimestamp) * time.Millisecond
 		rc.mr.msg.Type = rc.mr.c1.Type
-		rc.mr.msg.MessageStreamID = *rc.curMessageStreamID
+		rc.mr.msg.MessageStreamID = rc.curMessageStreamID
 		rc.mr.msg.Body = rc.mr.c1.Body
 		return &rc.mr.msg, nil
 
 	case 2:
-		if rc.curTimestamp == nil {
+		if !rc.curTimestampAvailable {
 			return nil, fmt.Errorf("received type 2 chunk without previous chunk")
 		}
 
-		if rc.curBody != nil {
+		if rc.curBodyRecv != 0 {
 			return nil, fmt.Errorf("received type 2 chunk but expected type 3 chunk")
 		}
 
-		chunkBodyLen := *rc.curBodyLen
+		chunkBodyLen := rc.curBodyLen
 		if chunkBodyLen > rc.mr.chunkSize {
 			chunkBodyLen = rc.mr.chunkSize
 		}
@@ -133,29 +157,27 @@ func (rc *readerChunkStream) readMessage(typ byte) (*Message, error) {
 			return nil, err
 		}
 
-		v1 := *rc.curTimestamp + rc.mr.c2.TimestampDelta
-		rc.curTimestamp = &v1
-		v2 := rc.mr.c2.TimestampDelta
-		rc.curTimestampDelta = &v2
+		rc.curTimestamp += rc.mr.c2.TimestampDelta
+		rc.curTimestampDelta = rc.mr.c2.TimestampDelta
+		rc.curTimestampDeltaAvailable = true
 
-		if *rc.curBodyLen != uint32(len(rc.mr.c2.Body)) {
-			rc.curBody = rc.mr.c2.Body
+		le := uint32(len(rc.mr.c2.Body))
+
+		if rc.curBodyLen != le {
+			rc.curBodyFragments = append(rc.curBodyFragments, rc.mr.c2.Body)
+			rc.curBodyRecv = le
 			return nil, errMoreChunksNeeded
 		}
 
-		rc.mr.msg.Timestamp = time.Duration(*rc.curTimestamp) * time.Millisecond
-		rc.mr.msg.Type = *rc.curType
-		rc.mr.msg.MessageStreamID = *rc.curMessageStreamID
+		rc.mr.msg.Timestamp = time.Duration(rc.curTimestamp) * time.Millisecond
+		rc.mr.msg.Type = rc.curType
+		rc.mr.msg.MessageStreamID = rc.curMessageStreamID
 		rc.mr.msg.Body = rc.mr.c2.Body
 		return &rc.mr.msg, nil
 
 	default: // 3
-		if rc.curBody == nil && rc.curTimestampDelta == nil {
-			return nil, fmt.Errorf("received type 3 chunk without previous chunk")
-		}
-
-		if rc.curBody != nil {
-			chunkBodyLen := (*rc.curBodyLen) - uint32(len(rc.curBody))
+		if rc.curBodyRecv != 0 {
+			chunkBodyLen := rc.curBodyLen - rc.curBodyRecv
 			if chunkBodyLen > rc.mr.chunkSize {
 				chunkBodyLen = rc.mr.chunkSize
 			}
@@ -165,23 +187,27 @@ func (rc *readerChunkStream) readMessage(typ byte) (*Message, error) {
 				return nil, err
 			}
 
-			rc.curBody = append(rc.curBody, rc.mr.c3.Body...)
+			rc.curBodyFragments = append(rc.curBodyFragments, rc.mr.c3.Body)
+			rc.curBodyRecv += uint32(len(rc.mr.c3.Body))
 
-			if *rc.curBodyLen != uint32(len(rc.curBody)) {
+			if rc.curBodyLen != rc.curBodyRecv {
 				return nil, errMoreChunksNeeded
 			}
 
-			body := rc.curBody
-			rc.curBody = nil
-
-			rc.mr.msg.Timestamp = time.Duration(*rc.curTimestamp) * time.Millisecond
-			rc.mr.msg.Type = *rc.curType
-			rc.mr.msg.MessageStreamID = *rc.curMessageStreamID
-			rc.mr.msg.Body = body
+			rc.mr.msg.Timestamp = time.Duration(rc.curTimestamp) * time.Millisecond
+			rc.mr.msg.Type = rc.curType
+			rc.mr.msg.MessageStreamID = rc.curMessageStreamID
+			rc.mr.msg.Body = joinFragments(rc.curBodyFragments, rc.curBodyRecv)
+			rc.curBodyFragments = rc.curBodyFragments[:0]
+			rc.curBodyRecv = 0
 			return &rc.mr.msg, nil
 		}
 
-		chunkBodyLen := (*rc.curBodyLen)
+		if !rc.curTimestampDeltaAvailable {
+			return nil, fmt.Errorf("received type 3 chunk without previous chunk")
+		}
+
+		chunkBodyLen := rc.curBodyLen
 		if chunkBodyLen > rc.mr.chunkSize {
 			chunkBodyLen = rc.mr.chunkSize
 		}
@@ -191,17 +217,19 @@ func (rc *readerChunkStream) readMessage(typ byte) (*Message, error) {
 			return nil, err
 		}
 
-		v1 := *rc.curTimestamp + *rc.curTimestampDelta
-		rc.curTimestamp = &v1
+		rc.curTimestamp += rc.curTimestampDelta
+
+		le := uint32(len(rc.mr.c3.Body))
 
-		if *rc.curBodyLen != uint32(len(rc.mr.c3.Body)) {
-			rc.curBody = rc.mr.c3.Body
+		if rc.curBodyLen != le {
+			rc.curBodyFragments = append(rc.curBodyFragments, rc.mr.c3.Body)
+			rc.curBodyRecv = le
 			return nil, errMoreChunksNeeded
 		}
 
-		rc.mr.msg.Timestamp = time.Duration(*rc.curTimestamp) * time.Millisecond
-		rc.mr.msg.Type = *rc.curType
-		rc.mr.msg.MessageStreamID = *rc.curMessageStreamID
+		rc.mr.msg.Timestamp = time.Duration(rc.curTimestamp) * time.Millisecond
+		rc.mr.msg.Type = rc.curType
+		rc.mr.msg.MessageStreamID = rc.curMessageStreamID
 		rc.mr.msg.Body = rc.mr.c3.Body
 		return &rc.mr.msg, nil
 	}
@@ -240,8 +268,13 @@ func NewReader(
 }
 
 // SetChunkSize sets the maximum chunk size.
-func (r *Reader) SetChunkSize(v uint32) {
+func (r *Reader) SetChunkSize(v uint32) error {
+	if v > maxBodySize {
+		return fmt.Errorf("chunk size (%d) exceeds maximum (%d)", v, maxBodySize)
+	}
+
 	r.chunkSize = v
+	return nil
 }
 
 // SetWindowAckSize sets the window acknowledgement size.

internal/rtmp/rawmessage/reader_test.go

@@ -235,7 +235,9 @@ func TestReaderAcknowledge(t *testing.T) {
 				r.lastAckCount = 4294967096
 			}
 
-			r.SetChunkSize(65536)
+			err := r.SetChunkSize(65536)
+			require.NoError(t, err)
+
 			r.SetWindowAckSize(100)
 
 			buf2, err := chunk.Chunk0{
@@ -256,3 +258,19 @@ func TestReaderAcknowledge(t *testing.T) {
 		})
 	}
 }
+
+func FuzzReader(f *testing.F) {
+	f.Fuzz(func(t *testing.T, b []byte) {
+		br := bytecounter.NewReader(bytes.NewReader(b))
+		r := NewReader(br, br, func(count uint32) error {
+			return nil
+		})
+
+		for {
+			_, err := r.Read()
+			if err != nil {
+				break
+			}
+		}
+	})
+}

internal/rtmp/rawmessage/testdata/fuzz/FuzzDecoder/19981bffc2abbaf1

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("A")

internal/rtmp/rawmessage/testdata/fuzz/FuzzDecoder/2470f01dca6d27ef

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\xe6")

internal/rtmp/rawmessage/testdata/fuzz/FuzzDecoder/2c0aa0d43cf2378b

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000\xf4000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000")

internal/rtmp/rawmessage/testdata/fuzz/FuzzDecoder/3225fb43d226570f

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000\x00\x00000000000000000000000000000000000000000000000000000000\xb0")

internal/rtmp/rawmessage/testdata/fuzz/FuzzDecoder/582528ddfad69eb5

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0")

internal/rtmp/rawmessage/testdata/fuzz/FuzzDecoder/649388b35b8d7d24

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\x00000\x00\x00\x04000000000\x800000000\xc0")

internal/rtmp/rawmessage/testdata/fuzz/FuzzDecoder/899c4ec5c6184841

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000\xf0")

internal/rtmp/rawmessage/testdata/fuzz/FuzzDecoder/a2d2a54b9b1b0098

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000")

internal/rtmp/rawmessage/testdata/fuzz/FuzzDecoder/ab461fd3f1e0b76d

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000p")

internal/rtmp/rawmessage/testdata/fuzz/FuzzDecoder/cf0f70a31328c9ba

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000\xb0")

internal/rtmp/rawmessage/testdata/fuzz/FuzzDecoder/f5aad145f6286289

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\x80")

internal/rtmp/rawmessage/testdata/fuzz/FuzzReader/2a3abe67115a80dc

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000\xbe000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000")

internal/rtmp/rawmessage/testdata/fuzz/FuzzReader/321edca93ba341df

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("0000\x00\x00\x0000000p000\xe200000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000")

internal/rtmp/rawmessage/testdata/fuzz/FuzzReader/7f07c167964a9467

@@ -0,0 +1,2 @@
+go test fuzz v1
+[]byte("\xd6")