Don't drop packets when probing Simulcast

Before any packets that we read during the probe would get lost

Co-authored-by: cptpcrd <[email protected]>

Author: Sean-Der

go.mod

@@ -13,7 +13,7 @@ require (
 	github.com/pion/rtp v1.8.25
 	github.com/pion/sctp v1.8.40
 	github.com/pion/sdp/v3 v3.0.16
-	github.com/pion/srtp/v3 v3.0.8
+	github.com/pion/srtp/v3 v3.0.9
 	github.com/pion/stun/v3 v3.0.1
 	github.com/pion/transport/v3 v3.1.1
 	github.com/pion/turn/v4 v4.1.3

go.sum

@@ -59,8 +59,8 @@ github.com/pion/sctp v1.8.40 h1:bqbgWYOrUhsYItEnRObUYZuzvOMsVplS3oNgzedBlG8=
 github.com/pion/sctp v1.8.40/go.mod h1:SPBBUENXE6ThkEksN5ZavfAhFYll+h+66ZiG6IZQuzo=
 github.com/pion/sdp/v3 v3.0.16 h1:0dKzYO6gTAvuLaAKQkC02eCPjMIi4NuAr/ibAwrGDCo=
 github.com/pion/sdp/v3 v3.0.16/go.mod h1:9tyKzznud3qiweZcD86kS0ff1pGYB3VX+Bcsmkx6IXo=
-github.com/pion/srtp/v3 v3.0.8 h1:RjRrjcIeQsilPzxvdaElN0CpuQZdMvcl9VZ5UY9suUM=
-github.com/pion/srtp/v3 v3.0.8/go.mod h1:2Sq6YnDH7/UDCvkSoHSDNDeyBcFgWL0sAVycVbAsXFg=
+github.com/pion/srtp/v3 v3.0.9 h1:lRGF4G61xxj+m/YluB3ZnBpiALSri2lTzba0kGZMrQY=
+github.com/pion/srtp/v3 v3.0.9/go.mod h1:E+AuWd7Ug2Fp5u38MKnhduvpVkveXJX6J4Lq4rxUYt8=
 github.com/pion/stun/v3 v3.0.1 h1:jx1uUq6BdPihF0yF33Jj2mh+C9p0atY94IkdnW174kA=
 github.com/pion/stun/v3 v3.0.1/go.mod h1:RHnvlKFg+qHgoKIqtQWMOJF52wsImCAf/Jh5GjX+4Tw=
 github.com/pion/transport/v3 v3.1.1 h1:Tr684+fnnKlhPceU+ICdrw6KKkTms+5qHMgw6bIkYOM=

peerconnection.go

@@ -12,7 +12,6 @@ import (
 	"crypto/rand"
 	"errors"
 	"fmt"
-	"io"
 	"slices"
 	"strconv"
 	"strings"
@@ -1688,7 +1687,7 @@ func (pc *PeerConnection) handleNonMediaBandwidthProbe() {
 	}
 }
 
-func (pc *PeerConnection) handleIncomingSSRC(rtpStream io.Reader, ssrc SSRC) error { //nolint:gocyclo,gocognit,cyclop
+func (pc *PeerConnection) handleIncomingSSRC(rtpStream *srtp.ReadStreamSRTP, ssrc SSRC) error { //nolint:gocyclo,gocognit,cyclop
 	remoteDescription := pc.RemoteDescription()
 	if remoteDescription == nil {
 		return errPeerConnRemoteDescriptionNil
@@ -1725,7 +1724,7 @@ func (pc *PeerConnection) handleIncomingSSRC(rtpStream io.Reader, ssrc SSRC) err
 	// We read the RTP packet to determine the payload type
 	b := make([]byte, pc.api.settingEngine.getReceiveMTU())
 
-	i, err := rtpStream.Read(b)
+	i, err := rtpStream.Peek(b)
 	if err != nil {
 		return err
 	}
@@ -1930,7 +1929,7 @@ func (pc *PeerConnection) undeclaredRTPMediaProcessor() { //nolint:cyclop
 			continue
 		}
 
-		go func(rtpStream io.Reader, ssrc SSRC) {
+		go func(rtpStream *srtp.ReadStreamSRTP, ssrc SSRC) {
 			if err := pc.handleIncomingSSRC(rtpStream, ssrc); err != nil {
 				pc.log.Errorf(incomingUnhandledRTPSsrc, ssrc, err)
 			}

peerconnection_media_test.go

@@ -10,6 +10,7 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	"crypto/rand"
 	"errors"
 	"fmt"
 	"io"
@@ -2062,14 +2063,13 @@ func TestPeerConnection_Simulcast_RTX(t *testing.T) { //nolint:cyclop
 	assert.NotZero(t, ridID)
 	assert.NotZero(t, rsid)
 
-	err = signalPairWithModification(pcOffer, pcAnswer, func(sdp string) string {
+	assert.NoError(t, signalPairWithModification(pcOffer, pcAnswer, func(sdp string) string {
 		// Original chrome sdp contains no ssrc info https://pastebin.com/raw/JTjX6zg6
 		re := regexp.MustCompile("(?m)[\r\n]+^.*a=ssrc.*$")
 		res := re.ReplaceAllString(sdp, "")
 
 		return res
-	})
-	assert.NoError(t, err)
+	}))
 
 	// padding only packets should not affect simulcast probe
 	var sequenceNumber uint16
@@ -2493,3 +2493,113 @@ func Test_PeerConnection_RTX_E2E(t *testing.T) { //nolint:cyclop
 	closePairNow(t, pcOffer, pcAnswer)
 	assert.NoError(t, wan.Stop())
 }
+
+// Simulcast probe would lose packets. This test asserts that we don't drop
+// any packets during the probe
+func TestPeerConnection_Simulcast_Probe_PacketLoss(t *testing.T) { //nolint:cyclop
+	t.Skip()
+
+	lim := test.TimeOut(time.Second * 30)
+	defer lim.Stop()
+
+	report := test.CheckRoutines(t)
+	defer report()
+
+	const rtpPktCount = 10
+	pcOffer, pcAnswer, wan := createVNetPair(t, nil)
+
+	rids := []string{"a", "b", "c"}
+	vp8WriterA, err := NewTrackLocalStaticRTP(
+		RTPCodecCapability{MimeType: MimeTypeVP8}, "video", "pion2", WithRTPStreamID(rids[0]),
+	)
+	assert.NoError(t, err)
+
+	vp8WriterB, err := NewTrackLocalStaticRTP(
+		RTPCodecCapability{MimeType: MimeTypeVP8}, "video", "pion2", WithRTPStreamID(rids[1]),
+	)
+	assert.NoError(t, err)
+
+	vp8WriterC, err := NewTrackLocalStaticRTP(
+		RTPCodecCapability{MimeType: MimeTypeVP8}, "video", "pion2", WithRTPStreamID(rids[2]),
+	)
+	assert.NoError(t, err)
+
+	sender, err := pcOffer.AddTrack(vp8WriterA)
+	assert.NoError(t, err)
+	assert.NotNil(t, sender)
+
+	assert.NoError(t, sender.AddEncoding(vp8WriterB))
+	assert.NoError(t, sender.AddEncoding(vp8WriterC))
+
+	expectedBuffers := [][]byte{
+		make([]byte, outboundMTU*rtpPktCount),
+		make([]byte, outboundMTU*rtpPktCount),
+		make([]byte, outboundMTU*rtpPktCount),
+	}
+	for i := range expectedBuffers {
+		_, err := rand.Read(expectedBuffers[i])
+		assert.NoError(t, err)
+	}
+
+	pcAnswer.OnTrack(func(trackRemote *TrackRemote, _ *RTPReceiver) {
+		actualBuffer := []byte{}
+
+		for i := 0; i < rtpPktCount; i++ {
+			pkt, _, err := trackRemote.ReadRTP()
+			assert.NoError(t, err)
+
+			actualBuffer = append(actualBuffer, pkt.Payload...)
+			fmt.Println(pkt.SequenceNumber)
+		}
+
+		assert.Equal(t, actualBuffer, expectedBuffers[0])
+	})
+
+	var midID, ridID uint8
+	for _, extension := range sender.GetParameters().HeaderExtensions {
+		switch extension.URI {
+		case sdp.SDESMidURI:
+			midID = uint8(extension.ID) //nolint:gosec // G115
+		case sdp.SDESRTPStreamIDURI:
+			ridID = uint8(extension.ID) //nolint:gosec // G115
+		}
+	}
+	assert.NotZero(t, midID)
+	assert.NotZero(t, ridID)
+
+	assert.NoError(t, signalPairWithModification(pcOffer, pcAnswer, func(sdp string) string {
+		// Original chrome sdp contains no ssrc info https://pastebin.com/raw/JTjX6zg6
+		re := regexp.MustCompile("(?m)[\r\n]+^.*a=ssrc.*$")
+		res := re.ReplaceAllString(sdp, "")
+
+		return res
+	}))
+
+	peerConnectionConnected := untilConnectionState(PeerConnectionStateConnected, pcOffer, pcAnswer)
+	peerConnectionConnected.Wait()
+
+	for sequenceNumber := uint16(0); sequenceNumber < rtpPktCount; sequenceNumber++ {
+		for i, track := range []*TrackLocalStaticRTP{vp8WriterA} {
+			pkt := &rtp.Packet{
+				Header: rtp.Header{
+					Version:        2,
+					PayloadType:    96,
+					SequenceNumber: sequenceNumber,
+				},
+			}
+
+			// Make sure that packets for Stream received before MID/RID don't get dropped
+			if sequenceNumber > 3 {
+				assert.NoError(t, pkt.SetExtension(midID, []byte("0")))
+				assert.NoError(t, pkt.SetExtension(ridID, []byte(track.RID())))
+			}
+
+			offset := int(sequenceNumber) * outboundMTU
+			pkt.Payload = expectedBuffers[i][offset : offset+outboundMTU]
+			assert.NoError(t, track.WriteRTP(pkt))
+		}
+	}
+
+	assert.NoError(t, wan.Stop())
+	closePairNow(t, pcOffer, pcAnswer)
+}

track_local_static_test.go

@@ -827,8 +827,7 @@ func Test_TrackRemote_ReadRTP_UnmarshalError(t *testing.T) {
 	tr := newTrackRemote(RTPCodecTypeVideo, 0, 0, "", recv)
 
 	tr.mu.Lock()
-	tr.peeked = []byte{0x80, 96}
-	tr.peekedAttributes = nil
+	tr.peekedPackets = []*peekedPacket{{payload: []byte{0x80, 96}}}
 	tr.mu.Unlock()
 
 	pkt, attrs, err := tr.ReadRTP()

track_remote.go

@@ -15,6 +15,11 @@ import (
 	"github.com/pion/rtp"
 )
 
+type peekedPacket struct {
+	payload    []byte
+	attributes interceptor.Attributes
+}
+
 // TrackRemote represents a single inbound source of media.
 type TrackRemote struct {
 	mu sync.RWMutex
@@ -30,9 +35,9 @@ type TrackRemote struct {
 	params      RTPParameters
 	rid         string
 
-	receiver         *RTPReceiver
-	peeked           []byte
-	peekedAttributes interceptor.Attributes
+	receiver *RTPReceiver
+
+	peekedPackets []*peekedPacket
 
 	audioPlayoutStatsProviders []AudioPlayoutStatsProvider
 }
@@ -116,25 +121,19 @@ func (t *TrackRemote) Codec() RTPCodecParameters {
 func (t *TrackRemote) Read(b []byte) (n int, attributes interceptor.Attributes, err error) {
 	t.mu.RLock()
 	receiver := t.receiver
-	peeked := t.peeked != nil
+	var peekedPkt *peekedPacket
+	if len(t.peekedPackets) != 0 {
+		peekedPkt = t.peekedPackets[0]
+		t.peekedPackets = t.peekedPackets[1:]
+	}
 	t.mu.RUnlock()
 
-	if peeked {
-		t.mu.Lock()
-		data := t.peeked
-		attributes = t.peekedAttributes
-
-		t.peeked = nil
-		t.peekedAttributes = nil
-		t.mu.Unlock()
-		// someone else may have stolen our packet when we
-		// released the lock.  Deal with it.
-		if data != nil {
-			n = copy(b, data)
-			err = t.checkAndUpdateTrack(b)
+	if peekedPkt != nil {
+		n = copy(b, peekedPkt.payload)
+		attributes = peekedPkt.attributes
+		err = t.checkAndUpdateTrack(b)
 
-			return n, attributes, err
-		}
+		return
 	}
 
 	// If there's a separate RTX track and an RTX packet is available, return that
@@ -212,8 +211,7 @@ func (t *TrackRemote) peek(b []byte) (n int, a interceptor.Attributes, err error
 	// that case.
 	data := make([]byte, n)
 	n = copy(data, b[:n])
-	t.peeked = data
-	t.peekedAttributes = a
+	t.peekedPackets = append(t.peekedPackets, &peekedPacket{payload: data, attributes: a})
 	t.mu.Unlock()
 
 	return