Support RTP extensions when writing samples

Author: fancycode

track_local_static.go

@@ -230,6 +230,15 @@ func (s *TrackLocalStaticRTP) Write(b []byte) (n int, err error) {
 	return len(b), s.writeRTP(packet)
 }
 
+// SampleRTPExtension is a RTP extension that can be added to packets generated
+// when writing Samples to a TrackLocalStaticSample.
+type SampleRTPExtension struct {
+	// ID is the negotiated extension id.
+	ID uint8
+	// Payload is the payload of the extension to add to the packet.
+	Payload []byte
+}
+
 // TrackLocalStaticSample is a TrackLocal that has a pre-set codec and accepts Samples.
 // If you wish to send a RTP Packet use TrackLocalStaticRTP.
 type TrackLocalStaticSample struct {
@@ -340,7 +349,7 @@ func (s *TrackLocalStaticSample) Unbind(t TrackLocalContext) error {
 // If one PeerConnection fails the packets will still be sent to
 // all PeerConnections. The error message will contain the ID of the failed
 // PeerConnections so you can remove them.
-func (s *TrackLocalStaticSample) WriteSample(sample media.Sample) error {
+func (s *TrackLocalStaticSample) WriteSample(sample media.Sample, extensions ...SampleRTPExtension) error {
 	s.rtpTrack.mu.RLock()
 	packetizer := s.packetizer
 	clockRate := s.clockRate
@@ -377,6 +386,11 @@ func (s *TrackLocalStaticSample) WriteSample(sample media.Sample) error {
 
 	writeErrs := []error{}
 	for _, p := range packets {
+		for _, e := range extensions {
+			if err := p.SetExtension(e.ID, e.Payload); err != nil {
+				writeErrs = append(writeErrs, err)
+			}
+		}
 		if err := s.rtpTrack.WriteRTP(p); err != nil {
 			writeErrs = append(writeErrs, err)
 		}

track_local_static_test.go

@@ -1027,3 +1027,105 @@ func (p *countingPacketizer) EnableAbsSendTime(value int)                  {}
 func (p *countingPacketizer) SkipSamples(skippedSamples uint32) {
 	p.totalSamples += uint64(skippedSamples)
 }
+
+type sampleExtensionChecker struct {
+	dummyWriter
+
+	t         *testing.T
+	errorTest bool
+	called    atomic.Uint32
+}
+
+func (c *sampleExtensionChecker) WriteRTP(header *rtp.Header, payload []byte) (int, error) {
+	c.called.Add(1)
+	assert.True(c.t, header.Extension)
+	assert.EqualValues(c.t, "hello", header.GetExtension(1))
+	if !c.errorTest {
+		assert.EqualValues(c.t, "world", header.GetExtension(2))
+	} else {
+		assert.Nil(c.t, header.GetExtension(2))
+	}
+
+	return 0, nil
+}
+
+func TestTrackLocalStaticSample_WriteSample_Extensions(t *testing.T) {
+	track, err := NewTrackLocalStaticSample(
+		RTPCodecCapability{MimeType: MimeTypeVP8},
+		"video",
+		"pion",
+	)
+	require.NoError(t, err)
+
+	checker := &sampleExtensionChecker{
+		t: t,
+	}
+
+	track.rtpTrack.mu.Lock()
+	track.rtpTrack.bindings = []trackBinding{{
+		id:          "b1",
+		ssrc:        0x1234,
+		payloadType: 96,
+		writeStream: checker,
+	}}
+	fp := &fakePacketizer{}
+	track.packetizer = fp
+	track.rtpTrack.mu.Unlock()
+
+	sample := media.Sample{
+		Data:     []byte("hi"),
+		Duration: 20 * time.Millisecond,
+	}
+	extension1 := SampleRTPExtension{
+		ID:      1,
+		Payload: []byte("hello"),
+	}
+	extension2 := SampleRTPExtension{
+		ID:      2,
+		Payload: []byte("world"),
+	}
+	err = track.WriteSample(sample, extension1, extension2)
+	require.NoError(t, err)
+	require.EqualValues(t, 2, checker.called.Load())
+}
+
+func TestTrackLocalStaticSample_WriteSample_ExtensionsError(t *testing.T) {
+	track, err := NewTrackLocalStaticSample(
+		RTPCodecCapability{MimeType: MimeTypeVP8},
+		"video",
+		"pion",
+	)
+	require.NoError(t, err)
+
+	checker := &sampleExtensionChecker{
+		t:         t,
+		errorTest: true,
+	}
+
+	track.rtpTrack.mu.Lock()
+	track.rtpTrack.bindings = []trackBinding{{
+		id:          "b1",
+		ssrc:        0x1234,
+		payloadType: 96,
+		writeStream: checker,
+	}}
+	fp := &fakePacketizer{}
+	track.packetizer = fp
+	track.rtpTrack.mu.Unlock()
+
+	sample := media.Sample{
+		Data:     []byte("hi"),
+		Duration: 20 * time.Millisecond,
+	}
+	extension1 := SampleRTPExtension{
+		ID:      1,
+		Payload: []byte("hello"),
+	}
+	extension2 := SampleRTPExtension{
+		ID:      2,
+		Payload: []byte("this is a long extension payload that will trigger an error"),
+	}
+	err = track.WriteSample(sample, extension1, extension2)
+	assert.ErrorContains(t, err, "one byte extension")
+	require.EqualValues(t, 2, checker.called.Load())
+}