Commonwealth-Online-Server/server/TcpTransport.cs

using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;

namespace CommonwealthOnline.Server;

internal sealed class TcpGameConnection : IGameConnection
{
    private readonly TcpClient _client;
    private readonly NetworkStream _stream;
    private readonly SemaphoreSlim _sendGate = new(1, 1);
    private int _closed;

    public TcpGameConnection(TcpClient client)
    {
        _client = client;
        _client.NoDelay = true;
        _stream = client.GetStream();
        RemoteEndpoint = (IPEndPoint)(_client.Client.RemoteEndPoint ?? throw new InvalidOperationException("TCP remote endpoint missing"));
        ConnectionKey = $"tcp:{RemoteEndpoint.Address}:{RemoteEndpoint.Port}:{Guid.NewGuid():N}";
    }

    public string ConnectionKey { get; }
    public string TransportName => "tcp";
    public IPEndPoint RemoteEndpoint { get; }
    public bool IsClosed => Volatile.Read(ref _closed) != 0;
    internal NetworkStream Stream => _stream;

    public async ValueTask<SendOutcome> SendAsync(EncodedPacket packet, CancellationToken cancellationToken = default)
    {
        if (IsClosed) return SendOutcome.NotConnected;
        if (packet.Payload.Length > ProtocolConstants.MaxMessageBytes) return SendOutcome.TooLarge;
        if (packet.Payload.AsSpan().IndexOf((byte)'\n') >= 0 || packet.Payload.AsSpan().IndexOf((byte)'\r') >= 0) return SendOutcome.Error;
        await _sendGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            if (IsClosed) return SendOutcome.NotConnected;
            await _stream.WriteAsync(packet.Payload, cancellationToken).ConfigureAwait(false);
            await _stream.WriteAsync(new byte[] { (byte)'\n' }, cancellationToken).ConfigureAwait(false);
            return SendOutcome.Sent;
        }
        catch (Exception ex) when (ex is IOException or SocketException or ObjectDisposedException)
        {
            return SendOutcome.NotConnected;
        }
        finally { _sendGate.Release(); }
    }

    public ValueTask DisconnectAsync(int reason, string debug)
    {
        if (Interlocked.Exchange(ref _closed, 1) != 0) return ValueTask.CompletedTask;
        try { _client.Client.Shutdown(SocketShutdown.Both); } catch { }
        try { _client.Close(); } catch { }
        return ValueTask.CompletedTask;
    }

    public async ValueTask DisposeAsync()
    {
        await DisconnectAsync(0, "dispose");
        _sendGate.Dispose();
        _stream.Dispose();
        _client.Dispose();
    }
}

internal sealed class TcpServerTransport : IAsyncDisposable
{
    private readonly ServerOptions _options;
    private readonly IServerIngress _server;
    private readonly CancellationTokenSource _shutdown = new();
    private readonly ConcurrentDictionary<string, Task> _clientTasks = new();
    private TcpListener? _listener;
    private Task? _acceptTask;

    public TcpServerTransport(ServerOptions options, IServerIngress server) { _options = options; _server = server; }

    public void Start()
    {
        if (_acceptTask is not null) return;
        if (!ServerOptions.TryResolveIpv4(_options.Host, out var address)) throw new InvalidOperationException($"Could not resolve IPv4 bind host {_options.Host}");
        _listener = new TcpListener(address, _options.Port);
        _listener.Start();
        _acceptTask = Task.Run(() => AcceptLoopAsync(_shutdown.Token));
        _server.Log($"TCP compatibility transport listening on {_options.Host}:{_options.Port}");
    }

    private async Task AcceptLoopAsync(CancellationToken cancellationToken)
    {
        while (!cancellationToken.IsCancellationRequested)
        {
            TcpClient client;
            try { client = await _listener!.AcceptTcpClientAsync(cancellationToken).ConfigureAwait(false); }
            catch (OperationCanceledException) { break; }
            catch (ObjectDisposedException) { break; }
            catch (SocketException ex)
            {
                if (!cancellationToken.IsCancellationRequested) _server.Log($"TCP accept error: {ex.Message}", "warning");
                continue;
            }
            var connection = new TcpGameConnection(client);
            bool accepted;
            try { accepted = await _server.AcceptConnectionAsync(connection, cancellationToken).ConfigureAwait(false); }
            catch (Exception ex)
            {
                _server.Log($"TCP admission failed for {connection.RemoteEndpoint}: {ex.Message}", "warning");
                await connection.DisposeAsync();
                continue;
            }
            if (!accepted) { await connection.DisposeAsync(); continue; }
            var task = RunClientAsync(connection, cancellationToken);
            _clientTasks[connection.ConnectionKey] = task;
            _ = task.ContinueWith(_ => _clientTasks.TryRemove(connection.ConnectionKey, out _), TaskScheduler.Default);
        }
    }

    private async Task RunClientAsync(TcpGameConnection connection, CancellationToken cancellationToken)
    {
        var readBuffer = new byte[4096];
        using var message = new MemoryStream(4096);
        try
        {
            while (!cancellationToken.IsCancellationRequested && !connection.IsClosed)
            {
                int count;
                try { count = await connection.Stream.ReadAsync(readBuffer, cancellationToken).ConfigureAwait(false); }
                catch (OperationCanceledException) { break; }
                catch (Exception ex) when (ex is IOException or SocketException or ObjectDisposedException) { break; }
                if (count == 0) break;
                for (var i = 0; i < count; i++)
                {
                    var value = readBuffer[i];
                    if (value == (byte)'\n')
                    {
                        var data = message.ToArray();
                        message.SetLength(0);
                        if (data.Length > 0 && data[^1] == (byte)'\r') Array.Resize(ref data, data.Length - 1);
                        if (data.Length == 0) continue;
                        await _server.HandleMessageAsync(connection, data, cancellationToken).ConfigureAwait(false);
                        if (connection.IsClosed) return;
                    }
                    else
                    {
                        message.WriteByte(value);
                        if (message.Length > ProtocolConstants.MaxMessageBytes)
                        {
                            await _server.EndSessionForTransportAsync(connection, "packet_too_large", "Packet exceeded maximum line size.").ConfigureAwait(false);
                            return;
                        }
                    }
                }
            }
        }
        finally
        {
            await _server.HandleConnectionClosedAsync(connection).ConfigureAwait(false);
            await connection.DisposeAsync();
        }
    }

    public async ValueTask DisposeAsync()
    {
        _shutdown.Cancel();
        try { _listener?.Stop(); } catch { }
        if (_acceptTask is not null) { try { await _acceptTask.ConfigureAwait(false); } catch { } }
        var tasks = _clientTasks.Values.ToArray();
        if (tasks.Length > 0) { try { await Task.WhenAll(tasks).WaitAsync(TimeSpan.FromSeconds(2)); } catch { } }
        _shutdown.Dispose();
    }
}