From a62099671c47d26bae461b9ebcd7ac403ceb11d3 Mon Sep 17 00:00:00 2001
From: Sean DuBois <sean@siobud.com>
Date: Sun, 21 Apr 2024 22:55:01 -0400
Subject: [PATCH] Add test for TrackRemote and RTX Packets

Relates to #2752
---
 rtpreceiver_go_test.go | 125 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 125 insertions(+)

diff --git a/rtpreceiver_go_test.go b/rtpreceiver_go_test.go
index 911f0c83822..6a8b2b23c2d 100644
--- a/rtpreceiver_go_test.go
+++ b/rtpreceiver_go_test.go
@@ -7,11 +7,21 @@
 package webrtc
 
 import (
+	"bufio"
 	"context"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+	"strconv"
+	"strings"
 	"testing"
 	"time"
 
+	"github.com/pion/randutil"
+	"github.com/pion/rtp"
 	"github.com/pion/sdp/v3"
+	"github.com/pion/transport/v3/test"
 	"github.com/pion/webrtc/v4/pkg/media"
 	"github.com/stretchr/testify/assert"
 )
@@ -70,3 +80,118 @@ func TestSetRTPParameters(t *testing.T) {
 	assert.NoError(t, wan.Stop())
 	closePairNow(t, sender, receiver)
 }
+
+// Assert the behavior of reading a RTX with a distinct SSRC
+// All the attributes should be populated and the packet unpacked
+func Test_RTX_Read(t *testing.T) {
+	defer test.TimeOut(time.Second * 30).Stop()
+
+	var ssrc *uint32
+	ssrcLines := ""
+	rtxSsrc := randutil.NewMathRandomGenerator().Uint32()
+
+	pcOffer, pcAnswer, err := newPair()
+	assert.NoError(t, err)
+
+	track, err := NewTrackLocalStaticRTP(RTPCodecCapability{MimeType: MimeTypeVP8}, "track-id", "stream-id")
+	assert.NoError(t, err)
+
+	_, err = pcOffer.AddTrack(track)
+	assert.NoError(t, err)
+
+	rtxRead, rtxReadCancel := context.WithCancel(context.Background())
+	pcAnswer.OnTrack(func(track *TrackRemote, _ *RTPReceiver) {
+		for {
+			pkt, attributes, readRTPErr := track.ReadRTP()
+			if errors.Is(readRTPErr, io.EOF) {
+				return
+			} else if pkt.PayloadType == 0 {
+				continue
+			}
+
+			assert.NoError(t, readRTPErr)
+			assert.NotNil(t, pkt)
+			assert.Equal(t, pkt.SSRC, *ssrc)
+			assert.Equal(t, pkt.PayloadType, uint8(96))
+			assert.Equal(t, pkt.Payload, []byte{0xB, 0xA, 0xD})
+
+			rtxPayloadType := attributes.Get(AttributeRtxPayloadType)
+			rtxSequenceNumber := attributes.Get(AttributeRtxSequenceNumber)
+			rtxSSRC := attributes.Get(AttributeRtxSsrc)
+			if rtxPayloadType != nil && rtxSequenceNumber != nil && rtxSSRC != nil {
+				assert.Equal(t, rtxPayloadType, uint8(97))
+				assert.Equal(t, rtxSSRC, rtxSsrc)
+				assert.Equal(t, rtxSequenceNumber, pkt.SequenceNumber+500)
+
+				rtxReadCancel()
+			}
+		}
+	})
+
+	assert.NoError(t, signalPairWithModification(pcOffer, pcAnswer, func(offer string) (modified string) {
+		scanner := bufio.NewScanner(strings.NewReader(offer))
+		for scanner.Scan() {
+			l := scanner.Text()
+
+			if strings.HasPrefix(l, "a=ssrc") {
+				if ssrc == nil {
+					lineSplit := strings.Split(l, " ")[0]
+					parsed, atoiErr := strconv.ParseUint(strings.TrimPrefix(lineSplit, "a=ssrc:"), 10, 32)
+					assert.NoError(t, atoiErr)
+
+					parsedSsrc := uint32(parsed)
+					ssrc = &parsedSsrc
+
+					modified += fmt.Sprintf("a=ssrc-group:FID %d %d\r\n", *ssrc, rtxSsrc)
+				}
+
+				ssrcLines += l + "\n"
+			} else if ssrcLines != "" {
+				ssrcLines = strings.ReplaceAll(ssrcLines, fmt.Sprintf("%d", *ssrc), fmt.Sprintf("%d", rtxSsrc))
+				modified += ssrcLines
+				ssrcLines = ""
+			}
+
+			modified += l + "\n"
+		}
+
+		return modified
+	}))
+
+	func() {
+		for i := uint16(0); ; i++ {
+			pkt := rtp.Packet{
+				Header: rtp.Header{
+					Version:        2,
+					SSRC:           *ssrc,
+					PayloadType:    96,
+					SequenceNumber: i,
+				},
+				Payload: []byte{0xB, 0xA, 0xD},
+			}
+
+			select {
+			case <-time.After(20 * time.Millisecond):
+				// Send the original packet
+				err = track.WriteRTP(&pkt)
+				assert.NoError(t, err)
+
+				rtxPayload := []byte{0x0, 0x0, 0xB, 0xA, 0xD}
+				binary.BigEndian.PutUint16(rtxPayload[0:2], pkt.Header.SequenceNumber)
+
+				// Send the RTX
+				_, err = track.bindings[0].writeStream.WriteRTP(&rtp.Header{
+					Version:        2,
+					SSRC:           rtxSsrc,
+					PayloadType:    97,
+					SequenceNumber: i + 500,
+				}, rtxPayload)
+				assert.NoError(t, err)
+			case <-rtxRead.Done():
+				return
+			}
+		}
+	}()
+
+	closePairNow(t, pcOffer, pcAnswer)
+}