Piik/internal/app/mediaedge/receiver_test.go

package mediaedge

import (
	"fmt"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/livekit/livekit-server/pkg/sfu"
	"github.com/pion/rtp"
	"github.com/pion/sdp/v3"
	"github.com/pion/stun/v3"
	"github.com/pion/webrtc/v4"
)

func TestOfferSendsCodecDistinguishesBundledAndRejectedMedia(t *testing.T) {
	for _, test := range []struct {
		name       string
		port       int
		attributes string
		want       bool
	}{
		{"ordinary", 9, "a=sendonly\r\n", true},
		{"bundled", 0, "a=bundle-only\r\na=sendonly\r\n", true},
		{"rejected", 0, "a=sendonly\r\n", false},
		{"bundled-inactive", 0, "a=bundle-only\r\na=inactive\r\n", false},
		{"bundled-receive-only", 0, "a=bundle-only\r\na=recvonly\r\n", false},
	} {
		t.Run(test.name, func(t *testing.T) {
			offer := "v=0\r\no=- 1 1 IN IP4 127.0.0.1\r\ns=-\r\nt=0 0\r\na=group:BUNDLE 0 1\r\n" +
				"m=video 9 UDP/TLS/RTP/SAVPF 96\r\na=mid:0\r\na=rtpmap:96 VP8/90000\r\n" +
				fmt.Sprintf("m=audio %d UDP/TLS/RTP/SAVPF 111\r\na=mid:1\r\na=rtpmap:111 opus/48000/2\r\n%s", test.port, test.attributes)
			got, err := offerSendsCodec(offer, "audio", "opus")
			if err != nil || got != test.want {
				t.Fatalf("audio=%v, wanted %v: %v", got, test.want, err)
			}
		})
	}
}

func TestReceiverAdvertisesLocalStereoPreference(t *testing.T) {
	for _, remotePreference := range []string{"", ";stereo=0", ";stereo=1"} {
		t.Run(remotePreference, func(t *testing.T) {
			engine, err := NewEngine(EngineOptions{BindAddress: "127.0.0.1:0", IncludeLoopback: true})
			if err != nil {
				t.Fatal(err)
			}
			t.Cleanup(func() { _ = engine.Close() })
			// Independent upstream capabilities: local registration must not make
			// the fixture advertise the receive preference we are checking.
			upstream, err := webrtc.NewPeerConnection(webrtc.Configuration{})
			if err != nil {
				t.Fatal(err)
			}
			t.Cleanup(func() { _ = upstream.Close() })
			for _, kind := range []webrtc.RTPCodecType{webrtc.RTPCodecTypeVideo, webrtc.RTPCodecTypeAudio} {
				transceiver, addErr := upstream.AddTransceiverFromKind(kind, webrtc.RTPTransceiverInit{Direction: webrtc.RTPTransceiverDirectionSendonly})
				if addErr != nil {
					t.Fatal(addErr)
				}
				if kind == webrtc.RTPCodecTypeAudio {
					for _, codec := range transceiver.Sender().GetParameters().Codecs {
						if codec.MimeType != webrtc.MimeTypeOpus {
							continue
						}
						codec.SDPFmtpLine = "minptime=10;useinbandfec=1" + remotePreference
						err = transceiver.SetCodecPreferences([]webrtc.RTPCodecParameters{codec})
						break
					}
				}
				if err != nil {
					t.Fatal(err)
				}
			}
			var receiver *Receiver
			for attempt := 0; attempt < 2; attempt++ {
				offer, err := upstream.CreateOffer(nil)
				if err != nil {
					t.Fatal(err)
				}
				if err = upstream.SetLocalDescription(offer); err != nil {
					t.Fatal(err)
				}
				var answer webrtc.SessionDescription
				if receiver == nil {
					receiver, answer, err = engine.NewReceiver(ReceiverOptions{Offer: offer, EdgeCapacity: 1})
					if err != nil {
						t.Fatal(err)
					}
					t.Cleanup(func() { _ = receiver.Close() })
				} else {
					var reused bool
					answer, reused, err = receiver.Renegotiate(offer, nil)
					if err != nil || !reused {
						t.Fatalf("renegotiate: reused=%v, error=%v", reused, err)
					}
				}
				var parsed sdp.SessionDescription
				if err = parsed.UnmarshalString(answer.SDP); err != nil {
					t.Fatal(err)
				}
				found := false
				for _, media := range parsed.MediaDescriptions {
					if media.MediaName.Media != "audio" {
						continue
					}
					for _, attribute := range media.Attributes {
						if attribute.Key == "fmtp" && strings.Contains(attribute.Value, "stereo=1") && strings.Contains(attribute.Value, "maxaveragebitrate=192000") {
							found = true
						}
					}
				}
				if !found {
					t.Fatalf("answer %d lost local stereo preference", attempt)
				}
				if err = upstream.SetRemoteDescription(answer); err != nil {
					t.Fatal(err)
				}
			}
		})
	}
}

func TestReceiverCodecMatchesTheSingleNegotiatedAnswer(t *testing.T) {
	for _, preferVP8 := range []bool{false, true} {
		engine, err := NewEngine(EngineOptions{BindAddress: "127.0.0.1:0", IncludeLoopback: true})
		if err != nil {
			t.Fatal(err)
		}
		t.Cleanup(func() { _ = engine.Close() })
		upstream, err := testICESocket(t, engine).newPeerConnection()
		if err != nil {
			t.Fatal(err)
		}
		t.Cleanup(func() { _ = upstream.Close() })
		transceiver, err := upstream.AddTransceiverFromKind(webrtc.RTPCodecTypeVideo,
			webrtc.RTPTransceiverInit{Direction: webrtc.RTPTransceiverDirectionSendonly})
		if err != nil {
			t.Fatal(err)
		}
		preferences := []webrtc.RTPCodecParameters{videoCodecs["h264"], videoCodecs["vp8"]}
		want := "h264"
		if preferVP8 {
			preferences[0], preferences[1] = preferences[1], preferences[0]
			want = "vp8"
		}
		if err = transceiver.SetCodecPreferences(preferences); err != nil {
			t.Fatal(err)
		}
		offer, err := upstream.CreateOffer(nil)
		if err != nil {
			t.Fatal(err)
		}
		receiver, answer, err := engine.NewReceiver(ReceiverOptions{Offer: offer, EdgeCapacity: 1})
		if err != nil {
			t.Fatal(err)
		}
		t.Cleanup(func() { _ = receiver.Close() })
		if receiver.Codec() != want || !strings.Contains(strings.ToLower(answer.SDP), want+"/90000") ||
			(strings.Contains(strings.ToLower(answer.SDP), "vp8/90000") && strings.Contains(strings.ToLower(answer.SDP), "h264/90000")) {
			t.Fatalf("receiver codec %q does not match its single-codec answer", receiver.Codec())
		}
	}
}

func TestReceiverNegotiatesForwardableH264(t *testing.T) {
	for _, test := range []struct {
		name    string
		formats []string
		want    string
	}{
		{"browser-auto", []string{"42e01f", "vp8"}, "42e01f"},
		{"safari-order", []string{"640c1f", "42e01f"}, "42e01f"},
		{"main", []string{"4d001f"}, ""},
		{"high", []string{"64001f"}, ""},
		{"baseline", []string{"42001f"}, ""},
		{"missing-profile", []string{""}, ""},
		{"equivalent-baseline", []string{"42c01f"}, "42c01f"},
		{"small-browser-output", []string{"42e00b"}, "42e00b"},
		{"level-ceiling", []string{"42e033"}, "42e033"},
		{"over-ceiling", []string{"42e034"}, ""},
		{"invalid-level", []string{"42e035"}, ""},
		{"invalid-constraints", []string{"42e11f"}, ""},
		{"mode-zero", []string{"42e01f;packetization-mode=0"}, ""},
	} {
		t.Run(test.name, func(t *testing.T) {
			engine, err := NewEngine(EngineOptions{BindAddress: "127.0.0.1:0", IncludeLoopback: true})
			if err != nil {
				t.Fatal(err)
			}
			t.Cleanup(func() { _ = engine.Close() })
			upstream, err := webrtc.NewPeerConnection(webrtc.Configuration{})
			if err != nil {
				t.Fatal(err)
			}
			t.Cleanup(func() { _ = upstream.Close() })
			transceiver, err := upstream.AddTransceiverFromKind(webrtc.RTPCodecTypeVideo,
				webrtc.RTPTransceiverInit{Direction: webrtc.RTPTransceiverDirectionSendonly})
			if err != nil {
				t.Fatal(err)
			}
			var preferences []webrtc.RTPCodecParameters
			for index, format := range test.formats {
				codec := videoCodecs["h264"]
				codec.SDPFmtpLine = "level-asymmetry-allowed=1;packetization-mode=1"
				if format != "" {
					codec.SDPFmtpLine += ";profile-level-id=" + format
				}
				if format == "vp8" {
					codec = videoCodecs["vp8"]
				}
				codec.PayloadType = webrtc.PayloadType(96 + index)
				preferences = append(preferences, codec)
			}
			if err = transceiver.SetCodecPreferences(preferences); err != nil {
				t.Fatal(err)
			}
			offer, err := upstream.CreateOffer(nil)
			if err != nil {
				t.Fatal(err)
			}
			receiver, answer, err := engine.NewReceiver(ReceiverOptions{Offer: offer, EdgeCapacity: 1})
			if receiver != nil {
				t.Cleanup(func() { _ = receiver.Close() })
			}
			if test.want == "" {
				if err == nil {
					t.Fatalf("unsupported format was admitted: %s", answer.SDP)
				}
				return
			}
			if err != nil {
				t.Fatal(err)
			}
			if receiver.Codec() != "h264" || !strings.Contains(answer.SDP, "profile-level-id="+test.want) {
				t.Fatalf("wanted %s, got %s: %s", test.want, receiver.Codec(), answer.SDP)
			}
			// A profile-only change is still a changed media shape. Reject it before
			// mutating the current peer/source, so the owner can prepare a replacement.
			previous := receiver.connection.RemoteDescription().SDP
			changed := offer
			changed.SDP = strings.ReplaceAll(changed.SDP, "profile-level-id="+test.want, "profile-level-id=4d001f")
			_, reused, renegotiateErr := receiver.Renegotiate(changed, nil)
			if renegotiateErr != nil || reused || receiver.connection.RemoteDescription().SDP != previous {
				t.Fatalf("unsupported profile altered the current source: reused=%v, error=%v", reused, renegotiateErr)
			}
			source := receiver.Source()
			_, reused, renegotiateErr = receiver.Renegotiate(offer, nil)
			if renegotiateErr != nil || !reused || receiver.Source() != source {
				t.Fatalf("compatible renegotiation replaced the source: reused=%v, error=%v", reused, renegotiateErr)
			}
		})
	}
}

func TestReceiverForwardsEncodedVideoToANativeEdge(t *testing.T) {
	for _, codec := range []string{"h264", "vp8"} {
		t.Run(codec, func(t *testing.T) { testReceiverForwarding(t, codec) })
	}
}

func testReceiverForwarding(t *testing.T, codec string) {
	engine, err := NewEngine(EngineOptions{
		BindAddress:     "127.0.0.1:0",
		IncludeLoopback: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = engine.Close() })

	upstreamEngine, err := NewEngine(EngineOptions{
		BindAddress:     "127.0.0.1:0",
		IncludeLoopback: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = upstreamEngine.Close() })
	upstream, err := testICESocket(t, upstreamEngine).newPeerConnection()
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = upstream.Close() })
	track, err := webrtc.NewTrackLocalStaticRTP(
		videoCodecs[codec].RTPCodecCapability, "screen", "upstream",
	)
	if err != nil {
		t.Fatal(err)
	}
	if _, err = upstream.AddTrack(track); err != nil {
		t.Fatal(err)
	}
	for _, transceiver := range upstream.GetTransceivers() {
		if transceiver.Kind() == webrtc.RTPCodecTypeVideo {
			if err = transceiver.SetCodecPreferences([]webrtc.RTPCodecParameters{videoCodecs[codec]}); err != nil {
				t.Fatal(err)
			}
		}
	}
	audioTrack, err := webrtc.NewTrackLocalStaticRTP(
		opusCapability, "audio", "upstream",
	)
	if err != nil {
		t.Fatal(err)
	}
	if _, err = upstream.AddTrack(audioTrack); err != nil {
		t.Fatal(err)
	}
	var nativeReceiver *Receiver
	var upstreamCandidates []*webrtc.ICECandidateInit
	var receiverCandidates []*webrtc.ICECandidateInit
	var candidateMu sync.Mutex
	remoteReady := false
	completed := make(chan struct{}, 4)
	gatheringCurrent := make(chan func() bool, 4)
	surveyCandidates := 0
	var bindingRequests atomic.Int32
	server := bindingServer(t, func(request *stun.Message, sender *net.UDPAddr, listener *net.UDPConn) {
		bindingRequests.Add(1)
		answerBinding(request, sender, listener, sender.Port)
	})
	servers := []webrtc.ICEServer{{URLs: []string{"stun:" + server.String()}}}
	upstream.OnICECandidate(func(candidate *webrtc.ICECandidate) {
		if candidate == nil {
			return
		}
		value := candidate.ToJSON()
		candidateMu.Lock()
		ready := remoteReady
		if !ready {
			upstreamCandidates = append(upstreamCandidates, &value)
		}
		candidateMu.Unlock()
		if ready {
			_ = nativeReceiver.AddRemoteCandidate(&value)
		}
	})
	upstreamGathered := webrtc.GatheringCompletePromise(upstream)
	offer, err := upstream.CreateOffer(nil)
	if err != nil {
		t.Fatal(err)
	}
	if err = upstream.SetLocalDescription(offer); err != nil {
		t.Fatal(err)
	}
	waitSignal(t, upstreamGathered, "upstream ICE gathering")

	nativeReceiver, answer, err := engine.NewReceiver(ReceiverOptions{
		Offer:        *upstream.LocalDescription(),
		ICEServers:   servers,
		EdgeCapacity: 1,
		Events: ReceiverEvents{
			LocalCandidate: func(candidate *webrtc.ICECandidateInit, current func() bool) {
				if candidate == nil {
					gatheringCurrent <- current
					completed <- struct{}{}
					return
				}
				candidateMu.Lock()
				if strings.Contains(candidate.Candidate, "candidate:ns") {
					surveyCandidates++
				}
				ready := remoteReady
				if !ready {
					receiverCandidates = append(receiverCandidates, candidate)
				}
				candidateMu.Unlock()
				if ready {
					_ = upstream.AddICECandidate(*candidate)
				}
			},
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = nativeReceiver.Close() })
	if nativeReceiver.Codec() != codec || nativeReceiver.Source().Codec() != codec {
		t.Fatalf("native receiver/source selected %q instead of %q", nativeReceiver.Codec(), codec)
	}
	if err = upstream.SetRemoteDescription(answer); err != nil {
		t.Fatal(err)
	}
	candidateMu.Lock()
	remoteReady = true
	pendingUpstream, pendingReceiver := upstreamCandidates, receiverCandidates
	upstreamCandidates, receiverCandidates = nil, nil
	candidateMu.Unlock()
	for _, candidate := range pendingUpstream {
		_ = nativeReceiver.AddRemoteCandidate(candidate)
	}
	for _, candidate := range pendingReceiver {
		_ = upstream.AddICECandidate(*candidate)
	}
	waitConnected(t, upstream, "native receiver")

	downstream, packets, audioPackets, err := connectedReceiverForSources(
		t, engine, nativeReceiver.Source(), nativeReceiver.AudioSource(), "native-relay-edge",
	)
	t.Cleanup(func() { _ = downstream.Close() })
	if err != nil {
		t.Fatal(err)
	}
	want := &rtp.Packet{
		Header: rtp.Header{
			Version:        2,
			PayloadType:    uint8(videoCodecs[codec].PayloadType),
			SequenceNumber: 7,
			Timestamp:      90_000,
			SSRC:           42,
			Marker:         true,
		},
		Payload: []byte{0x65, 0x88, 0x84, 0x00},
	}
	if codec == "vp8" {
		want.Payload = []byte{0x10, 0x10, 0, 0, 0x9d, 0x01, 0x2a, 0x80, 0x02, 0xe0, 0x01, 0}
	} else {
		want.Payload = []byte{0x78}
		for _, nalu := range sfu.H264KeyFrame2x2 {
			want.Payload = append(want.Payload, byte(len(nalu)>>8), byte(len(nalu)))
			want.Payload = append(want.Payload, nalu...)
		}
	}
	if err = want.SetExtension(9, []byte{0xde, 0xad}); err != nil {
		t.Fatal(err)
	}
	if err = track.WriteRTP(want); err != nil {
		t.Fatal(err)
	}
	got := waitPacket(t, packets)
	if string(got.Payload) != string(want.Payload) || !got.Marker ||
		got.GetExtension(9) != nil {
		t.Fatalf("forwarded packet = %#v", got)
	}
	wantAudio := &rtp.Packet{
		Header: rtp.Header{
			Version: 2, PayloadType: 111, SequenceNumber: 8,
			Timestamp: 960, SSRC: 43, Marker: true,
		},
		Payload: []byte{0xf8, 0xff, 0xfe},
	}
	if err = wantAudio.SetExtension(9, []byte{0xbe, 0xef}); err != nil {
		t.Fatal(err)
	}
	if err = audioTrack.WriteRTP(wantAudio); err != nil {
		t.Fatal(err)
	}
	gotAudio := waitPacket(t, audioPackets)
	if string(gotAudio.Payload) != string(wantAudio.Payload) ||
		gotAudio.GetExtension(9) != nil {
		t.Fatalf("forwarded audio packet = %#v", gotAudio)
	}
	for index, input := range []struct {
		track   *webrtc.TrackLocalStaticRTP
		media   *rtp.Packet
		packets <-chan *rtp.Packet
	}{{track, want, packets}, {audioTrack, wantAudio, audioPackets}} {
		padding := &rtp.Packet{Header: input.media.Header, PaddingSize: 4}
		padding.Padding = true
		padding.SequenceNumber++
		padding.Marker = false
		if err = input.track.WriteRTP(padding); err != nil {
			t.Fatal(err)
		}
		if index == 1 {
			if packet := waitPacket(t, input.packets); !packet.Padding || packet.PaddingSize != padding.PaddingSize || len(packet.Payload) != 0 {
				t.Fatal("audio lost RTP padding continuity")
			}
		}
		next := input.media.Clone()
		next.SequenceNumber += 2
		next.Timestamp += 3_000
		if err = input.track.WriteRTP(next); err != nil {
			t.Fatal(err)
		}
		packet := waitPacket(t, input.packets)
		if index == 0 && packet.SequenceNumber != got.SequenceNumber+1 {
			t.Fatal("video projection did not close the filtered padding sequence gap")
		}
		if string(packet.Payload) != string(next.Payload) {
			t.Fatalf("stream %d stopped delivering media after padding", index)
		}
	}
	waitSignal(t, completed, "initial receiver gathering")
	firstGatheringCurrent := <-gatheringCurrent
	if !firstGatheringCurrent() {
		t.Fatal("current receiver gathering lost delivery authority")
	}
	defer func() {
		_ = nativeReceiver.Close()
		if nativeReceiver.active() || firstGatheringCurrent() {
			t.Error("retired receiver still authorizes queued events")
		}
	}()
	source, audio := nativeReceiver.Source(), nativeReceiver.AudioSource()
	socket := nativeReceiver.socket
	for _, restart := range []bool{false, true} {
		previousGathering := nativeReceiver.localCandidates
		candidateMu.Lock()
		remoteReady = false
		candidateMu.Unlock()
		offer, offerErr := upstream.CreateOffer(&webrtc.OfferOptions{ICERestart: restart})
		if offerErr != nil {
			t.Fatal(offerErr)
		}
		if err = upstream.SetLocalDescription(offer); err != nil {
			t.Fatal(err)
		}
		answer, reused, renegotiateErr := nativeReceiver.Renegotiate(offer, servers)
		if renegotiateErr != nil || !reused {
			t.Fatalf("receiver renegotiation: reused=%v, %v", reused, renegotiateErr)
		}
		if err = upstream.SetRemoteDescription(answer); err != nil {
			t.Fatal(err)
		}
		candidateMu.Lock()
		remoteReady = true
		pendingUpstream, pendingReceiver = upstreamCandidates, receiverCandidates
		upstreamCandidates, receiverCandidates = nil, nil
		candidateMu.Unlock()
		for _, candidate := range pendingUpstream {
			_ = nativeReceiver.AddRemoteCandidate(candidate)
		}
		for _, candidate := range pendingReceiver {
			_ = upstream.AddICECandidate(*candidate)
		}
		waitConnected(t, upstream, "renegotiated receiver")
		if (previousGathering != nativeReceiver.localCandidates) != restart {
			t.Fatalf("collector replaced=%v, ICE restart=%v", previousGathering != nativeReceiver.localCandidates, restart)
		}
		if restart {
			waitSignal(t, completed, "restarted receiver gathering")
			if firstGatheringCurrent() || !(<-gatheringCurrent)() {
				t.Fatal("queued candidates did not retain their gathering owner")
			}
			previousGathering.addPion(nil)
			select {
			case <-completed:
				t.Fatal("retired collector emitted completion")
			default:
			}
			candidateMu.Lock()
			count := surveyCandidates
			candidateMu.Unlock()
			if count != 2 {
				t.Fatalf("STUN survey candidates=%d, want one per gathering", count)
			}
			if bindingRequests.Load() < 2 {
				t.Fatal("receiver ICE restart reused a completed STUN observation")
			}
		} else if bindingRequests.Load() != 1 {
			t.Fatal("ordinary SDP renegotiation restarted STUN discovery")
		}
		if nativeReceiver.socket != socket || nativeReceiver.Source() != source || nativeReceiver.AudioSource() != audio ||
			downstream.connection.ConnectionState() != webrtc.PeerConnectionStateConnected {
			t.Fatal("renegotiation retired the source or its healthy downstream")
		}
		for _, input := range []struct {
			track   *webrtc.TrackLocalStaticRTP
			media   *rtp.Packet
			packets <-chan *rtp.Packet
		}{{track, want, packets}, {audioTrack, wantAudio, audioPackets}} {
			input.media.SequenceNumber += 10
			input.media.Timestamp += 90_000
			if err = input.track.WriteRTP(input.media); err != nil {
				t.Fatal(err)
			}
			if packet := waitPacket(t, input.packets); string(packet.Payload) != string(input.media.Payload) {
				t.Fatal("media did not resume on the retained downstream")
			}
		}
	}
}

func connectedReceiverForSources(
	t *testing.T,
	engine *Engine,
	source *Source,
	audio *AudioSource,
	connectionID string,
) (*Edge, <-chan *rtp.Packet, <-chan *rtp.Packet, error) {
	options := EdgeOptions{ConnectionID: connectionID, Audio: audio}
	edge, err := engine.NewEdge(source, options)
	if err != nil {
		return nil, nil, nil, err
	}
	receiver := newReceiverWithAudio(t, audio != nil, "127.0.0.1:0")
	t.Cleanup(func() { _ = receiver.Close() })
	packets := make(chan *rtp.Packet, 4)
	audioPackets := make(chan *rtp.Packet, 4)
	receiver.OnTrack(func(track *webrtc.TrackRemote, _ *webrtc.RTPReceiver) {
		go func() {
			for {
				packet, _, readErr := track.ReadRTP()
				if readErr != nil {
					return
				}
				if track.Kind() == webrtc.RTPCodecTypeAudio {
					audioPackets <- packet
				} else {
					packets <- packet
				}
			}
		}()
	})
	connectEdgeToReceiver(t, edge, receiver)
	return edge, packets, audioPackets, nil
}

func waitConnected(t *testing.T, connection *webrtc.PeerConnection, label string) {
	deadline := time.Now().Add(5 * time.Second)
	for connection.ConnectionState() != webrtc.PeerConnectionStateConnected {
		if time.Now().After(deadline) {
			t.Fatalf("%s state = %s", label, connection.ConnectionState())
		}
		time.Sleep(10 * time.Millisecond)
	}
}