From fa1f5d913d560d0d30127f3e399772f4b831f199 Mon Sep 17 00:00:00 2001
From: Atsushi Watanabe <atsushi.w@ieee.org>
Date: Tue, 27 Feb 2024 17:55:42 +0900
Subject: [PATCH] SampleBuilder: Return header of sample head packet

To get RTP extension data associated with the sample.
---
 pkg/media/media.go                            |  1 +
 pkg/media/samplebuilder/samplebuilder.go      |  9 ++++-
 pkg/media/samplebuilder/samplebuilder_test.go | 38 +++++++++----------
 3 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/pkg/media/media.go b/pkg/media/media.go
index 386e65de3ef..a7e5ed09965 100644
--- a/pkg/media/media.go
+++ b/pkg/media/media.go
@@ -18,6 +18,7 @@ type Sample struct {
 	PacketTimestamp    uint32
 	PrevDroppedPackets uint16
 	Metadata           interface{}
+	RTPHeader          *rtp.Header
 }
 
 // Writer defines an interface to handle
diff --git a/pkg/media/samplebuilder/samplebuilder.go b/pkg/media/samplebuilder/samplebuilder.go
index 74964f30490..9581973f589 100644
--- a/pkg/media/samplebuilder/samplebuilder.go
+++ b/pkg/media/samplebuilder/samplebuilder.go
@@ -272,13 +272,17 @@ func (s *SampleBuilder) buildSample(purgingBuffers bool) *media.Sample {
 	// merge all the buffers into a sample
 	data := []byte{}
 	var metadata interface{}
+	var rtpHeader rtp.Header
 	for i := consume.head; i != consume.tail; i++ {
 		p, err := s.depacketizer.Unmarshal(s.buffer[i].Payload)
 		if err != nil {
 			return nil
 		}
-		if i == consume.head && s.packetHeadHandler != nil {
-			metadata = s.packetHeadHandler(s.depacketizer)
+		if i == consume.head {
+			if s.packetHeadHandler != nil {
+				metadata = s.packetHeadHandler(s.depacketizer)
+			}
+			rtpHeader = s.buffer[i].Header.Clone()
 		}
 
 		data = append(data, p...)
@@ -291,6 +295,7 @@ func (s *SampleBuilder) buildSample(purgingBuffers bool) *media.Sample {
 		PacketTimestamp:    sampleTimestamp,
 		PrevDroppedPackets: s.droppedPackets,
 		Metadata:           metadata,
+		RTPHeader:          &rtpHeader,
 	}
 
 	s.droppedPackets = 0
diff --git a/pkg/media/samplebuilder/samplebuilder_test.go b/pkg/media/samplebuilder/samplebuilder_test.go
index 276406ddc50..ceffa10dcc6 100644
--- a/pkg/media/samplebuilder/samplebuilder_test.go
+++ b/pkg/media/samplebuilder/samplebuilder_test.go
@@ -84,8 +84,8 @@ func TestSampleBuilder(t *testing.T) {
 				{Header: rtp.Header{SequenceNumber: 5002, Timestamp: 7}, Payload: []byte{0x03}},
 			},
 			samples: []*media.Sample{
-				{Data: []byte{0x01}, Duration: time.Second, PacketTimestamp: 5},
-				{Data: []byte{0x02}, Duration: time.Second, PacketTimestamp: 6},
+				{Data: []byte{0x01}, Duration: time.Second, PacketTimestamp: 5, RTPHeader: &rtp.Header{SequenceNumber: 5000, Timestamp: 5}},
+				{Data: []byte{0x02}, Duration: time.Second, PacketTimestamp: 6, RTPHeader: &rtp.Header{SequenceNumber: 5001, Timestamp: 6}},
 			},
 			maxLate:          50,
 			maxLateTimestamp: 0,
@@ -102,7 +102,7 @@ func TestSampleBuilder(t *testing.T) {
 				{Header: rtp.Header{SequenceNumber: 5012, Timestamp: 17}, Payload: []byte{0x07}},
 			},
 			samples: []*media.Sample{
-				{Data: []byte{0x01}, Duration: time.Second * 2, PacketTimestamp: 5},
+				{Data: []byte{0x01}, Duration: time.Second * 2, PacketTimestamp: 5, RTPHeader: &rtp.Header{SequenceNumber: 5000, Timestamp: 5, Marker: true}},
 			},
 			maxLate:          5,
 			maxLateTimestamp: 0,
@@ -134,8 +134,8 @@ func TestSampleBuilder(t *testing.T) {
 				{Header: rtp.Header{SequenceNumber: 5012, Timestamp: 17}, Payload: []byte{0x07}},
 			},
 			samples: []*media.Sample{
-				{Data: []byte{0x01}, Duration: time.Second * 2, PacketTimestamp: 5},
-				{Data: []byte{0x02}, Duration: time.Second * 2, PacketTimestamp: 7, PrevDroppedPackets: 1},
+				{Data: []byte{0x01}, Duration: time.Second * 2, PacketTimestamp: 5, RTPHeader: &rtp.Header{SequenceNumber: 5000, Timestamp: 5, Marker: true}},
+				{Data: []byte{0x02}, Duration: time.Second * 2, PacketTimestamp: 7, PrevDroppedPackets: 1, RTPHeader: &rtp.Header{SequenceNumber: 5002, Timestamp: 7, Marker: true}},
 			},
 			maxLate:          5,
 			maxLateTimestamp: 0,
@@ -149,8 +149,8 @@ func TestSampleBuilder(t *testing.T) {
 				{Header: rtp.Header{SequenceNumber: 5003, Timestamp: 7}, Payload: []byte{0x04}},
 			},
 			samples: []*media.Sample{
-				{Data: []byte{0x01}, Duration: time.Second, PacketTimestamp: 5},
-				{Data: []byte{0x02, 0x03}, Duration: time.Second, PacketTimestamp: 6},
+				{Data: []byte{0x01}, Duration: time.Second, PacketTimestamp: 5, RTPHeader: &rtp.Header{SequenceNumber: 5000, Timestamp: 5}},
+				{Data: []byte{0x02, 0x03}, Duration: time.Second, PacketTimestamp: 6, RTPHeader: &rtp.Header{SequenceNumber: 5001, Timestamp: 6}},
 			},
 			maxLate:          50,
 			maxLateTimestamp: 0,
@@ -203,11 +203,11 @@ func TestSampleBuilder(t *testing.T) {
 				{Header: rtp.Header{SequenceNumber: 5005, Timestamp: 6}, Payload: []byte{0x06}},
 			},
 			samples: []*media.Sample{
-				{Data: []byte{0x01}, Duration: time.Second, PacketTimestamp: 1},
-				{Data: []byte{0x02}, Duration: time.Second, PacketTimestamp: 2},
-				{Data: []byte{0x03}, Duration: time.Second, PacketTimestamp: 3},
-				{Data: []byte{0x04}, Duration: time.Second, PacketTimestamp: 4},
-				{Data: []byte{0x05}, Duration: time.Second, PacketTimestamp: 5},
+				{Data: []byte{0x01}, Duration: time.Second, PacketTimestamp: 1, RTPHeader: &rtp.Header{SequenceNumber: 5000, Timestamp: 1}},
+				{Data: []byte{0x02}, Duration: time.Second, PacketTimestamp: 2, RTPHeader: &rtp.Header{SequenceNumber: 5001, Timestamp: 2}},
+				{Data: []byte{0x03}, Duration: time.Second, PacketTimestamp: 3, RTPHeader: &rtp.Header{SequenceNumber: 5002, Timestamp: 3}},
+				{Data: []byte{0x04}, Duration: time.Second, PacketTimestamp: 4, RTPHeader: &rtp.Header{SequenceNumber: 5003, Timestamp: 4}},
+				{Data: []byte{0x05}, Duration: time.Second, PacketTimestamp: 5, RTPHeader: &rtp.Header{SequenceNumber: 5004, Timestamp: 5}},
 			},
 			maxLate:          50,
 			maxLateTimestamp: 0,
@@ -225,7 +225,7 @@ func TestSampleBuilder(t *testing.T) {
 				{Header: rtp.Header{SequenceNumber: 5017, Timestamp: 7001}, Payload: []byte{0x05}},
 			},
 			samples: []*media.Sample{
-				{Data: []byte{0x04, 0x05}, Duration: time.Second * time.Duration(2), PacketTimestamp: 4000, PrevDroppedPackets: 13},
+				{Data: []byte{0x04, 0x05}, Duration: time.Second * time.Duration(2), PacketTimestamp: 4000, PrevDroppedPackets: 13, RTPHeader: &rtp.Header{SequenceNumber: 5013, Timestamp: 4000}},
 			},
 			withHeadChecker:  true,
 			headBytes:        []byte{0x04},
@@ -247,7 +247,7 @@ func TestSampleBuilder(t *testing.T) {
 			withHeadChecker: true,
 			headBytes:       []byte{1},
 			samples: []*media.Sample{
-				{Data: []byte{1, 2, 3}, Duration: 0, PacketTimestamp: 1, PrevDroppedPackets: 0}, // first sample
+				{Data: []byte{1, 2, 3}, Duration: 0, PacketTimestamp: 1, PrevDroppedPackets: 0, RTPHeader: &rtp.Header{SequenceNumber: 5000, Timestamp: 1}}, // first sample
 			},
 			maxLate:          50,
 			maxLateTimestamp: 2000,
@@ -265,7 +265,7 @@ func TestSampleBuilder(t *testing.T) {
 			withHeadChecker: true,
 			headBytes:       []byte{1},
 			samples: []*media.Sample{
-				{Data: []byte{1, 2}, Duration: 0, PacketTimestamp: 1, PrevDroppedPackets: 0}, // 1st sample
+				{Data: []byte{1, 2}, Duration: 0, PacketTimestamp: 1, PrevDroppedPackets: 0, RTPHeader: &rtp.Header{SequenceNumber: 5000, Timestamp: 1}}, // 1st sample
 			},
 			maxLate:          50,
 			maxLateTimestamp: 2000,
@@ -309,18 +309,18 @@ func TestSampleBuilderMaxLate(t *testing.T) {
 	s.Push(&rtp.Packet{Header: rtp.Header{SequenceNumber: 0, Timestamp: 1}, Payload: []byte{0x01}})
 	s.Push(&rtp.Packet{Header: rtp.Header{SequenceNumber: 1, Timestamp: 2}, Payload: []byte{0x01}})
 	s.Push(&rtp.Packet{Header: rtp.Header{SequenceNumber: 2, Timestamp: 3}, Payload: []byte{0x01}})
-	assert.Equal(&media.Sample{Data: []byte{0x01}, Duration: time.Second, PacketTimestamp: 1}, s.Pop(), "Failed to build samples before gap")
+	assert.Equal(&media.Sample{Data: []byte{0x01}, Duration: time.Second, PacketTimestamp: 1, RTPHeader: &rtp.Header{SequenceNumber: 0, Timestamp: 1}}, s.Pop(), "Failed to build samples before gap")
 
 	s.Push(&rtp.Packet{Header: rtp.Header{SequenceNumber: 5000, Timestamp: 500}, Payload: []byte{0x02}})
 	s.Push(&rtp.Packet{Header: rtp.Header{SequenceNumber: 5001, Timestamp: 501}, Payload: []byte{0x02}})
 	s.Push(&rtp.Packet{Header: rtp.Header{SequenceNumber: 5002, Timestamp: 502}, Payload: []byte{0x02}})
 
-	assert.Equal(&media.Sample{Data: []byte{0x01}, Duration: time.Second, PacketTimestamp: 2}, s.Pop(), "Failed to build samples after large gap")
+	assert.Equal(&media.Sample{Data: []byte{0x01}, Duration: time.Second, PacketTimestamp: 2, RTPHeader: &rtp.Header{SequenceNumber: 1, Timestamp: 2}}, s.Pop(), "Failed to build samples after large gap")
 	assert.Equal((*media.Sample)(nil), s.Pop(), "Failed to build samples after large gap")
 
 	s.Push(&rtp.Packet{Header: rtp.Header{SequenceNumber: 6000, Timestamp: 600}, Payload: []byte{0x03}})
-	assert.Equal(&media.Sample{Data: []byte{0x02}, Duration: time.Second, PacketTimestamp: 500, PrevDroppedPackets: 4998}, s.Pop(), "Failed to build samples after large gap")
-	assert.Equal(&media.Sample{Data: []byte{0x02}, Duration: time.Second, PacketTimestamp: 501}, s.Pop(), "Failed to build samples after large gap")
+	assert.Equal(&media.Sample{Data: []byte{0x02}, Duration: time.Second, PacketTimestamp: 500, PrevDroppedPackets: 4998, RTPHeader: &rtp.Header{SequenceNumber: 5000, Timestamp: 500}}, s.Pop(), "Failed to build samples after large gap")
+	assert.Equal(&media.Sample{Data: []byte{0x02}, Duration: time.Second, PacketTimestamp: 501, RTPHeader: &rtp.Header{SequenceNumber: 5001, Timestamp: 501}}, s.Pop(), "Failed to build samples after large gap")
 }
 
 func TestSeqnumDistance(t *testing.T) {