voice-cat/tests/VoiceCat.Tests/MediaRelayTests.cs

using VoiceCat.Transport;
using System.Net;
using System.Net.Sockets;
using VoiceCat.Crypto;
using VoiceCat.Protocol;
using VoiceCat.Server.Transport;
using Voicecat.V1;
using static VoiceCat.Tests.ServerTests;

namespace VoiceCat.Tests;

public sealed class MediaRelayTests
{
    [Fact]
    public async Task AutomaticModeReportsAuthenticatedSenderUplinkLossWithCap()
    {
        var clock = new ManualClock();
        await using var fixture = new ServerFixture(timeProvider: clock);
        Client admin = await ChannelManagementTests.AdminAsync(fixture);
        using (var store = new VoiceCat.Server.Data.AccountStore(Path.Combine(fixture.Directory, "voicecat.db")))
        {
            Channel lobby = store.LoadChannels().Single(channel => channel.Id == 1);
            lobby.Audio.PacketLossMode = PacketLossMode.PacketLossAutoFast;
            Assert.True(await ChannelManagementTests.ResultAsync(admin, new() { EditChannel = new() { Channel = lobby } }));
        }
        await using var alice = await VoicePeer.AttachAsync(fixture, admin);
        await using var bob = await VoicePeer.ConnectAsync(fixture, "Bob");
        StreamAnnounceResult stream = await alice.AnnounceAsync(StreamKind.StreamMic);

        await alice.SendAsync(alice.Seal(stream.Ssrc, [1]));
        await bob.ReceiveVoiceAsync();
        for (int i = 0; i < 99; i++) alice.Seal(stream.Ssrc, [2]);
        clock.Advance(TimeSpan.FromSeconds(3));
        await alice.SendAsync(alice.Seal(stream.Ssrc, [3]));

        PacketLossUpdate update = (await admin.ReadUntilAsync(envelope => envelope.PacketLossUpdate is not null)).PacketLossUpdate;
        Assert.Equal(1U, update.ChannelId);
        Assert.Equal(99U, update.MeasuredPercent);
        Assert.Equal(30U, update.AppliedPercent);
    }

    [Fact]
    public async Task DisconnectInvalidatesBothBindingAndActiveStreams()
    {
        await using var fixture = new ServerFixture();
        await using var alice = await VoicePeer.ConnectAsync(fixture, "Alice");
        await using var bob = await VoicePeer.ConnectAsync(fixture, "Bob");
        var stream = await alice.AnnounceAsync(StreamKind.StreamMic);
        alice.Client.Send(new() { Disconnect = new() });
        await bob.Client.ReadUntilAsync(e => e.UserEvent?.Kind == UserEvent.Types.Kind.Left && e.UserEvent.LeftId == alice.Client.Authentication!.Self.Id);
        await alice.SendAsync(alice.Seal(stream.Ssrc, [1]));
        await bob.AssertNoVoiceAsync();
    }

    [Fact]
    public async Task AnnounceRequiresSubscriptionAndUsesAuthoritativeMusicSettings()
    {
        await using var fixture = new ServerFixture();
        await using var client = await fixture.ConnectAsync();
        await client.LoginAsync("Alice");
        client.Send(new() { StreamAnnounce = new() { Kind = StreamKind.StreamMic } });
        Assert.False((await client.ReadUntilAsync(e => e.StreamAnnounceResult is not null)).StreamAnnounceResult.Ok);
        client.Send(new() { SubscribeVoice = new() });
        await client.ReadUntilAsync(e => e.VoiceSubscriptionResult is not null);
        client.Send(new() { JoinChannel = new() { ChannelId = 2 } });
        await client.ReadUntilAsync(e => e.JoinChannelResult is not null);
        client.Send(new() { RequestId = 21, StreamAnnounce = new() { Kind = StreamKind.StreamScreenAudio, RequestedAudio = new() { SampleRate = 8000, BitrateBps = 64000 } } });
        var announced = await client.ReadUntilAsync(e => e.StreamAnnounceResult is not null);
        Assert.Equal(21UL, announced.RequestId);
        Assert.True(announced.StreamAnnounceResult.Ok);
        Assert.Equal(64000U, announced.StreamAnnounceResult.EffectiveAudio.BitrateBps);
        Assert.Equal(48000U, announced.StreamAnnounceResult.EffectiveAudio.SampleRate);
        Assert.Equal(ChannelMode.ModeStereo, announced.StreamAnnounceResult.EffectiveAudio.Mode);
        client.Send(new() { StreamAnnounce = new() { Kind = (StreamKind)99 } });
        Assert.False((await client.ReadUntilAsync(e => e.StreamAnnounceResult is not null)).StreamAnnounceResult.Ok);
    }

    [Fact]
    public async Task EncryptedOpusIsResealedWithRecipientCountersAcrossMultipleStreamsAndSenders()
    {
        await using var fixture = new ServerFixture();
        await using var alice = await VoicePeer.ConnectAsync(fixture, "Alice");
        await using var bob = await VoicePeer.ConnectAsync(fixture, "Bob");
        await using var carol = await VoicePeer.ConnectAsync(fixture, "Carol");
        StreamAnnounceResult mic = await alice.AnnounceAsync(StreamKind.StreamMic);
        StreamAnnounceResult screen = await alice.AnnounceAsync(StreamKind.StreamScreenAudio);
        StreamAnnounceResult other = await carol.AnnounceAsync(StreamKind.StreamMic);
        Assert.NotEqual(mic.StreamId, screen.StreamId);
        Assert.NotEqual(mic.Ssrc, screen.Ssrc);
        Assert.Equal(48000U, screen.EffectiveAudio.SampleRate);
        Assert.Equal(24000U, screen.EffectiveAudio.BitrateBps);
        using var encoder = new Codec.OpusEncoder(new());
        short[] samples = Enumerable.Range(0, 960).Select(i => (short)(8000 * Math.Sin(i * 0.1))).ToArray();
        byte[] payload = new byte[4000];
        int length = encoder.Encode(samples, payload);
        payload = payload[..length];
        foreach (var (sender, stream) in new[] { (alice, mic), (carol, other), (alice, screen) })
        {
            byte[] packet = sender.Seal(stream.Ssrc, payload, 960, VoiceFrameFlags.Marker | VoiceFrameFlags.FecPresent);
            await sender.SendAsync(packet);
            var received = await bob.ReceiveVoiceAsync();
            Assert.Equal(payload, received.Payload);
            Assert.Equal(stream.Ssrc, received.Header.Ssrc);
            Assert.Equal(960U, received.Header.Timestamp);
            Assert.Equal(VoiceFrameFlags.Marker | VoiceFrameFlags.FecPresent, received.Header.Flags);
        }
        Assert.Equal(2UL, bob.LastSequence);
    }

    [Fact]
    public async Task ReplayForgeryAndSpoofedStreamsAreDroppedWithoutBreakingValidMedia()
    {
        await using var fixture = new ServerFixture();
        await using var alice = await VoicePeer.ConnectAsync(fixture, "Alice");
        await using var bob = await VoicePeer.ConnectAsync(fixture, "Bob");
        StreamAnnounceResult stream = await alice.AnnounceAsync(StreamKind.StreamMic);
        byte[] packet = alice.Seal(stream.Ssrc, [1, 2, 3]);
        await alice.SendAsync(packet);
        Assert.Equal(new byte[] { 1, 2, 3 }, (await bob.ReceiveVoiceAsync()).Payload);
        await alice.SendAsync(packet);
        byte[] forged = alice.Seal(stream.Ssrc, [4]);
        forged[^1] ^= 1;
        await alice.SendAsync(forged);
        await alice.SendAsync(alice.Seal(stream.Ssrc + 1000, [5]));
        await alice.SendAsync([1]);
        await alice.SendAsync(alice.Seal(stream.Ssrc, [6]));
        Assert.Equal(new byte[] { 6 }, (await bob.ReceiveVoiceAsync()).Payload);
        Assert.Equal(1UL, bob.LastSequence);
    }

    [Fact]
    public async Task SubscriptionChannelMovementAndStreamStopIsolateMedia()
    {
        await using var fixture = new ServerFixture();
        await using var alice = await VoicePeer.ConnectAsync(fixture, "Alice");
        await using var bob = await VoicePeer.ConnectAsync(fixture, "Bob");
        var mic = await alice.AnnounceAsync(StreamKind.StreamMic);
        await bob.SubscribeAsync(false);
        await alice.SendAsync(alice.Seal(mic.Ssrc, [1]));
        await bob.AssertNoVoiceAsync();
        await bob.SubscribeAsync(true);
        bob.Client.Send(new() { JoinChannel = new() { ChannelId = 2 } });
        Assert.True((await bob.Client.ReadUntilAsync(e => e.JoinChannelResult is not null)).JoinChannelResult.Ok);
        await alice.SendAsync(alice.Seal(mic.Ssrc, [2]));
        await bob.AssertNoVoiceAsync();
        bob.Client.Send(new() { JoinChannel = new() { ChannelId = 1 } });
        await bob.Client.ReadUntilAsync(e => e.JoinChannelResult is not null);
        alice.Client.Send(new() { StreamStop = new() { StreamId = mic.StreamId } });
        await alice.Client.ReadUntilAsync(e => e.UserEvent?.User?.Id == alice.Client.Authentication!.Self.Id && e.UserEvent.User.Streams.Count == 0);
        await alice.SendAsync(alice.Seal(mic.Ssrc, [3]));
        await bob.AssertNoVoiceAsync();
        var replacement = await alice.AnnounceAsync(StreamKind.StreamMic);
        await alice.SendAsync(alice.Seal(replacement.Ssrc, [4]));
        Assert.Equal(new byte[] { 4 }, (await bob.ReceiveVoiceAsync()).Payload);
    }

    [Fact]
    public async Task BadTokensCannotBindAndExistingBindingCannotBeStolen()
    {
        await using var fixture = new ServerFixture();
        await using var alice = await VoicePeer.ConnectAsync(fixture, "Alice");
        using var rogue = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
        rogue.Bind(new IPEndPoint(IPAddress.Loopback, 0));
        byte[] binding = new byte[VoiceFrameHeader.Size + 16];
        new VoiceFrameHeader(MediaFrameType.UdpBinding, 0, 0, 0, 0, 0).Write(binding);
        await rogue.SendToAsync(binding, SocketFlags.None, fixture.Server.MediaEndPoint);
        alice.Client.Authentication!.UdpToken.Span.CopyTo(binding.AsSpan(VoiceFrameHeader.Size));
        await rogue.SendToAsync(binding, SocketFlags.None, fixture.Server.MediaEndPoint);
        byte[] keepalive = new byte[VoiceFrameHeader.Size];
        new VoiceFrameHeader(MediaFrameType.Keepalive, 0, 0, 0, 0, 0).Write(keepalive);
        await rogue.SendToAsync(keepalive, SocketFlags.None, fixture.Server.MediaEndPoint);
        using var timeout = new CancellationTokenSource(200);
        await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await rogue.ReceiveAsync(new byte[100], SocketFlags.None, timeout.Token));
        await alice.SendAsync(keepalive);
        Assert.Equal(keepalive, await alice.ReceivePacketAsync());
    }

    internal sealed class VoicePeer : IAsyncDisposable
    {
        public Client Client { get; }
        private Socket udp = new(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
        private readonly IPEndPoint endpoint;
        private readonly MediaSessionCrypto crypto;
        public ulong LastSequence { get; private set; }
        private VoicePeer(Client client, IPEndPoint endpoint, MediaSessionCrypto crypto)
        {
            Client = client; this.endpoint = endpoint; this.crypto = crypto;
            udp.Bind(new IPEndPoint(IPAddress.Loopback, 0));
        }
        public static async Task<VoicePeer> ConnectAsync(ServerFixture fixture, string nickname)
        {
            Client client = await fixture.ConnectAsync();
            await client.LoginAsync(nickname);
            return await AttachAsync(fixture, client);
        }
        internal static async Task<VoicePeer> AttachAsync(ServerFixture fixture, Client client)
        {
            var peer = new VoicePeer(client, fixture.Server.MediaEndPoint, await client.TakeMediaCryptoAsync());
            client.Send(new() { UdpBinding = new() { UdpToken = client.Authentication!.UdpToken } });
            Assert.True((await client.ReadUntilAsync(e => e.UdpBinding is not null)).UdpBinding.Ack);
            byte[] binding = new byte[VoiceFrameHeader.Size + 16];
            new VoiceFrameHeader(MediaFrameType.UdpBinding, 0, 0, 0, 0, 0).Write(binding);
            client.Authentication.UdpToken.Span.CopyTo(binding.AsSpan(VoiceFrameHeader.Size));
            await peer.SendAsync(binding);
            byte[] keepalive = new byte[VoiceFrameHeader.Size];
            new VoiceFrameHeader(MediaFrameType.Keepalive, 0, 0, 0, 0, 0).Write(keepalive);
            await peer.SendAsync(keepalive);
            Assert.Equal(keepalive, await peer.ReceivePacketAsync());
            await peer.SubscribeAsync(true);
            return peer;
        }
        public async Task SubscribeAsync(bool subscribed)
        {
            Client.Send(subscribed ? new() { SubscribeVoice = new() } : new() { UnsubscribeVoice = new() });
            var result = (await Client.ReadUntilAsync(e => e.VoiceSubscriptionResult is not null)).VoiceSubscriptionResult;
            Assert.True(result.Ok); Assert.Equal(subscribed, result.Subscribed);
        }
        public async Task<StreamAnnounceResult> AnnounceAsync(StreamKind kind)
        {
            Client.Send(new() { StreamAnnounce = new() { Kind = kind, RequestedAudio = new() { SampleRate = 8000, BitrateBps = 900000 } } });
            var result = (await Client.ReadUntilAsync(e => e.StreamAnnounceResult is not null)).StreamAnnounceResult;
            Assert.True(result.Ok, result.Error);
            return result;
        }
        public byte[] Seal(uint ssrc, byte[] payload, uint timestamp = 0, VoiceFrameFlags flags = 0)
        {
            byte[] packet = new byte[VoiceFrameHeader.Size + payload.Length + 16];
            crypto.Encryptor.Encrypt(new(MediaFrameType.Voice, flags, 0, ssrc, 0, timestamp), payload, packet);
            return packet;
        }
        public async Task SendAsync(byte[] packet) => await udp.SendToAsync(packet, SocketFlags.None, endpoint, Client.Timeout.Token);

        // Models a Wi-Fi/cellular handover: the peer keeps its TLS control session but its media
        // source address changes, then it re-offers its UDP binding token from the new address.
        internal async Task HandoverAsync()
        {
            udp.Dispose();
            udp = new(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
            udp.Bind(new IPEndPoint(IPAddress.Loopback, 0));
            byte[] rebind = new byte[MediaEncryptor.RebindSize];
            int size = crypto.Encryptor.EncryptRebind(Client.Authentication!.UdpToken.Span, rebind);
            await SendAsync(rebind[..size]);
        }

        // A rebind captured off the wire must not let anyone else claim the peer's downlink.
        internal async Task<byte[]> CaptureRebindAsync()
        {
            byte[] rebind = new byte[MediaEncryptor.RebindSize];
            int size = crypto.Encryptor.EncryptRebind(Client.Authentication!.UdpToken.Span, rebind);
            await SendAsync(rebind[..size]);
            return rebind[..size];
        }

        // Returns true when the relay echoes a keepalive to the peer's current source address,
        // which is the only signal that the server will route downlink media back to it.
        internal async Task<bool> KeepaliveEchoesAsync(int timeoutMilliseconds = 1000)
        {
            byte[] keepalive = new byte[VoiceFrameHeader.Size];
            new VoiceFrameHeader(MediaFrameType.Keepalive, 0, 0, 0, 0, 0).Write(keepalive);
            await SendAsync(keepalive);
            using var timeout = new CancellationTokenSource(timeoutMilliseconds);
            try { byte[] buffer = new byte[65535]; await udp.ReceiveAsync(buffer, SocketFlags.None, timeout.Token); return true; }
            catch (OperationCanceledException) { return false; }
        }
        public async Task<byte[]> ReceivePacketAsync()
        {
            byte[] buffer = new byte[65535];
            int size = await udp.ReceiveAsync(buffer, SocketFlags.None, Client.Timeout.Token);
            return buffer[..size];
        }
        public async Task<(VoiceFrameHeader Header, byte[] Payload)> ReceiveVoiceAsync()
        {
            byte[] packet = await ReceivePacketAsync();
            byte[] plain = new byte[65535];
            Assert.True(crypto.Decryptor.TryDecrypt(packet, plain, out var header, out int length));
            LastSequence = header.Sequence;
            return (header, plain[..length]);
        }
        public async Task AssertNoVoiceAsync()
        {
            using var timeout = new CancellationTokenSource(200);
            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await udp.ReceiveAsync(new byte[65535], SocketFlags.None, timeout.Token));
        }
        public async ValueTask DisposeAsync() { udp.Dispose(); crypto.Dispose(); await Client.DisposeAsync(); }
    }

    private sealed class ManualClock : TimeProvider
    {
        private long timestamp;
        public override long TimestampFrequency => 1_000;
        public override long GetTimestamp() => Volatile.Read(ref timestamp);
        internal void Advance(TimeSpan duration) => Interlocked.Add(ref timestamp, (long)duration.TotalMilliseconds);
    }
}