From a43143e7ac79e3cbb68c65309fa0967769ac79c2 Mon Sep 17 00:00:00 2001
From: Alex Pokotilo <gjakop@gmail.com>
Date: Wed, 10 Apr 2024 14:32:42 +0600
Subject: [PATCH] Add padding support to TrackLocalStaticSample

To add padding-only samples call GeneratePadding
---
 track_local_static.go      | 25 ++++++++++++++++
 track_local_static_test.go | 61 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/track_local_static.go b/track_local_static.go
index 17ce0587af..1c06feba94 100644
--- a/track_local_static.go
+++ b/track_local_static.go
@@ -304,3 +304,28 @@ func (s *TrackLocalStaticSample) WriteSample(sample media.Sample) error {
 
 	return util.FlattenErrs(writeErrs)
 }
+
+// GeneratePadding writes padding-only samples to the TrackLocalStaticSample
+// 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) GeneratePadding(samples uint32) error {
+	s.rtpTrack.mu.RLock()
+	p := s.packetizer
+	s.rtpTrack.mu.RUnlock()
+
+	if p == nil {
+		return nil
+	}
+
+	packets := p.GeneratePadding(samples)
+
+	writeErrs := []error{}
+	for _, p := range packets {
+		if err := s.rtpTrack.WriteRTP(p); err != nil {
+			writeErrs = append(writeErrs, err)
+		}
+	}
+
+	return util.FlattenErrs(writeErrs)
+}
diff --git a/track_local_static_test.go b/track_local_static_test.go
index 5136eac67e..5c4f1ddd67 100644
--- a/track_local_static_test.go
+++ b/track_local_static_test.go
@@ -252,3 +252,64 @@ func BenchmarkTrackLocalWrite(b *testing.B) {
 		assert.NoError(b, err)
 	}
 }
+
+func Test_TrackLocalStatic_Padding(t *testing.T) {
+	mediaEngineOne := &MediaEngine{}
+	assert.NoError(t, mediaEngineOne.RegisterCodec(RTPCodecParameters{
+		RTPCodecCapability: RTPCodecCapability{MimeType: "video/VP8", ClockRate: 90000, Channels: 0, SDPFmtpLine: "", RTCPFeedback: nil},
+		PayloadType:        100,
+	}, RTPCodecTypeVideo))
+
+	mediaEngineTwo := &MediaEngine{}
+	assert.NoError(t, mediaEngineTwo.RegisterCodec(RTPCodecParameters{
+		RTPCodecCapability: RTPCodecCapability{MimeType: "video/VP8", ClockRate: 90000, Channels: 0, SDPFmtpLine: "", RTCPFeedback: nil},
+		PayloadType:        200,
+	}, RTPCodecTypeVideo))
+
+	offerer, err := NewAPI(WithMediaEngine(mediaEngineOne)).NewPeerConnection(Configuration{})
+	assert.NoError(t, err)
+
+	answerer, err := NewAPI(WithMediaEngine(mediaEngineTwo)).NewPeerConnection(Configuration{})
+	assert.NoError(t, err)
+
+	track, err := NewTrackLocalStaticSample(RTPCodecCapability{MimeType: MimeTypeVP8}, "video", "pion")
+	assert.NoError(t, err)
+
+	_, err = offerer.AddTransceiverFromKind(RTPCodecTypeVideo)
+	assert.NoError(t, err)
+
+	_, err = answerer.AddTrack(track)
+	assert.NoError(t, err)
+
+	onTrackFired, onTrackFiredFunc := context.WithCancel(context.Background())
+
+	offerer.OnTrack(func(track *TrackRemote, _ *RTPReceiver) {
+		assert.Equal(t, track.PayloadType(), PayloadType(100))
+		assert.Equal(t, track.Codec().RTPCodecCapability.MimeType, "video/VP8")
+
+		for i := 0; i < 20; i++ {
+			// Padding payload
+			p, _, e := track.ReadRTP()
+			assert.NoError(t, e)
+			assert.True(t, p.Padding)
+			assert.Equal(t, p.PaddingSize, byte(255))
+		}
+
+		onTrackFiredFunc()
+	})
+
+	assert.NoError(t, signalPair(offerer, answerer))
+
+	exit := false
+
+	for !exit {
+		select {
+		case <-time.After(1 * time.Millisecond):
+			assert.NoError(t, track.GeneratePadding(1))
+		case <-onTrackFired.Done():
+			exit = true
+		}
+	}
+
+	closePairNow(t, offerer, answerer)
+}