Rendezvous/src/FinalFactory.Rendezvous.TestClient/DirectEchoProtocol.cs

using System.Security.Cryptography;
using System.Text;
using LiteNetLib;
using LiteNetLib.Utils;

namespace FinalFactory.Rendezvous.TestClient;

internal sealed class DirectEchoProtocol : IDisposable
{
    private const string PingPrefix = "rv1-ping:";
    private const string EchoPrefix = "rv1-echo:";
    private const string AckPrefix = "rv1-ack:";
    private const string DonePrefix = "rv1-done:";
    private readonly EventBasedNetListener _events;
    private readonly bool _host;
    private readonly Dictionary<NetPeer, string> _hostNonces = [];
    private readonly TaskCompletionSource<bool> _completed = new(
        TaskCreationOptions.RunContinuationsAsynchronously);
    private string? _nonce;
    private bool _disposed;

    internal DirectEchoProtocol(EventBasedNetListener events, bool host)
    {
        _events = events ?? throw new ArgumentNullException(nameof(events));
        _host = host;
        _events.NetworkReceiveEvent += OnReceive;
        _events.PeerDisconnectedEvent += OnPeerDisconnected;
    }

    internal Task Completion => _completed.Task;
    internal int PendingHostExchangeCount => _hostNonces.Count;
    internal event Action<NetPeer>? ExchangeCompleted;

    internal void BeginJoin(NetPeer peer)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);
        if (_host || _nonce is not null)
        {
            throw new InvalidOperationException("The direct echo exchange is already active.");
        }

        _nonce = Convert.ToHexString(RandomNumberGenerator.GetBytes(16)).ToLowerInvariant();
        Send(peer, PingPrefix + _nonce);
    }

    public void Dispose()
    {
        if (_disposed)
        {
            return;
        }
        _events.NetworkReceiveEvent -= OnReceive;
        _events.PeerDisconnectedEvent -= OnPeerDisconnected;
        _hostNonces.Clear();
        _disposed = true;
    }

    private void OnReceive(
        NetPeer peer,
        NetPacketReader reader,
        byte channel,
        DeliveryMethod deliveryMethod)
    {
        try
        {
            ReadOnlySpan<byte> payload = reader.GetRemainingBytesSpan();
            if (payload.Length is < 9 or > 64)
            {
                return;
            }
            string message = Encoding.ASCII.GetString(payload);
            if (_host && TryNonce(message, PingPrefix, out string? pingNonce))
            {
                _hostNonces[peer] = pingNonce!;
                Send(peer, EchoPrefix + pingNonce);
            }
            else if (_host
                     && _hostNonces.TryGetValue(peer, out string? hostNonce)
                     && string.Equals(message, AckPrefix + hostNonce, StringComparison.Ordinal))
            {
                _hostNonces.Remove(peer);
                Send(peer, DonePrefix + hostNonce);
                ExchangeCompleted?.Invoke(peer);
                _completed.TrySetResult(true);
            }
            else if (!_host
                     && _nonce is not null
                     && string.Equals(message, EchoPrefix + _nonce, StringComparison.Ordinal))
            {
                Send(peer, AckPrefix + _nonce);
            }
            else if (!_host
                     && _nonce is not null
                     && string.Equals(message, DonePrefix + _nonce, StringComparison.Ordinal))
            {
                ExchangeCompleted?.Invoke(peer);
                _completed.TrySetResult(true);
            }
        }
        finally
        {
            reader.Recycle();
        }
    }

    private static bool TryNonce(string message, string prefix, out string? nonce)
    {
        nonce = null;
        if (!message.StartsWith(prefix, StringComparison.Ordinal)
            || message.Length != prefix.Length + 32)
        {
            return false;
        }
        string candidate = message[prefix.Length..];
        if (!candidate.All(static character => character is >= '0' and <= '9'
                or >= 'a' and <= 'f'))
        {
            return false;
        }
        nonce = candidate;
        return true;
    }

    private static void Send(NetPeer peer, string message) => peer.Send(
        Encoding.ASCII.GetBytes(message),
        DeliveryMethod.ReliableOrdered);

    private void OnPeerDisconnected(NetPeer peer, DisconnectInfo disconnectInfo) =>
        _hostNonces.Remove(peer);
}