Implement RTCP forwarding

Signed-off-by: Šimon Brandner <[email protected]>

Author: SimonBrandner

pkg/conference/matrix.go

@@ -65,7 +65,7 @@ func (c *Conference) onNewParticipant(participantID ParticipantID, inviteEvent *
 			logger:          logger,
 			remoteSessionID: inviteEvent.SenderSessionID,
 			streamMetadata:  inviteEvent.SDPStreamMetadata,
-			publishedTracks: make(map[event.SFUTrackDescription]*webrtc.TrackLocalStaticRTP),
+			publishedTracks: make(map[event.SFUTrackDescription]PublishedTrack),
 		}
 
 		c.participants[participantID] = participant

pkg/conference/participant.go

@@ -2,6 +2,7 @@ package conference
 
 import (
 	"encoding/json"
+	"sync/atomic"
 
 	"github.com/matrix-org/waterfall/pkg/peer"
 	"github.com/matrix-org/waterfall/pkg/signaling"
@@ -19,14 +20,19 @@ type ParticipantID struct {
 	CallID   string
 }
 
+type PublishedTrack struct {
+	Track            *webrtc.TrackLocalStaticRTP
+	LastPLITimestamp atomic.Int64
+}
+
 // Participant represents a participant in the conference.
 type Participant struct {
 	id              ParticipantID
 	logger          *logrus.Entry
 	peer            *peer.Peer[ParticipantID]
 	remoteSessionID id.SessionID
 	streamMetadata  event.CallSDPStreamMetadata
-	publishedTracks map[event.SFUTrackDescription]*webrtc.TrackLocalStaticRTP
+	publishedTracks map[event.SFUTrackDescription]PublishedTrack
 }
 
 func (p *Participant) asMatrixRecipient() signaling.MatrixRecipient {

pkg/conference/processor.go

@@ -67,7 +67,7 @@ func (c *Conference) processPeerMessage(message common.Message[ParticipantID, pe
 			return
 		}
 
-		participant.publishedTracks[key] = msg.Track
+		participant.publishedTracks[key] = PublishedTrack{Track: msg.Track}
 		c.resendMetadataToAllExcept(participant.id)
 
 	case peer.PublishedTrackFailed:
@@ -129,6 +129,22 @@ func (c *Conference) processPeerMessage(message common.Message[ParticipantID, pe
 			Op:       event.SFUOperationMetadata,
 			Metadata: c.getAvailableStreamsFor(participant.id),
 		})
+	case peer.ForwardRTCP:
+		for _, participant := range c.participants {
+			for _, publishedTrack := range participant.publishedTracks {
+				if publishedTrack.Track.StreamID() == msg.StreamID && publishedTrack.Track.ID() == msg.TrackID {
+					participant.peer.WriteRTCP(msg.Packets, msg.StreamID, msg.TrackID, publishedTrack.LastPLITimestamp.Load())
+				}
+			}
+		}
+	case peer.PLISent:
+		for _, participant := range c.participants {
+			for _, publishedTrack := range participant.publishedTracks {
+				if publishedTrack.Track.StreamID() == msg.StreamID && publishedTrack.Track.ID() == msg.TrackID {
+					publishedTrack.LastPLITimestamp.Store(msg.Timestamp)
+				}
+			}
+		}
 
 	default:
 		c.logger.Errorf("Unknown message type: %T", msg)

pkg/conference/state.go

@@ -6,7 +6,6 @@ import (
 	"github.com/matrix-org/waterfall/pkg/signaling"
 	"github.com/pion/webrtc/v3"
 	"github.com/sirupsen/logrus"
-	"golang.org/x/exp/maps"
 	"maunium.net/go/mautrix/event"
 )
 
@@ -56,7 +55,10 @@ func (c *Conference) removeParticipant(participantID ParticipantID) {
 	c.resendMetadataToAllExcept(participantID)
 
 	// Remove the participant's tracks from all participants who might have subscribed to them.
-	obsoleteTracks := maps.Values(participant.publishedTracks)
+	obsoleteTracks := []*webrtc.TrackLocalStaticRTP{}
+	for _, publishedTrack := range participant.publishedTracks {
+		obsoleteTracks = append(obsoleteTracks, publishedTrack.Track)
+	}
 	for _, otherParticipant := range c.participants {
 		otherParticipant.peer.UnsubscribeFrom(obsoleteTracks)
 	}
@@ -72,7 +74,7 @@ func (c *Conference) getAvailableStreamsFor(forParticipant ParticipantID) event.
 			// Now, find out which of published tracks belong to the streams for which we have metadata
 			// available and construct a metadata map for a given participant based on that.
 			for _, track := range participant.publishedTracks {
-				trackID, streamID := track.ID(), track.StreamID()
+				trackID, streamID := track.Track.ID(), track.Track.StreamID()
 
 				if metadata, ok := streamsMetadata[streamID]; ok {
 					metadata.Tracks[trackID] = event.CallSDPStreamMetadataTrack{}
@@ -97,7 +99,7 @@ func (c *Conference) getTracks(identifiers []event.SFUTrackDescription) []*webrt
 		// Check if this participant has any of the tracks that we're looking for.
 		for _, identifier := range identifiers {
 			if track, ok := participant.publishedTracks[identifier]; ok {
-				tracks = append(tracks, track)
+				tracks = append(tracks, track.Track)
 			}
 		}
 	}

pkg/peer/messages.go

@@ -1,6 +1,7 @@
 package peer
 
 import (
+	"github.com/pion/rtcp"
 	"github.com/pion/webrtc/v3"
 	"maunium.net/go/mautrix/event"
 )
@@ -38,3 +39,15 @@ type DataChannelMessage struct {
 }
 
 type DataChannelAvailable struct{}
+
+type ForwardRTCP struct {
+	Packets  []rtcp.Packet
+	StreamID string
+	TrackID  string
+}
+
+type PLISent struct {
+	Timestamp int64
+	StreamID  string
+	TrackID   string
+}

pkg/peer/peer.go

@@ -2,10 +2,12 @@ package peer
 
 import (
 	"errors"
+	"io"
 	"sync"
 	"time"
 
 	"github.com/matrix-org/waterfall/pkg/common"
+	"github.com/pion/rtcp"
 	"github.com/pion/webrtc/v3"
 	"github.com/sirupsen/logrus"
 	"maunium.net/go/mautrix/event"
@@ -22,6 +24,8 @@ var (
 	ErrCantSubscribeToTrack       = errors.New("can't subscribe to track")
 )
 
+const minimalPLIInterval = time.Millisecond * 500
+
 // A wrapped representation of the peer connection (single peer in the call).
 // The peer gets information about the things happening outside via public methods
 // and informs the outside world about the things happening inside the peer by posting
@@ -98,17 +102,69 @@ func (p *Peer[ID]) SubscribeTo(track *webrtc.TrackLocalStaticRTP) error {
 	// Before these packets are returned they are processed by interceptors. For things
 	// like NACK this needs to be called.
 	go func() {
-		rtcpBuf := make([]byte, 1500)
 		for {
-			if _, _, rtcpErr := rtpSender.Read(rtcpBuf); rtcpErr != nil {
-				return
+			packets, _, err := rtpSender.ReadRTCP()
+			if err != nil {
+				if errors.Is(err, io.ErrClosedPipe) || errors.Is(err, io.EOF) {
+					return
+				}
+
+				p.logger.WithError(err).Warn("failed to read RTCP on track")
 			}
+
+			p.sink.Send(ForwardRTCP{Packets: packets, TrackID: track.ID(), StreamID: track.StreamID()})
 		}
 	}()
 
 	return nil
 }
 
+func (p *Peer[ID]) WriteRTCP(packets []rtcp.Packet, streamID string, trackID string, lastPLITimestamp int64) {
+	packetsToSend := []rtcp.Packet{}
+	var mediaSSRC uint32
+	for _, receiver := range p.peerConnection.GetReceivers() {
+		if receiver.Track().ID() == trackID && receiver.Track().StreamID() == streamID {
+			mediaSSRC = uint32(receiver.Track().SSRC())
+			break
+		}
+	}
+
+	for _, packet := range packets {
+		switch typedPacket := packet.(type) {
+		// We mung the packets here, so that the SSRCs match what the
+		// receiver expects:
+		// The media SSRC is the SSRC of the media about which the packet is
+		// reporting; therefore, we mung it to be the SSRC of the publishing
+		// participant's track. Without this, it would be SSRC of the SFU's
+		// track which isn't right
+		case *rtcp.PictureLossIndication:
+			// Since we sometimes spam the sender with PLIs, make sure we don't send
+			// them way too often
+			if time.Now().UnixNano()-lastPLITimestamp < minimalPLIInterval.Nanoseconds() {
+				continue
+			}
+
+			p.sink.Send(PLISent{Timestamp: time.Now().UnixNano(), StreamID: streamID, TrackID: trackID})
+
+			typedPacket.MediaSSRC = mediaSSRC
+			packetsToSend = append(packetsToSend, typedPacket)
+		case *rtcp.FullIntraRequest:
+			typedPacket.MediaSSRC = mediaSSRC
+			packetsToSend = append(packetsToSend, typedPacket)
+		}
+
+		packetsToSend = append(packetsToSend, packet)
+	}
+
+	if len(packetsToSend) != 0 {
+		if err := p.peerConnection.WriteRTCP(packetsToSend); err != nil {
+			if !errors.Is(err, io.ErrClosedPipe) {
+				p.logger.WithError(err).Warn("failed to write RTCP on track")
+			}
+		}
+	}
+}
+
 // Unsubscribes from the given list of tracks.
 func (p *Peer[ID]) UnsubscribeFrom(tracks []*webrtc.TrackLocalStaticRTP) {
 	// That's unfortunately an O(m*n) operation, but we don't expect the number of tracks to be big.