Commonwealth-Online-Server/server/GnsTransport.cs

using System.Net;
using System.Runtime.InteropServices;
using System.Text;
using System.Text.Json.Nodes;

namespace CommonwealthOnline.Server;

internal enum GnsEventType : uint { None = 0, Connected = 1, Disconnected = 2, Message = 3, OversizeMessage = 4 }

internal sealed unsafe class GnsNativeServer : IDisposable
{
    [StructLayout(LayoutKind.Sequential)]
    private struct NativeEvent
    {
        public uint Type;
        public uint ConnectionId;
        public int Reason;
        public uint PayloadSize;
        public fixed byte Debug[128];
    }

    [UnmanagedFunctionPointer(CallingConvention.Cdecl)] private delegate int CreateDelegate([MarshalAs(UnmanagedType.LPUTF8Str)] string bindHost, ushort port, out IntPtr handle, IntPtr errorBuffer, nuint errorBufferSize);
    [UnmanagedFunctionPointer(CallingConvention.Cdecl)] private delegate void DestroyDelegate(IntPtr handle);
    [UnmanagedFunctionPointer(CallingConvention.Cdecl)] private delegate ushort LocalPortDelegate(IntPtr handle);
    [UnmanagedFunctionPointer(CallingConvention.Cdecl)] private delegate uint ConnectionCountDelegate(IntPtr handle);
    [UnmanagedFunctionPointer(CallingConvention.Cdecl)] private delegate int PollDelegate(IntPtr handle, NativeEvent* outEvent, IntPtr payloadBuffer, uint payloadCapacity);
    [UnmanagedFunctionPointer(CallingConvention.Cdecl)] private delegate int SendDelegate(IntPtr handle, uint connectionId, IntPtr payload, uint payloadSize, uint delivery);
    [UnmanagedFunctionPointer(CallingConvention.Cdecl)] private delegate int DisconnectDelegate(IntPtr handle, uint connectionId, int reason, [MarshalAs(UnmanagedType.LPUTF8Str)] string debug);
    [UnmanagedFunctionPointer(CallingConvention.Cdecl)] private delegate int RemoteIpv4Delegate(IntPtr handle, uint connectionId, out uint ipv4HostOrder, out ushort port);

    private readonly IntPtr _library;
    private IntPtr _handle;
    private readonly DestroyDelegate _destroy;
    private readonly LocalPortDelegate _localPort;
    private readonly ConnectionCountDelegate _connectionCount;
    private readonly PollDelegate _poll;
    private readonly SendDelegate _send;
    private readonly DisconnectDelegate _disconnect;
    private readonly RemoteIpv4Delegate _remoteIpv4;
    private readonly IntPtr _payloadBuffer = Marshal.AllocHGlobal(ProtocolConstants.MaxMessageBytes);
    private int _disposed;

    public GnsNativeServer(string bindHost, int port, string? configuredPath, string serverBaseDirectory)
    {
        var libraryPath = ResolveLibrary(configuredPath, serverBaseDirectory);
        _library = NativeLibrary.Load(libraryPath);
        var create = Get<CreateDelegate>("co_gns_server_create");
        _destroy = Get<DestroyDelegate>("co_gns_server_destroy");
        _localPort = Get<LocalPortDelegate>("co_gns_server_local_port");
        _connectionCount = Get<ConnectionCountDelegate>("co_gns_server_connection_count");
        _poll = Get<PollDelegate>("co_gns_server_poll");
        _send = Get<SendDelegate>("co_gns_server_send");
        _disconnect = Get<DisconnectDelegate>("co_gns_server_disconnect");
        _remoteIpv4 = Get<RemoteIpv4Delegate>("co_gns_server_remote_ipv4");

        var errorBuffer = Marshal.AllocHGlobal(512);
        try
        {
            new Span<byte>((void*)errorBuffer, 512).Clear();
            var result = create(bindHost, checked((ushort)port), out _handle, errorBuffer, 512);
            if (result != 1 || _handle == IntPtr.Zero)
                throw new InvalidOperationException(Marshal.PtrToStringUTF8(errorBuffer) ?? "GNS native bridge failed to start");
        }
        finally { Marshal.FreeHGlobal(errorBuffer); }
    }

    public ushort LocalPort => _localPort(_handle);
    public uint ConnectionCount => _connectionCount(_handle);

    public (GnsEventType Type, uint ConnectionId, int Reason, byte[] Payload, string Debug)? Poll()
    {
        NativeEvent native = default;
        var result = _poll(_handle, &native, _payloadBuffer, ProtocolConstants.MaxMessageBytes);
        if (result == 0) return null;
        if (result < 0) throw new IOException($"GNS native poll failed with result {result}");
        if (!Enum.IsDefined(typeof(GnsEventType), native.Type)) throw new IOException($"GNS native bridge returned unknown event type {native.Type}");
        var type = (GnsEventType)native.Type;
        if (native.PayloadSize > ProtocolConstants.MaxMessageBytes && type != GnsEventType.OversizeMessage) throw new IOException("GNS native bridge returned an oversized message payload");
        var payload = Array.Empty<byte>();
        if (type == GnsEventType.Message && native.PayloadSize > 0)
        {
            payload = new byte[native.PayloadSize];
            Marshal.Copy(_payloadBuffer, payload, 0, payload.Length);
        }
        string debug;
        {
            byte* pointer = native.Debug;
            var length = 0;
            while (length < 128 && pointer[length] != 0) length++;
            debug = Encoding.UTF8.GetString(pointer, length);
        }
        return (type, native.ConnectionId, native.Reason, payload, debug);
    }

    public SendOutcome Send(uint connectionId, ReadOnlySpan<byte> payload, Delivery delivery)
    {
        if (payload.Length > ProtocolConstants.MaxMessageBytes) return SendOutcome.TooLarge;
        fixed (byte* pointer = payload)
        {
            return _send(_handle, connectionId, (IntPtr)pointer, (uint)payload.Length, (uint)delivery) switch
            {
                0 => SendOutcome.Sent,
                1 => SendOutcome.Dropped,
                2 => SendOutcome.Backpressure,
                3 => SendOutcome.NotConnected,
                4 => SendOutcome.TooLarge,
                _ => SendOutcome.Error
            };
        }
    }

    public void Disconnect(uint connectionId, int reason, string debug) => _disconnect(_handle, connectionId, reason, debug);

    public IPEndPoint? RemoteEndpoint(uint connectionId)
    {
        if (_remoteIpv4(_handle, connectionId, out var ipv4, out var port) != 1) return null;
        var bytes = new[] { (byte)(ipv4 >> 24), (byte)(ipv4 >> 16), (byte)(ipv4 >> 8), (byte)ipv4 };
        return new IPEndPoint(new IPAddress(bytes), port);
    }

    private T Get<T>(string name) where T : Delegate => Marshal.GetDelegateForFunctionPointer<T>(NativeLibrary.GetExport(_library, name));

    private static string ResolveLibrary(string? configuredPath, string serverBaseDirectory)
    {
        if (!string.IsNullOrWhiteSpace(configuredPath))
        {
            var full = Path.GetFullPath(configuredPath, serverBaseDirectory);
            if (File.Exists(full)) return full;
            throw new FileNotFoundException("Configured GNS bridge was not found", full);
        }
        var name = OperatingSystem.IsWindows() ? "commonwealth_online_gns_bridge.dll" : OperatingSystem.IsMacOS() ? "libcommonwealth_online_gns_bridge.dylib" : "libcommonwealth_online_gns_bridge.so";
        var candidates = new[]
        {
            Path.Combine(AppContext.BaseDirectory, name),
            Path.Combine(AppContext.BaseDirectory, "native_transport", name),
            Path.Combine(serverBaseDirectory, name),
            Path.Combine(serverBaseDirectory, "native_transport", name),
            Path.Combine(Environment.CurrentDirectory, "native_transport", name)
        };
        return candidates.FirstOrDefault(File.Exists) ?? throw new FileNotFoundException($"Commonwealth Online GNS native bridge was not found. Searched: {string.Join(", ", candidates)}");
    }

    public void Dispose()
    {
        if (Interlocked.Exchange(ref _disposed, 1) != 0) return;
        if (_handle != IntPtr.Zero) { _destroy(_handle); _handle = IntPtr.Zero; }
        Marshal.FreeHGlobal(_payloadBuffer);
        if (_library != IntPtr.Zero) NativeLibrary.Free(_library);
    }
}

internal sealed class GnsGameConnection : IGameConnection
{
    private readonly GnsNativeServer _native;
    private readonly uint _id;
    private readonly object _gate = new();
    private readonly Dictionary<string, SequenceCounter> _outgoingSequences = new(StringComparer.Ordinal)
    {
        ["transform"] = new SequenceCounter(), ["npcState"] = new SequenceCounter()
    };
    private int _closed;

    public GnsGameConnection(GnsNativeServer native, uint id, IPEndPoint remoteEndpoint) { _native = native; _id = id; RemoteEndpoint = remoteEndpoint; ConnectionKey = $"gns:{id}"; }
    public string ConnectionKey { get; }
    public string TransportName => "gns";
    public IPEndPoint RemoteEndpoint { get; }
    public bool IsClosed => Volatile.Read(ref _closed) != 0;
    public uint NativeId => _id;

    public ValueTask<SendOutcome> SendAsync(EncodedPacket packet, CancellationToken cancellationToken = default)
    {
        if (IsClosed) return ValueTask.FromResult(SendOutcome.NotConnected);
        lock (_gate)
        {
            if (IsClosed) return ValueTask.FromResult(SendOutcome.NotConnected);
            ReadOnlySpan<byte> wire = packet.Payload;
            byte[]? envelope = null;
            if (TransportPolicy.IsSnapshot(packet.PacketType))
            {
                envelope = GnsSnapshotEnvelope.Encode(packet.PacketType, packet.Payload, _outgoingSequences[packet.PacketType].Advance());
                wire = envelope;
            }
            return ValueTask.FromResult(_native.Send(_id, wire, packet.Delivery));
        }
    }

    public ValueTask DisconnectAsync(int reason, string debug)
    {
        if (Interlocked.Exchange(ref _closed, 1) == 0) { try { _native.Disconnect(_id, reason, debug); } catch { } }
        return ValueTask.CompletedTask;
    }
    internal void MarkRemoteClosed() => Interlocked.Exchange(ref _closed, 1);
    public ValueTask DisposeAsync() => DisconnectAsync(0, "dispose");
}

internal sealed class GnsServerTransport : IAsyncDisposable
{
    private readonly ServerOptions _options;
    private readonly IServerIngress _server;
    private readonly CancellationTokenSource _shutdown = new();
    private readonly Dictionary<uint, GnsGameConnection> _connections = new();
    private readonly Dictionary<(uint ConnectionId, string PacketType), SequenceWindow> _incomingSequences = new();
    private readonly object _gate = new();
    private GnsNativeServer? _native;
    private Task? _pumpTask;

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

    public void Start()
    {
        if (_pumpTask is not null) return;
        _native = new GnsNativeServer(_options.Host, _options.Port, _options.GnsBridgePath, _options.BaseDirectory);
        if (_native.LocalPort != _options.Port) throw new InvalidOperationException($"GNS transport bound UDP {_native.LocalPort}, expected UDP {_options.Port}.");
        _pumpTask = Task.Run(() => PumpAsync(_shutdown.Token));
        _server.Log($"GameNetworkingSockets gameplay transport listening on UDP {_options.Port}.");
    }

    private async Task PumpAsync(CancellationToken cancellationToken)
    {
        while (!cancellationToken.IsCancellationRequested)
        {
            (GnsEventType Type, uint ConnectionId, int Reason, byte[] Payload, string Debug)? evt;
            try { evt = _native!.Poll(); }
            catch (Exception ex)
            {
                if (!cancellationToken.IsCancellationRequested) _server.Log($"GNS poll error: {ex.Message}", "error");
                try { await Task.Delay(10, cancellationToken).ConfigureAwait(false); } catch (OperationCanceledException) { break; }
                continue;
            }
            if (evt is null)
            {
                try { await Task.Delay(2, cancellationToken).ConfigureAwait(false); } catch (OperationCanceledException) { break; }
                continue;
            }
            var value = evt.Value;
            switch (value.Type)
            {
                case GnsEventType.Connected: await HandleConnectedAsync(value.ConnectionId, cancellationToken).ConfigureAwait(false); break;
                case GnsEventType.Message: await HandleMessageAsync(value.ConnectionId, value.Payload, cancellationToken).ConfigureAwait(false); break;
                case GnsEventType.OversizeMessage:
                    if (TryGetConnection(value.ConnectionId, out var oversized)) await _server.EndSessionForTransportAsync(oversized, "packet_too_large", "Packet exceeded maximum message size.").ConfigureAwait(false);
                    else _native!.Disconnect(value.ConnectionId, 0, "Oversized pre-session packet");
                    break;
                case GnsEventType.Disconnected: await HandleDisconnectedAsync(value.ConnectionId).ConfigureAwait(false); break;
            }
        }
    }

    private async Task HandleConnectedAsync(uint id, CancellationToken cancellationToken)
    {
        var endpoint = _native!.RemoteEndpoint(id);
        if (endpoint is null) { _native.Disconnect(id, 0, "Remote endpoint unavailable"); return; }
        var connection = new GnsGameConnection(_native, id, endpoint);
        lock (_gate) _connections[id] = connection;
        bool accepted;
        try { accepted = await _server.AcceptConnectionAsync(connection, cancellationToken).ConfigureAwait(false); }
        catch (Exception ex) { _server.Log($"GNS admission failed for {endpoint}: {ex.Message}", "warning"); accepted = false; }
        if (!accepted) { lock (_gate) _connections.Remove(id); await connection.DisposeAsync(); }
    }

    private async Task HandleMessageAsync(uint id, byte[] payload, CancellationToken cancellationToken)
    {
        if (!TryGetConnection(id, out var connection)) { _native!.Disconnect(id, 0, "Message before GNS admission"); return; }
        ReadOnlyMemory<byte> gameplayPayload = payload;
        if (GnsSnapshotEnvelope.TryDecode(payload, out var envelope, out var envelopeError))
        {
            if (envelopeError is not null) { await _server.HandleTransportRejectAsync(connection, $"Malformed GNS snapshot envelope: {envelopeError}").ConfigureAwait(false); return; }
            bool accepted;
            lock (_gate)
            {
                if (!_incomingSequences.TryGetValue((id, envelope.PacketType), out var window)) _incomingSequences[(id, envelope.PacketType)] = window = new SequenceWindow();
                accepted = window.Accept(envelope.Sequence);
            }
            if (!accepted) { await _server.HandleTransportRejectAsync(connection, "Stale or duplicate GNS snapshot sequence", false).ConfigureAwait(false); return; }
            JsonObject packet;
            try { packet = PacketCodec.Decode(envelope.Payload); }
            catch (PacketCodecException ex) { await _server.HandleTransportRejectAsync(connection, $"Invalid GNS snapshot payload: {ex.Message}").ConfigureAwait(false); return; }
            if (JsonHelpers.String(packet["type"]) != envelope.PacketType) { await _server.HandleTransportRejectAsync(connection, "GNS snapshot envelope family does not match packet type").ConfigureAwait(false); return; }
            gameplayPayload = envelope.Payload;
        }
        else
        {
            try
            {
                var packet = PacketCodec.Decode(payload);
                var type = JsonHelpers.String(packet["type"])!;
                if (TransportPolicy.IsSnapshot(type)) { await _server.HandleTransportRejectAsync(connection, "GNS snapshot missing required sequence envelope").ConfigureAwait(false); return; }
            }
            catch (PacketCodecException) { }
        }
        await _server.HandleMessageAsync(connection, gameplayPayload, cancellationToken).ConfigureAwait(false);
    }

    private async Task HandleDisconnectedAsync(uint id)
    {
        GnsGameConnection? connection;
        lock (_gate)
        {
            _connections.Remove(id, out connection);
            foreach (var key in _incomingSequences.Keys.Where(x => x.ConnectionId == id).ToArray()) _incomingSequences.Remove(key);
        }
        if (connection is null) return;
        connection.MarkRemoteClosed();
        await _server.HandleConnectionClosedAsync(connection).ConfigureAwait(false);
    }

    private bool TryGetConnection(uint id, out GnsGameConnection connection) { lock (_gate) return _connections.TryGetValue(id, out connection!); }

    public async ValueTask DisposeAsync()
    {
        _shutdown.Cancel();
        if (_pumpTask is not null) { try { await _pumpTask.ConfigureAwait(false); } catch { } }
        GnsGameConnection[] connections;
        lock (_gate) { connections = _connections.Values.ToArray(); _connections.Clear(); _incomingSequences.Clear(); }
        foreach (var connection in connections) await connection.DisposeAsync();
        _native?.Dispose();
        _native = null;
        _shutdown.Dispose();
    }
}