Rendezvous/src/FinalFactory.Rendezvous.Server/Transport/UdpMediatorService.cs

using System.Diagnostics;
using System.Net;
using System.Net.Sockets;
using FinalFactory.Rendezvous.Contracts;
using LiteNetLib;
using LiteNetLib.Layers;
using Microsoft.Extensions.Options;

namespace FinalFactory.Rendezvous.Server.Transport;

internal sealed partial class UdpMediatorService : BackgroundService
{
    private readonly ILogger<UdpMediatorService> _logger;
    private readonly UdpMediatorOptions _options;
    private readonly NatMediationProcessor _processor;
    private LiteNetManager? _manager;
    private LiteNetIntroductionSink? _introductionSink;

    public UdpMediatorService(
        IOptions<UdpMediatorOptions> options,
        ILogger<UdpMediatorService> logger,
        NatMediationProcessor processor)
    {
        _options = options.Value;
        _logger = logger;
        _processor = processor;
    }

    public IPEndPoint? LocalEndpoint { get; private set; }
    public IPEndPoint? LocalIpv6Endpoint { get; private set; }

    public override Task StartAsync(CancellationToken cancellationToken)
    {
        cancellationToken.ThrowIfCancellationRequested();
        if (_manager is not null)
        {
            throw new InvalidOperationException("The UDP mediator is already running.");
        }

        IPAddress listenAddress = IPAddress.Parse(_options.ListenAddress);
        if (listenAddress.AddressFamily != AddressFamily.InterNetwork)
        {
            throw new InvalidOperationException("The required UDP listen address must be IPv4.");
        }

        IPAddress? ipv6ListenAddress = string.IsNullOrWhiteSpace(_options.Ipv6ListenAddress)
            ? null
            : IPAddress.Parse(_options.Ipv6ListenAddress);
        if (ipv6ListenAddress is not null
            && ipv6ListenAddress.AddressFamily != AddressFamily.InterNetworkV6)
        {
            throw new InvalidOperationException("The optional UDP IPv6 listen address must be IPv6.");
        }

        EventBasedLiteNetListener listener = new();
        RendezvousPacketLayer packetLayer = new(_processor);
        LiteNetManager manager = new(listener, packetLayer)
        {
            NatPunchEnabled = true,
            IPv6Enabled = ipv6ListenAddress is not null,
            UnsyncedEvents = true,
            MaxPacketPerManualReceive = _options.MaxDatagramsPerPoll,
        };
        manager.NatPunchModule.UnsyncedEvents = true;
        _introductionSink = new(manager.NatPunchModule);
        packetLayer.Attach(_introductionSink);

        if (!manager.StartInManualMode(
                listenAddress,
                ipv6ListenAddress ?? IPAddress.IPv6Any,
                _options.Port))
        {
            _introductionSink = null;
            manager.Stop();
            throw new InvalidOperationException("The UDP mediator could not bind its LiteNetLib socket.");
        }

        _manager = manager;
        LocalEndpoint = new(listenAddress, manager.LocalPort);
        LocalIpv6Endpoint = ipv6ListenAddress is null
            ? null
            : new(ipv6ListenAddress, manager.LocalPort);
        LogMediatorListening(_logger, listenAddress, manager.LocalPort);
        return base.StartAsync(cancellationToken);
    }

    public override async Task StopAsync(CancellationToken cancellationToken)
    {
        await base.StopAsync(cancellationToken).ConfigureAwait(false);
        StopManager();
        LogMediatorStopped(_logger);
    }

    public override void Dispose()
    {
        StopManager();
        base.Dispose();
    }

    protected override async Task ExecuteAsync(CancellationToken stoppingToken)
    {
        LiteNetManager manager = _manager
            ?? throw new InvalidOperationException("The UDP mediator socket was not initialized.");
        long previous = Stopwatch.GetTimestamp();
        try
        {
            while (!stoppingToken.IsCancellationRequested)
            {
                manager.PollEvents();
                manager.NatPunchModule.PollEvents();
                long current = Stopwatch.GetTimestamp();
                manager.ManualUpdate((float)Stopwatch.GetElapsedTime(previous, current).TotalMilliseconds);
                previous = current;
                await Task.Delay(_options.PollIntervalMilliseconds, stoppingToken).ConfigureAwait(false);
            }
        }
        catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
        {
        }
        finally
        {
            LocalEndpoint = null;
            LocalIpv6Endpoint = null;
        }
    }

    private void StopManager()
    {
        LiteNetManager? manager = Interlocked.Exchange(ref _manager, null);
        _introductionSink = null;
        LocalEndpoint = null;
        LocalIpv6Endpoint = null;
        manager?.Stop();
    }

    [LoggerMessage(
        EventId = 1,
        Level = LogLevel.Information,
        Message = "UDP mediator listening on {ListenAddress}:{ListenPort}")]
    private static partial void LogMediatorListening(
        ILogger logger,
        IPAddress listenAddress,
        int listenPort);

    [LoggerMessage(
        EventId = 2,
        Level = LogLevel.Information,
        Message = "UDP mediator stopped")]
    private static partial void LogMediatorStopped(ILogger logger);

    private sealed class LiteNetIntroductionSink(NatPunchModule module) : INatIntroductionSink
    {
        public void Introduce(NatIntroductionPlan plan) => module.NatIntroduce(
            plan.HostLocal,
            plan.HostPublic,
            plan.ClientLocal,
            plan.ClientPublic,
            plan.IntroductionToken);
    }

    private sealed class RendezvousPacketLayer(NatMediationProcessor processor) : PacketLayerBase(0)
    {
        private INatIntroductionSink? _sink;

        public void Attach(INatIntroductionSink sink) => _sink = sink;

        public override void ProcessInboundPacket(
            ref IPEndPoint endPoint,
            ref byte[] data,
            ref int length)
        {
            bool isFrozenEnvelope = length >= 2
                && data[0] == RendezvousUdpCodec.MagicFirst
                && data[1] == RendezvousUdpCodec.MagicSecond;
            INatIntroductionSink? sink = _sink;
            if (isFrozenEnvelope)
            {
                if (sink is not null)
                {
                    _ = processor.ProcessDatagram(data.AsSpan(0, length), endPoint, sink);
                }
            }
            else if (sink is not null
                     && LiteNetNatRequestCodec.TryDecode(
                         data.AsSpan(0, length),
                         out IPEndPoint? claimedLocalEndpoint,
                         out string? token)
                     && claimedLocalEndpoint is not null
                     && token is not null)
            {
                _ = processor.ProcessRequest(claimedLocalEndpoint, endPoint, token, sink);
            }

            // Every inbound packet is consumed here. NatPunchModule is used only for outbound introductions.
            Drop(ref length);
        }

        public override void ProcessOutBoundPacket(
            ref IPEndPoint endPoint,
            ref byte[] data,
            ref int offset,
            ref int length)
        {
        }

        private static void Drop(ref int length) => length = 0;
    }
}