Fix panics in h264 Spreader.Process

codecs/h264/spreader.go

@@ -207,6 +207,10 @@ func (s *Spreader) Process(payload []byte) (outPayloads [][]byte, err error) { /
 		}
 	}
 
+	if s.Mtu-rtpHeaderSize-fuaOverhead <= 0 {
+		return nil, fmt.Errorf("mtu(%d) is too small for rtp header(%d)", s.Mtu, rtpHeaderSize) //nolint:err113
+	}
+
 	outPayloads, _, err = s.handleNalTooBigOrFua(outPayloads, seqNum, naluType, rtpHeaderData, nalData)
 	if err != nil {
 		return nil, err
@@ -230,7 +234,7 @@ func (s *Spreader) handleNalTooBigOrFua(cumulRTP [][]byte, seqNum uint16, naluTy
 }
 
 // relying on continuous seq number & start/end FU bits to sync ourselve, so not looking at RtpTimestamp.
-func (s *Spreader) spreadFua(cumulRTP [][]byte, firtSeqNum uint16, rtpHeader []byte, fua []byte) ([][]byte, uint16, error) { //nolint:lll
+func (s *Spreader) spreadFua(cumulRTP [][]byte, firtSeqNum uint16, rtpHeader []byte, fua []byte) ([][]byte, uint16, error) { //nolint:lll,cyclop
 	seqNum := firtSeqNum
 	if s.fuInProgress != nil {
 		expectedSeq := s.fuInProgress.LastSeq + 1
@@ -262,7 +266,8 @@ func (s *Spreader) spreadFua(cumulRTP [][]byte, firtSeqNum uint16, rtpHeader []b
 	reqSubSize := s.Mtu - lenRTPHeader - fuaOverhead
 	newData := fua[fuaOverhead:]
 	currentDataSize := len(s.fuInProgress.Trailing) + len(newData)
-	for currentDataSize > reqSubSize || (mustFinish && currentDataSize > 0) {
+	// an ending FU is always emitted, even empty, so that the end bit is not lost
+	for currentDataSize > reqSubSize || (mustFinish && (currentDataSize > 0 || lastFuHeader == nil)) {
 		bufSize := min(s.Mtu, lenRTPHeader+fuaOverhead+currentDataSize)
 		rtp := make([]byte, bufSize)
 		binary.BigEndian.PutUint16(rtpHeader[rtpSeqNumOffset:rtpSeqNumOffset+rtpSeqNumLength], seqNum)
@@ -293,11 +298,16 @@ func (s *Spreader) spreadFua(cumulRTP [][]byte, firtSeqNum uint16, rtpHeader []b
 		*lastFuHeader |= fuEndBitmask
 		s.fuInProgress = nil
 	} else {
-		copy(s.trailingBuf, newData)
-		s.fuInProgress.Trailing = s.trailingBuf[:len(newData)]
+		// trailing is still pending when nothing was emitted, new data goes after it
+		lenTrailing := len(s.fuInProgress.Trailing)
+		copy(s.trailingBuf[lenTrailing:], newData)
+		s.fuInProgress.Trailing = s.trailingBuf[:lenTrailing+len(newData)]
 	}
 
-	cumulRTP[len(cumulRTP)-1][rtpMPtOffset] |= entryMarker
+	// nothing is emitted when the whole fragment fits in the trailing buffer
+	if lastFuHeader != nil {
+		cumulRTP[len(cumulRTP)-1][rtpMPtOffset] |= entryMarker
+	}
 
 	return cumulRTP, seqNum, nil
 }
@@ -384,6 +394,10 @@ func (s *Spreader) explodeStapA(
 	seqNum := firtSeqNum
 	var err error
 	for currOffset < lenStapA {
+		if lenStapA < currOffset+stapaNALULengthSize {
+			return nil, seqNum, fmt.Errorf("STAP-A NALU size is truncated: %d", lenStapA-currOffset) //nolint:err113
+		}
+
 		naluSize := int(binary.BigEndian.Uint16(stapa[currOffset:]))
 		currOffset += stapaNALULengthSize
 

codecs/h264/spreader_test.go

@@ -251,6 +251,112 @@ func TestChunkSlicing(t *testing.T) {
 	assert.Equal(slice17to4[4], []byte{17}, "Bad Slicing")
 }
 
+func TestStapATruncatedNaluSize(t *testing.T) {
+	assert := assert.New(t)
+
+	// arrange
+	spreader := NewSpreader(100)
+	stapa := append([]byte{0x18, 0x00, 0x78, 0x65}, make([]byte, 119)...)
+	stapa = append(stapa, 0x00) // stray byte where the next NALU size should be
+	rtpPkt := append(getRTPHeader(false), stapa...)
+
+	// act
+	out, err := spreader.Process(rtpPkt)
+
+	// assert
+	assert.Error(err, "should reject truncated STAP-A")
+	assert.Nil(out, "should not output pkts")
+}
+
+func TestFuaSmallerThanTrailingRoomIsBuffered(t *testing.T) {
+	assert := assert.New(t)
+
+	// arrange
+	spreader := NewSpreader(100)
+	fuaStart := append([]byte{0x7C, 0x85}, make([]byte, 136)...)
+	fua1 := append(getRTPHeaderOffSet(0, false), fuaStart...)
+	fua2 := append(getRTPHeaderOffSet(1, false), 0x7C, 0x05, 0xAA)
+	fua3 := append(getRTPHeaderOffSet(2, true), 0x7C, 0x45, 0xBB)
+
+	// act
+	out1, err1 := spreader.Process(fua1)
+	out2, err2 := spreader.Process(fua2)
+	out3, err3 := spreader.Process(fua3)
+
+	// assert
+	assert.NoError(err1, "failed to process pkt")
+	assert.NoError(err2, "failed to process pkt")
+	assert.NoError(err3, "failed to process pkt")
+
+	assert.Equal(1, len(out1), "bad pkt out count")
+	assert.Equal(0, len(out2), "middle FUA should be kept as trailing")
+	if assert.Equal(1, len(out3), "bad pkt out count") {
+		expectedEnd := append([]byte{0x7C, 0x45}, make([]byte, 50)...)
+		expectedEnd = append(expectedEnd, 0xAA, 0xBB)
+		assert.Equal(append(getRTPHeaderOffSet(1, true), expectedEnd...), out3[0], "bad payload end FUA")
+	}
+	assert.Equal(uint16(0xFFFF), spreader.RTPOffset, "Bad offset")
+}
+
+func TestFuaEmptyEndIsEmitted(t *testing.T) {
+	assert := assert.New(t)
+
+	// arrange
+	spreader := NewSpreader(100)
+	rtpPkt := append(getRTPHeader(true), 0x7C, 0x45)
+	rtpPkt[0] |= 0x20 // only the padding makes it larger than the mtu
+	rtpPkt = append(rtpPkt, make([]byte, 99)...)
+	rtpPkt = append(rtpPkt, 100)
+
+	// act
+	out, err := spreader.Process(rtpPkt)
+
+	// assert
+	if assert.NoError(err, "failed to process pkt") && assert.Equal(1, len(out), "bad pkt out count") {
+		assert.Equal(append(getRTPHeader(true), 0x7C, 0x45), out[0], "bad payload end FUA")
+	}
+}
+
+func TestRTPHeaderLargerThanMtu(t *testing.T) {
+	assert := assert.New(t)
+
+	// arrange
+	spreader := NewSpreader(100)
+	rtpPkt := []byte{0x90, 0x62, 0x00, 0x13, 0x00, 0x00, 0x00, 0x00, 0x70, 0x1c, 0x77, 0x48, 0xBE, 0xDE, 0x00, 0x16}
+	rtpPkt = append(rtpPkt, make([]byte, 88)...) // extension
+	rtpPkt = append(rtpPkt, 0x65, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07)
+
+	// act
+	out, err := spreader.Process(rtpPkt)
+
+	// assert
+	assert.Error(err, "should reject header that leaves no room for payload")
+	assert.Nil(out, "should not output pkts")
+}
+
+func FuzzSpreaderProcess(f *testing.F) {
+	stapa := append(getRTPHeader(false), getH264StapADelimiterAndSpsAndPps()...)
+	sps := append(getRTPHeader(true), getH264SpsNal()...)
+	fuas := getRtpsFuaNonIdr()
+	f.Add(uint16(26), stapa, sps, stapa)
+	f.Add(uint16(24), sps, stapa, sps)
+	f.Add(uint16(1200), fuas[0], fuas[1], fuas[2])
+	f.Add(uint16(100), fuas[0], fuas[1], sps)
+
+	f.Fuzz(func(t *testing.T, mtu uint16, pkt1, pkt2, pkt3 []byte) {
+		spreader := NewSpreader(mtu)
+		for _, pkt := range [][]byte{pkt1, pkt2, pkt3} {
+			out, err := spreader.Process(pkt)
+			if err != nil {
+				continue
+			}
+			for _, rtpOut := range out {
+				assert.LessOrEqual(t, len(rtpOut), int(mtu), "pkt out is larger than mtu")
+			}
+		}
+	})
+}
+
 func getRTPHeader(mark bool) []byte {
 	rtpMarkerBitMask := byte(0x80)
 	rtpHeader := []byte{0x80, 0x62, 0x00, 0x13, 0x00, 0x00, 0x00, 0x00, 0x70, 0x1c, 0x77, 0x48}