ShrinkNetwork/Runtime/Transport/Kcp/ShrinkKcpClientTransport.cs

#nullable enable

using System;
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 ShrinkKcpClientTransport : IShrinkNetworkAsyncTransport
    {
        private readonly string _host;
        private readonly int _port;
        private readonly long _sessionId;
        private readonly ShrinkKcpTransportOptions _options;
        private readonly object _syncRoot = new();

        private CancellationTokenSource? _cts;
        private UdpClient? _udpClient;
        private ShrinkKcpPeer? _peer;
        private long _handshakeNonce;
        private int _started;
        private int _connected;

        public ShrinkKcpClientTransport(string host, int port, ShrinkKcpTransportOptions? options = null, long sessionId = 1)
        {
            if (string.IsNullOrWhiteSpace(host))
                throw new ArgumentException("Host cannot be empty.", nameof(host));
            if (port <= 0 || port > 65535)
                throw new ArgumentOutOfRangeException(nameof(port));

            _host = host.Trim();
            _port = port;
            _sessionId = sessionId;
            _options = (options ?? new ShrinkKcpTransportOptions()).Clone();
            _options.Validate();
        }

        public bool IsStarted => Volatile.Read(ref _started) == 1;

        public event Action<ShrinkNetworkTransportEvent>? OnEvent;

        public void Start()
        {
            if (Interlocked.Exchange(ref _started, 1) == 1)
                return;

            _cts = new CancellationTokenSource();
            _udpClient = new UdpClient(AddressFamily.InterNetwork);
            _udpClient.Client.ReceiveBufferSize = _options.ReceiveBufferSize;
            _udpClient.Connect(_host, _port);

            _handshakeNonce = CreateHandshakeNonce();
            ReceiveLoopAsync(_udpClient, _cts.Token).Forget();
            HandshakeLoopAsync(_cts.Token).Forget();
            UpdateLoopAsync(_cts.Token).Forget();
        }

        public void Stop()
        {
            var remote = $"{_host}:{_port}";
            var wasConnected = Interlocked.Exchange(ref _connected, 0) == 1;

            if (Interlocked.Exchange(ref _started, 0) == 0)
                return;

            try
            {
                var peer = _peer;
                if (peer != null)
                    SendDatagram(ShrinkKcpTransportProtocol.CreateDisconnect(peer.ConversationId));
            }
            catch
            {
            }

            _cts?.Cancel();

            lock (_syncRoot)
            {
                _peer?.Dispose();
                _peer = null;

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

                _udpClient = null;
            }

            if (wasConnected)
                OnEvent?.Invoke(ShrinkNetworkTransportEvent.Disconnected(_sessionId, remote));
        }

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

        public UniTask SendAsync(long sessionId, byte[] packetData)
        {
            if (!IsStarted)
                throw new InvalidOperationException("Transport is not started.");
            if (sessionId != _sessionId)
                throw new InvalidOperationException($"Unsupported session id {sessionId}. This transport only supports {_sessionId}.");
            if (Volatile.Read(ref _connected) != 1)
                throw new InvalidOperationException("KCP client is not connected.");

            lock (_syncRoot)
            {
                _peer?.Send(packetData ?? Array.Empty<byte>());
            }

            return UniTask.CompletedTask;
        }

        private async UniTaskVoid HandshakeLoopAsync(CancellationToken cancellationToken)
        {
            var startedAt = DateTime.UtcNow;

            try
            {
                while (!cancellationToken.IsCancellationRequested && Volatile.Read(ref _connected) == 0)
                {
                    if ((DateTime.UtcNow - startedAt).TotalMilliseconds > _options.ConnectTimeoutMs)
                        throw new TimeoutException($"KCP connect timed out: {_host}:{_port}");

                    SendDatagram(ShrinkKcpTransportProtocol.CreateConnectRequest(_handshakeNonce, _options.ConversationId));
                    await Task.Delay(_options.HandshakeRetryMs, cancellationToken);
                }
            }
            catch (OperationCanceledException)
            {
            }
            catch (Exception ex)
            {
                ShrinkNetworkLogger.Exception(ex);
                ShrinkNetworkLogger.Error($"[ShrinkNetwork] KCP connect failed: {_host}:{_port} {ex.Message}");
                Stop();
            }
        }

        private async UniTaskVoid ReceiveLoopAsync(UdpClient udpClient, CancellationToken cancellationToken)
        {
            try
            {
                while (!cancellationToken.IsCancellationRequested)
                {
                    UdpReceiveResult result;
                    try
                    {
                        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 receive loop ended: {ex.Message}");
            }
            finally
            {
                Stop();
            }
        }

        private async UniTaskVoid UpdateLoopAsync(CancellationToken cancellationToken)
        {
            try
            {
                while (!cancellationToken.IsCancellationRequested)
                {
                    var shouldStop = false;

                    lock (_syncRoot)
                    {
                        _peer?.Tick(packet => OnEvent?.Invoke(ShrinkNetworkTransportEvent.Packet(_sessionId, packet)));

                        if (_peer != null &&
                            Volatile.Read(ref _connected) == 1 &&
                            DateTime.UtcNow.Ticks - _peer.LastReceiveUtcTicks > TimeSpan.FromMilliseconds(_options.IdleTimeoutMs).Ticks)
                        {
                            ShrinkNetworkLogger.Warn($"[ShrinkNetwork] KCP idle timeout: {_host}:{_port}");
                            shouldStop = true;
                        }
                    }

                    if (shouldStop)
                    {
                        Stop();
                        return;
                    }

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

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

            if (ShrinkKcpTransportProtocol.TryReadConnectAccept(datagram, out var nonce, out var conversationId))
            {
                if (nonce != _handshakeNonce)
                    return;

                lock (_syncRoot)
                {
                    if (_peer == null)
                    {
                        _peer = new ShrinkKcpPeer(conversationId, _options,
                            payload => SendDatagram(ShrinkKcpTransportProtocol.CreateDataPacket(conversationId, payload)));
                    }
                }

                if (Interlocked.Exchange(ref _connected, 1) == 0)
                    OnEvent?.Invoke(ShrinkNetworkTransportEvent.Connected(_sessionId, remoteEndPoint.ToString()));

                return;
            }

            if (ShrinkKcpTransportProtocol.TryReadDisconnect(datagram, out var disconnectedConversationId))
            {
                if (_peer != null && _peer.ConversationId == disconnectedConversationId)
                    Stop();
                return;
            }

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

            lock (_syncRoot)
            {
                if (_peer == null || _peer.ConversationId != dataConversationId)
                    return;

                _peer.Input(datagram, payloadOffset, payloadLength,
                    packet => OnEvent?.Invoke(ShrinkNetworkTransportEvent.Packet(_sessionId, packet)));
            }
        }

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

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

        private static long CreateHandshakeNonce()
        {
            Span<byte> buffer = stackalloc byte[sizeof(long)];
            RandomNumberGenerator.Fill(buffer);
            return BitConverter.ToInt64(buffer);
        }
    }
}