Rendezvous/tests/FinalFactory.Rendezvous.Tests/Server/UdpMediatorServiceTests.cs

using System.Net;
using System.Net.Sockets;
using FinalFactory.Rendezvous.Contracts;
using FinalFactory.Rendezvous.Server.Transport;
using FinalFactory.Rendezvous.Tests.JoinAttempts;
using LiteNetLib;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;

namespace FinalFactory.Rendezvous.Tests.Server;

public sealed class UdpMediatorServiceTests
{
    [Fact]
    public async Task ServiceBindsAnEphemeralLiteNetLibPortAndStopsCleanly()
    {
        using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(5));
        using JoinAttemptFixture fixture = new();
        NatMediationProcessor processor = new(
            fixture.Sessions.Store,
            fixture.Sessions.Capabilities,
            fixture.Service);
        using UdpMediatorService service = new(
            Options.Create(new UdpMediatorOptions
            {
                ListenAddress = IPAddress.Loopback.ToString(),
                Port = 0,
                MaxDatagramsPerPoll = 8,
                PollIntervalMilliseconds = 1,
            }),
            NullLogger<UdpMediatorService>.Instance,
            processor);

        await service.StartAsync(timeout.Token);

        IPEndPoint? boundEndpoint = service.LocalEndpoint;
        Assert.NotNull(boundEndpoint);
        Assert.Equal(IPAddress.Loopback, boundEndpoint.Address);
        Assert.InRange(boundEndpoint.Port, 1, 65_535);
        Assert.Null(service.LocalIpv6Endpoint);

        await service.StopAsync(timeout.Token);

        Assert.Null(service.LocalEndpoint);
    }

    [Fact]
    public async Task OptionalIpv6BindingNeverWidensTheRequiredIpv4Binding()
    {
        if (!Socket.OSSupportsIPv6)
        {
            return;
        }

        using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(5));
        using JoinAttemptFixture fixture = new();
        NatMediationProcessor processor = new(
            fixture.Sessions.Store,
            fixture.Sessions.Capabilities,
            fixture.Service);
        using UdpMediatorService service = new(
            Options.Create(new UdpMediatorOptions
            {
                ListenAddress = IPAddress.Loopback.ToString(),
                Ipv6ListenAddress = IPAddress.IPv6Loopback.ToString(),
                Port = 0,
            }),
            NullLogger<UdpMediatorService>.Instance,
            processor);

        await service.StartAsync(timeout.Token);

        Assert.Equal(IPAddress.Loopback, service.LocalEndpoint!.Address);
        Assert.Equal(IPAddress.IPv6Loopback, service.LocalIpv6Endpoint!.Address);
        Assert.Equal(service.LocalEndpoint.Port, service.LocalIpv6Endpoint.Port);
        IPAddress? otherIpv4 = Dns.GetHostAddresses(Dns.GetHostName())
            .FirstOrDefault(address =>
                address.AddressFamily == AddressFamily.InterNetwork
                && !IPAddress.IsLoopback(address));
        if (otherIpv4 is not null)
        {
            using UdpClient scopeProbe = new(new IPEndPoint(otherIpv4, service.LocalEndpoint.Port));
            Assert.Equal(otherIpv4, ((IPEndPoint)scopeProbe.Client.LocalEndPoint!).Address);
        }

        await service.StopAsync(timeout.Token);
    }

    [Theory]
    [InlineData(false)]
    [InlineData(true)]
    public async Task NativeLiteNetLibRequestsIntroduceTheAuthorizedPair(bool restartMediator)
    {
        using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(5));
        using JoinAttemptFixture fixture = new();
        (RegisterSessionResponse registration, _) = fixture.CreateHost();
        CreateJoinAttemptResponse created = fixture.Create(registration.ListingId, "native-litenet");
        HostJoinAttempt hostAttempt = fixture.Service.BrowseForHost(
            registration.ListingId,
            ContractLimits.ContractVersion,
            registration.LeaseToken,
            ContractLimits.BrowserPageMaxItems,
            null).Value!.Items.Single(item => item.AttemptId == created.AttemptId);
        NatMediationProcessor processor = new(
            fixture.Sessions.Store,
            fixture.Sessions.Capabilities,
            fixture.Service);
        EventBasedNetListener hostListener = new();
        EventBasedNetListener clientListener = new();
        NetManager host = new(hostListener) { NatPunchEnabled = true };
        NetManager client = new(clientListener) { NatPunchEnabled = true };
        EventBasedNatPunchListener hostPunch = new();
        EventBasedNatPunchListener clientPunch = new();
        List<string> hostTickets = [];
        List<string> clientTickets = [];
        hostPunch.NatIntroductionSuccess += (_, _, ticket) => hostTickets.Add(ticket);
        clientPunch.NatIntroductionSuccess += (_, _, ticket) => clientTickets.Add(ticket);
        host.NatPunchModule.Init(hostPunch);
        client.NatPunchModule.Init(clientPunch);
        UdpMediatorService? service = null;
        bool serviceStarted = false;

        try
        {
            service = CreateMediator(processor, port: 0);
            await service.StartAsync(timeout.Token);
            serviceStarted = true;
            Assert.True(host.Start(0));
            Assert.True(client.Start(0));
            if (restartMediator)
            {
                await AssertNativeIntroductionAsync(
                    service,
                    host,
                    client,
                    hostTickets,
                    clientTickets,
                    created,
                    hostAttempt,
                    expectedCount: 1,
                    cancellationToken: timeout.Token);
                created = fixture.Create(registration.ListingId, "native-litenet-after-restart");
                hostAttempt = fixture.Service.BrowseForHost(
                    registration.ListingId,
                    ContractLimits.ContractVersion,
                    registration.LeaseToken,
                    ContractLimits.BrowserPageMaxItems,
                    null).Value!.Items.Single(item => item.AttemptId == created.AttemptId);
                int boundPort = Assert.IsType<IPEndPoint>(service.LocalEndpoint).Port;
                await service.StopAsync(timeout.Token);
                serviceStarted = false;
                service.Dispose();
                service = null;
                service = CreateMediator(processor, boundPort);
                await service.StartAsync(timeout.Token);
                serviceStarted = true;
                Assert.Equal(boundPort, Assert.IsType<IPEndPoint>(service.LocalEndpoint).Port);
            }

            await AssertNativeIntroductionAsync(
                service,
                host,
                client,
                hostTickets,
                clientTickets,
                created,
                hostAttempt,
                expectedCount: restartMediator ? 2 : 1,
                cancellationToken: timeout.Token);
        }
        finally
        {
            host.Stop();
            client.Stop();
            if (service is not null)
            {
                try
                {
                    if (serviceStarted)
                    {
                        using CancellationTokenSource cleanup = new(TimeSpan.FromSeconds(5));
                        await service.StopAsync(cleanup.Token);
                    }
                }
                finally
                {
                    service.Dispose();
                }
            }
        }
    }

    private static async Task AssertNativeIntroductionAsync(
        UdpMediatorService service,
        NetManager host,
        NetManager client,
        List<string> hostTickets,
        List<string> clientTickets,
        CreateJoinAttemptResponse created,
        HostJoinAttempt hostAttempt,
        int expectedCount,
        CancellationToken cancellationToken)
    {
        IPEndPoint mediator = Assert.IsType<IPEndPoint>(service.LocalEndpoint);
        host.NatPunchModule.SendNatIntroduceRequest(
            mediator,
            NatPunchRequestTokenCodec.Encode(
                NatPunchPeerRole.Host,
                created.MediationHandle,
                hostAttempt.HostPunchCapability));
        client.NatPunchModule.SendNatIntroduceRequest(
            mediator,
            NatPunchRequestTokenCodec.Encode(
                NatPunchPeerRole.Client,
                created.MediationHandle,
                created.ClientPunchCapability));

        while ((hostTickets.Distinct(StringComparer.Ordinal).Count() < expectedCount
                || clientTickets.Distinct(StringComparer.Ordinal).Count() < expectedCount)
               && !cancellationToken.IsCancellationRequested)
        {
            host.PollEvents();
            host.NatPunchModule.PollEvents();
            client.PollEvents();
            client.NatPunchModule.PollEvents();
            await Task.Delay(5, cancellationToken);
        }

        List<string> distinctHostTickets = hostTickets.Distinct(StringComparer.Ordinal).ToList();
        List<string> distinctClientTickets = clientTickets.Distinct(StringComparer.Ordinal).ToList();
        Assert.Equal(expectedCount, distinctHostTickets.Count);
        Assert.Equal(expectedCount, distinctClientTickets.Count);
        string hostTicket = distinctHostTickets[^1];
        string clientTicket = distinctClientTickets[^1];
        Assert.Equal(hostTicket, clientTicket);
        Assert.True(NatIntroductionTokenCodec.TryDecode(
            hostTicket,
            out NatIntroductionToken? introduction));
        Assert.NotNull(introduction);
        Assert.Equal(created.AttemptId, introduction.AttemptId);
        Assert.Equal(43, introduction.ConnectionTicket.Length);
    }

    private static UdpMediatorService CreateMediator(NatMediationProcessor processor, int port) => new(
        Options.Create(new UdpMediatorOptions
        {
            ListenAddress = IPAddress.Loopback.ToString(),
            Port = port,
            MaxDatagramsPerPoll = 8,
            PollIntervalMilliseconds = 1,
        }),
        NullLogger<UdpMediatorService>.Instance,
        processor);

    [Fact]
    public async Task FrozenV1EnvelopeIsConsumedOnTheLiteNetSocketWithinAmplificationBudget()
    {
        using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(5));
        using JoinAttemptFixture fixture = new();
        (RegisterSessionResponse registration, _) = fixture.CreateHost();
        CreateJoinAttemptResponse created = fixture.Create(registration.ListingId, "v1-envelope");
        HostJoinAttempt hostAttempt = fixture.Service.BrowseForHost(
            registration.ListingId,
            ContractLimits.ContractVersion,
            registration.LeaseToken,
            ContractLimits.BrowserPageMaxItems,
            null).Value!.Items.Single(item => item.AttemptId == created.AttemptId);
        NatMediationProcessor processor = new(
            fixture.Sessions.Store,
            fixture.Sessions.Capabilities,
            fixture.Service);
        using UdpMediatorService service = new(
            Options.Create(new UdpMediatorOptions
            {
                ListenAddress = IPAddress.Loopback.ToString(),
                Port = 0,
                MaxDatagramsPerPoll = 8,
                PollIntervalMilliseconds = 1,
            }),
            NullLogger<UdpMediatorService>.Instance,
            processor);
        await service.StartAsync(timeout.Token);
        using UdpClient host = new(new IPEndPoint(IPAddress.Loopback, 0));
        using UdpClient client = new(new IPEndPoint(IPAddress.Loopback, 0));
        IPEndPoint mediator = Assert.IsType<IPEndPoint>(service.LocalEndpoint);
        byte[] hostDatagram = RendezvousUdpCodec.Encode(new PresenceDatagram
        {
            MessageType = UdpPresenceMessageType.HostPresence,
            MediationHandle = created.MediationHandle,
            AddressFamily = AddressFamilyKind.Ipv4,
            LocalAddress = "192.168.1.10",
            LocalPort = 41_000,
            Capability = hostAttempt.HostPunchCapability,
        });
        byte[] clientDatagram = RendezvousUdpCodec.Encode(new PresenceDatagram
        {
            MessageType = UdpPresenceMessageType.ClientPresence,
            MediationHandle = created.MediationHandle,
            AddressFamily = AddressFamilyKind.Ipv4,
            LocalAddress = "192.168.1.11",
            LocalPort = 42_000,
            Capability = created.ClientPunchCapability,
        });

        try
        {
            await host.SendAsync(hostDatagram, mediator, timeout.Token);
            await client.SendAsync(clientDatagram, mediator, timeout.Token);
            UdpReceiveResult hostIntroduction = await host.ReceiveAsync(timeout.Token);
            UdpReceiveResult clientIntroduction = await client.ReceiveAsync(timeout.Token);

            Assert.True(
                hostIntroduction.Buffer.Length + clientIntroduction.Buffer.Length
                <= clientDatagram.Length * 2,
                $"The completing authenticated contribution exceeded the 2.0 response-byte budget: "
                + $"responses={hostIntroduction.Buffer.Length + clientIntroduction.Buffer.Length}, "
                + $"request={clientDatagram.Length}.");
        }
        finally
        {
            await service.StopAsync(CancellationToken.None);
        }
    }

    [Fact]
    public async Task OversizedMalformedAndGameplayDatagramsReceiveNoResponse()
    {
        using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(5));
        using JoinAttemptFixture fixture = new();
        NatMediationProcessor processor = new(
            fixture.Sessions.Store,
            fixture.Sessions.Capabilities,
            fixture.Service);
        using UdpMediatorService service = new(
            Options.Create(new UdpMediatorOptions
            {
                ListenAddress = IPAddress.Loopback.ToString(),
                Port = 0,
                MaxDatagramsPerPoll = 8,
                PollIntervalMilliseconds = 1,
            }),
            NullLogger<UdpMediatorService>.Instance,
            processor);
        await service.StartAsync(timeout.Token);
        using UdpClient sender = new(new IPEndPoint(IPAddress.Loopback, 0));
        IPEndPoint mediator = Assert.IsType<IPEndPoint>(service.LocalEndpoint);
        byte[] oversized = new byte[ContractLimits.UdpDatagramMaxBytes + 1];
        oversized[0] = RendezvousUdpCodec.MagicFirst;
        oversized[1] = RendezvousUdpCodec.MagicSecond;
        byte[] gameplayPayload = [0x01, 0x02, 0x03, 0x04];
        byte[] malformedNative = [17, 0];

        try
        {
            await sender.SendAsync(oversized, mediator, timeout.Token);
            await sender.SendAsync(gameplayPayload, mediator, timeout.Token);
            await sender.SendAsync(malformedNative, mediator, timeout.Token);
            using CancellationTokenSource noResponse = new(TimeSpan.FromMilliseconds(150));
            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () =>
                await sender.ReceiveAsync(noResponse.Token));
        }
        finally
        {
            await service.StopAsync(CancellationToken.None);
        }
    }

    [Fact]
    public async Task ForgedNativeIntroductionResponseCannotReflectToPayloadEndpoint()
    {
        using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(5));
        using JoinAttemptFixture fixture = new();
        NatMediationProcessor processor = new(
            fixture.Sessions.Store,
            fixture.Sessions.Capabilities,
            fixture.Service);
        using UdpMediatorService service = new(
            Options.Create(new UdpMediatorOptions
            {
                ListenAddress = IPAddress.Loopback.ToString(),
                Port = 0,
                MaxDatagramsPerPoll = 8,
                PollIntervalMilliseconds = 1,
            }),
            NullLogger<UdpMediatorService>.Instance,
            processor);
        await service.StartAsync(timeout.Token);
        using UdpClient reflectedTarget = new(new IPEndPoint(IPAddress.Loopback, 0));
        using UdpClient responseCapture = new(new IPEndPoint(IPAddress.Loopback, 0));
        using UdpClient attacker = new(new IPEndPoint(IPAddress.Loopback, 0));
        LiteNetManager generator = new(new EventBasedLiteNetListener()) { NatPunchEnabled = true };

        try
        {
            Assert.True(generator.Start(0));
            IPEndPoint target = (IPEndPoint)reflectedTarget.Client.LocalEndPoint!;
            IPEndPoint capture = (IPEndPoint)responseCapture.Client.LocalEndPoint!;
            generator.NatPunchModule.NatIntroduce(
                target,
                new IPEndPoint(IPAddress.Loopback, 9),
                capture,
                capture,
                "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA");
            byte[] forgedResponse = (await responseCapture.ReceiveAsync(timeout.Token)).Buffer;

            await attacker.SendAsync(
                forgedResponse,
                Assert.IsType<IPEndPoint>(service.LocalEndpoint),
                timeout.Token);

            using CancellationTokenSource noReflection = new(TimeSpan.FromMilliseconds(150));
            await Assert.ThrowsAnyAsync<OperationCanceledException>(async () =>
                await reflectedTarget.ReceiveAsync(noReflection.Token));
        }
        finally
        {
            generator.Stop();
            await service.StopAsync(CancellationToken.None);
        }
    }
}