Samplebuilder: Properly handle padding packets

See: webrtc-rs/webrtc#309
pion · Sep 15, 2023 · 9e336ad · 9e336ad
1 parent c0437dc
commit 9e336ad
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 0 deletions.
diff --git a/pkg/media/samplebuilder/samplebuilder.go b/pkg/media/samplebuilder/samplebuilder.go
@@ -38,9 +38,14 @@ type SampleBuilder struct {
 	// prepared contains the samples that have been processed to date
 	prepared sampleSequenceLocation
 
+	lastSampleTimestamp *uint32
+
 	// number of packets forced to be dropped
 	droppedPackets uint16
 
+	// number of padding packets detected and dropped (this will be a subset of `droppedPackets`)
+	paddingPackets uint16
+
 	// allows inspecting head packets of each sample and then returns a custom metadata
 	packetHeadHandler func(headPacket interface{}) interface{}
 }
@@ -191,6 +196,7 @@ const secondToNanoseconds = 1000000000
 // buildSample creates a sample from a valid collection of RTP Packets by
 // walking forwards building a sample if everything looks good clear and
 // update buffer+values
+// nolint: gocognit
 func (s *SampleBuilder) buildSample(purgingBuffers bool) *media.Sample {
 	if s.active.empty() {
 		s.active = s.filled
@@ -248,7 +254,16 @@ func (s *SampleBuilder) buildSample(purgingBuffers bool) *media.Sample {
 	// prior to decoding all the packets, check if this packet
 	// would end being disposed anyway
 	if !s.depacketizer.IsPartitionHead(s.buffer[consume.head].Payload) {
+		isPadding := false
+		for i := consume.head; i != consume.tail; i++ {
+			if s.lastSampleTimestamp != nil && *s.lastSampleTimestamp == s.buffer[i].Timestamp && len(s.buffer[i].Payload) == 0 {
+				isPadding = true
+			}
+		}
 		s.droppedPackets += consume.count()
+		if isPadding {
+			s.paddingPackets += consume.count()
+		}
 		s.purgeConsumedLocation(consume, true)
 		s.purgeConsumedBuffers()
 		return nil
@@ -279,6 +294,9 @@ func (s *SampleBuilder) buildSample(purgingBuffers bool) *media.Sample {
 	}
 
 	s.droppedPackets = 0
+	s.paddingPackets = 0
+	s.lastSampleTimestamp = new(uint32)
+	*s.lastSampleTimestamp = sampleTimestamp
 
 	s.preparedSamples[s.prepared.tail] = sample
 	s.prepared.tail++

diff --git a/pkg/media/samplebuilder/samplebuilder_test.go b/pkg/media/samplebuilder/samplebuilder_test.go
@@ -38,6 +38,12 @@ func (f *fakeDepacketizer) IsPartitionHead(payload []byte) bool {
 		// the tests should be fixed to not assume the bug
 		return true
 	}
+
+	// skip padding
+	if len(payload) < 1 {
+		return false
+	}
+
 	for _, b := range f.headBytes {
 		if payload[0] == b {
 			return true
@@ -226,6 +232,46 @@ func TestSampleBuilder(t *testing.T) {
 			maxLate:          50,
 			maxLateTimestamp: 2000,
 		},
+		{
+			// shamelessly stolen from webrtc-rs
+			message: "Sample builder should recognize padding packets",
+			packets: []*rtp.Packet{
+				{Header: rtp.Header{SequenceNumber: 5000, Timestamp: 1}, Payload: []byte{1}},               // 1st packet
+				{Header: rtp.Header{SequenceNumber: 5001, Timestamp: 1}, Payload: []byte{2}},               // 2nd packet
+				{Header: rtp.Header{SequenceNumber: 5002, Timestamp: 1, Marker: true}, Payload: []byte{3}}, // 3rd packet
+				{Header: rtp.Header{SequenceNumber: 5003, Timestamp: 1}, Payload: []byte{}},                // Padding packet 1
+				{Header: rtp.Header{SequenceNumber: 5004, Timestamp: 1}, Payload: []byte{}},                // Padding packet 2
+				{Header: rtp.Header{SequenceNumber: 5005, Timestamp: 3}, Payload: []byte{1}},               // 6th packet
+				{Header: rtp.Header{SequenceNumber: 5006, Timestamp: 3, Marker: true}, Payload: []byte{7}}, // 7th packet
+				{Header: rtp.Header{SequenceNumber: 5007, Timestamp: 4}, Payload: []byte{1}},               // 7th packet
+			},
+			withHeadChecker: true,
+			headBytes:       []byte{1},
+			samples: []*media.Sample{
+				{Data: []byte{1, 2, 3}, Duration: 0, PacketTimestamp: 1, PrevDroppedPackets: 0}, // first sample
+			},
+			maxLate:          50,
+			maxLateTimestamp: 2000,
+		},
+		{
+			// shamelessly stolen from webrtc-rs
+			message: "Sample builder should build a sample out of a packet that's both start and end following a run of padding packets",
+			packets: []*rtp.Packet{
+				{Header: rtp.Header{SequenceNumber: 5000, Timestamp: 1}, Payload: []byte{1}},               // 1st valid packet
+				{Header: rtp.Header{SequenceNumber: 5001, Timestamp: 1, Marker: true}, Payload: []byte{2}}, // 2nd valid packet
+				{Header: rtp.Header{SequenceNumber: 5002, Timestamp: 1}, Payload: []byte{}},                // 1st padding packet
+				{Header: rtp.Header{SequenceNumber: 5003, Timestamp: 1}, Payload: []byte{}},                // 2nd padding packet
+				{Header: rtp.Header{SequenceNumber: 5004, Timestamp: 2, Marker: true}, Payload: []byte{1}}, // 3rd valid packet
+				{Header: rtp.Header{SequenceNumber: 5005, Timestamp: 3}, Payload: []byte{1}},               // 4th valid packet, start of next sample
+			},
+			withHeadChecker: true,
+			headBytes:       []byte{1},
+			samples: []*media.Sample{
+				{Data: []byte{1, 2}, Duration: 0, PacketTimestamp: 1, PrevDroppedPackets: 0}, // 1st sample
+			},
+			maxLate:          50,
+			maxLateTimestamp: 2000,
+		},
 	}
 
 	t.Run("Pop", func(t *testing.T) {