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,lll
 	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
 	}
@@ -1802,6 +1801,8 @@ func (pc *PeerConnection) handleIncomingSSRC(rtpStream io.Reader, ssrc SSRC) err
 		return err
 	}
 
+	peekedPackets := []*peekedPacket{}
+
 	// if the first packet didn't contain simuilcast IDs, then probe more packets
 	var paddingOnly bool
 	for readCount := 0; readCount <= simulcastProbeCount; readCount++ {
@@ -1811,11 +1812,16 @@ func (pc *PeerConnection) handleIncomingSSRC(rtpStream io.Reader, ssrc SSRC) err
 				readCount--
 			}
 
-			i, _, err := interceptor.Read(b, nil)
+			i, attributes, err := interceptor.Read(b, nil)
 			if err != nil {
 				return err
 			}
 
+			peekedPackets = append(peekedPackets, &peekedPacket{
+				payload:    slices.Clone(b[:i]),
+				attributes: attributes,
+			})
+
 			if paddingOnly, err = handleUnknownRTPPacket(
 				b[:i], uint8(midExtensionID), //nolint:gosec // G115
 				uint8(streamIDExtensionID),       //nolint:gosec // G115
@@ -1851,6 +1857,7 @@ func (pc *PeerConnection) handleIncomingSSRC(rtpStream io.Reader, ssrc SSRC) err
 				interceptor,
 				rtcpReadStream,
 				rtcpInterceptor,
+				peekedPackets,
 			)
 			if err != nil {
 				return err
@@ -1930,7 +1937,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,105 @@ 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
+	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))
+
+	expectedBuffer := make([]byte, outboundMTU*rtpPktCount)
+	_, err = rand.Read(expectedBuffer)
+	assert.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	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...)
+		}
+
+		assert.Equal(t, actualBuffer, expectedBuffer)
+		cancel()
+	})
+
+	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++ {
+		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(vp8WriterA.RID())))
+		}
+
+		offset := int(sequenceNumber) * outboundMTU
+		pkt.Payload = expectedBuffer[offset : offset+outboundMTU]
+		assert.NoError(t, vp8WriterA.WriteRTP(pkt))
+	}
+
+	<-ctx.Done()
+	assert.NoError(t, wan.Stop())
+	closePairNow(t, pcOffer, pcAnswer)
+}

rtpreceiver.go

@@ -540,6 +540,7 @@ func (r *RTPReceiver) receiveForRid(
 	rtpInterceptor interceptor.RTPReader,
 	rtcpReadStream *srtp.ReadStreamSRTCP,
 	rtcpInterceptor interceptor.RTCPReader,
+	peekedPackets []*peekedPacket,
 ) (*TrackRemote, error) {
 	r.mu.Lock()
 	defer r.mu.Unlock()
@@ -551,6 +552,7 @@ func (r *RTPReceiver) receiveForRid(
 			r.tracks[i].track.codec = params.Codecs[0]
 			r.tracks[i].track.params = params
 			r.tracks[i].track.ssrc = SSRC(streamInfo.SSRC)
+			r.tracks[i].track.peekedPackets = peekedPackets
 			r.tracks[i].track.mu.Unlock()
 
 			r.tracks[i].streamInfo = streamInfo

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,18 @@ 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)
+		err = t.checkAndUpdateTrack(b)
 
-			return n, attributes, err
-		}
+		return n, peekedPkt.attributes, err
 	}
 
 	// If there's a separate RTX track and an RTX packet is available, return that
@@ -212,8 +210,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