ShrinkNetwork/Runtime/Transport/Kcp/ShrinkKcpServerTransport.cs

#nullable enable
using System;
using System.Buffers.Binary;
using System.Collections.Concurrent;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Threading;
using System.Threading.Tasks;
using Cysharp.Threading.Tasks;

namespace ShrinkNetwork
{
    public sealed class ShrinkKcpServerTransport : IShrinkNetworkAsyncTransport, IShrinkNetworkSessionControlTransport
    {
        private sealed class SessionState : IDisposable
        {
            public long SessionId;
            public uint ConversationId;
            public IPEndPoint RemoteEndPoint = null!;
            public ShrinkKcpPeer Peer = null!;
            public readonly object SyncRoot = new();

            public void Dispose()
            {
                Peer.Dispose();
            }
        }

        private readonly ConcurrentDictionary<long, SessionState> _sessions = new();
        private readonly ConcurrentDictionary<string, long> _endpointToSessionId = new();
        private readonly ConcurrentDictionary<uint, long> _conversationToSessionId = new();
        private readonly ShrinkKcpTransportOptions _options;
        private readonly IPAddress _listeningAddress;
        private readonly int _port;
        private CancellationTokenSource? _cts;
        private UdpClient? _udpClient;
        private long _sessionIdGenerator;

        public ShrinkKcpServerTransport(IPAddress listeningAddress, int port, ShrinkKcpTransportOptions? options = null)
        {
            if (port <= 0 || port > 65535)
                throw new ArgumentOutOfRangeException(nameof(port));

            _listeningAddress = listeningAddress ?? throw new ArgumentNullException(nameof(listeningAddress));
            _port = port;
            _options = (options ?? new ShrinkKcpTransportOptions()).Clone();
            _options.Validate();
        }

        public bool IsStarted { get; private set; }

        public event Action<ShrinkNetworkTransportEvent>? OnEvent;

        public void Start()
        {
            if (IsStarted)
                return;

            _udpClient = new UdpClient(new IPEndPoint(_listeningAddress, _port));
            _udpClient.Client.ReceiveBufferSize = _options.ReceiveBufferSize;
            _cts = new CancellationTokenSource();
            IsStarted = true;
            ReceiveLoopAsync(_cts.Token).Forget();
            UpdateLoopAsync(_cts.Token).Forget();
        }

        public void Stop()
        {
            if (!IsStarted)
                return;

            IsStarted = false;
            _cts?.Cancel();

            try
            {
                _udpClient?.Close();
            }
            catch
            {
            }

            _udpClient = null;

            foreach (var state in _sessions.Values.ToArray())
                RemoveSession(state.SessionId, true);

            _endpointToSessionId.Clear();
            _conversationToSessionId.Clear();
        }

        public void Send(long sessionId, byte[] packetData)
        {
            SendAsync(sessionId, packetData).Forget();
        }

        public bool DisconnectSession(long sessionId, string? reason = null)
        {
            if (!_sessions.ContainsKey(sessionId))
                return false;

            RemoveSession(sessionId, true);
            if (!string.IsNullOrWhiteSpace(reason))
                ShrinkNetworkLogger.Warn($"[ShrinkNetwork][KCP-Server] Disconnect session {sessionId}: {reason}");
            return true;
        }

        public UniTask SendAsync(long sessionId, byte[] packetData)
        {
            if (!_sessions.TryGetValue(sessionId, out var session))
                throw new InvalidOperationException($"Session {sessionId} is not connected.");

            lock (session.SyncRoot)
            {
                session.Peer.Send(packetData ?? Array.Empty<byte>());
            }

            return UniTask.CompletedTask;
        }

        private async UniTaskVoid ReceiveLoopAsync(CancellationToken cancellationToken)
        {
            try
            {
                while (!cancellationToken.IsCancellationRequested)
                {
                    UdpReceiveResult result;
                    try
                    {
                        if (_udpClient == null)
                            break;

                        result = await _udpClient.ReceiveAsync();
                    }
                    catch (ObjectDisposedException)
                    {
                        break;
                    }
                    catch (SocketException) when (!IsStarted)
                    {
                        break;
                    }

                    HandleDatagram(result.Buffer, result.RemoteEndPoint);
                }
            }
            catch (Exception ex)
            {
                ShrinkNetworkLogger.Exception(ex);
                ShrinkNetworkLogger.Warn($"[ShrinkNetwork] KCP server receive loop ended: {ex.Message}");
            }
            finally
            {
                Stop();
            }
        }

        private async UniTaskVoid UpdateLoopAsync(CancellationToken cancellationToken)
        {
            try
            {
                while (!cancellationToken.IsCancellationRequested)
                {
                    var sessions = _sessions.Values.ToArray();
                    foreach (var session in sessions)
                    {
                        var shouldDisconnect = false;
                        lock (session.SyncRoot)
                        {
                            session.Peer.Tick(packet => OnEvent?.Invoke(
                                ShrinkNetworkTransportEvent.Packet(session.SessionId, packet)));

                            if (DateTime.UtcNow.Ticks - session.Peer.LastReceiveUtcTicks >
                                TimeSpan.FromMilliseconds(_options.IdleTimeoutMs).Ticks)
                            {
                                shouldDisconnect = true;
                            }
                        }

                        if (shouldDisconnect)
                            RemoveSession(session.SessionId, true);
                    }

                    await Task.Delay(_options.UpdateIntervalMs, cancellationToken);
                }
            }
            catch (OperationCanceledException)
            {
            }
        }

        private void HandleDatagram(byte[] datagram, IPEndPoint remoteEndPoint)
        {
            if (datagram == null || datagram.Length == 0)
                return;

            if (ShrinkKcpTransportProtocol.TryReadConnectRequest(datagram, out var nonce, out _))
            {
                HandleConnectRequest(remoteEndPoint, nonce);
                return;
            }

            if (ShrinkKcpTransportProtocol.TryReadDisconnect(datagram, out var disconnectedConversationId))
            {
                if (_conversationToSessionId.TryGetValue(disconnectedConversationId, out var disconnectedSessionId))
                {
                    if (_sessions.TryGetValue(disconnectedSessionId, out var disconnectedSession) &&
                        string.Equals(disconnectedSession.RemoteEndPoint.ToString(), remoteEndPoint.ToString(), StringComparison.Ordinal))
                    {
                        RemoveSession(disconnectedSessionId, false);
                    }
                }
                return;
            }

            if (!ShrinkKcpTransportProtocol.TryReadDataPacket(datagram, out var conversationId, out var payloadOffset,
                    out var payloadLength))
                return;

            if (!_conversationToSessionId.TryGetValue(conversationId, out var sessionId))
                return;

            if (!_sessions.TryGetValue(sessionId, out var session))
                return;

            if (!string.Equals(session.RemoteEndPoint.ToString(), remoteEndPoint.ToString(), StringComparison.Ordinal))
                return;

            lock (session.SyncRoot)
            {
                session.Peer.Input(datagram, payloadOffset, payloadLength,
                    packet => OnEvent?.Invoke(ShrinkNetworkTransportEvent.Packet(session.SessionId, packet)));
            }
        }

        private void HandleConnectRequest(IPEndPoint remoteEndPoint, long nonce)
        {
            var endpointKey = remoteEndPoint.ToString();
            if (_endpointToSessionId.TryGetValue(endpointKey, out var existingSessionId) &&
                _sessions.TryGetValue(existingSessionId, out var existingSession))
            {
                SendDatagram(remoteEndPoint,
                    ShrinkKcpTransportProtocol.CreateConnectAccept(nonce, existingSession.ConversationId));
                return;
            }

            var sessionId = Interlocked.Increment(ref _sessionIdGenerator);
            var conversationId = NextConversationId();
            var session = new SessionState
            {
                SessionId = sessionId,
                ConversationId = conversationId,
                RemoteEndPoint = remoteEndPoint,
                Peer = new ShrinkKcpPeer(conversationId, _options,
                    payload => SendDatagram(remoteEndPoint,
                        ShrinkKcpTransportProtocol.CreateDataPacket(conversationId, payload)))
            };

            _sessions[sessionId] = session;
            _endpointToSessionId[endpointKey] = sessionId;
            _conversationToSessionId[conversationId] = sessionId;
            SendDatagram(remoteEndPoint, ShrinkKcpTransportProtocol.CreateConnectAccept(nonce, conversationId));
            OnEvent?.Invoke(ShrinkNetworkTransportEvent.Connected(sessionId, endpointKey));
        }

        private void RemoveSession(long sessionId, bool notifyDisconnect)
        {
            if (!_sessions.TryRemove(sessionId, out var session))
                return;

            _endpointToSessionId.TryRemove(session.RemoteEndPoint.ToString(), out _);
            _conversationToSessionId.TryRemove(session.ConversationId, out _);

            if (notifyDisconnect)
            {
                try
                {
                    SendDatagram(session.RemoteEndPoint,
                        ShrinkKcpTransportProtocol.CreateDisconnect(session.ConversationId));
                }
                catch
                {
                }
            }

            session.Dispose();
            OnEvent?.Invoke(ShrinkNetworkTransportEvent.Disconnected(sessionId, session.RemoteEndPoint.ToString()));
        }

        private void SendDatagram(IPEndPoint remoteEndPoint, byte[] datagram)
        {
            if (_udpClient == null)
                return;

            _udpClient.Send(datagram, datagram.Length, remoteEndPoint);
        }

        private uint NextConversationId()
        {
            Span<byte> buffer = stackalloc byte[sizeof(uint)];
            while (true)
            {
                RandomNumberGenerator.Fill(buffer);
                var next = BinaryPrimitives.ReadUInt32LittleEndian(buffer);
                if (next == 0)
                    continue;
                if (!_conversationToSessionId.ContainsKey(next))
                    return next;
            }
        }
    }
}