ShrinkNetwork/Editor/Scaffolding/ServerProjectTemplate/KcpServerTransport.cs.txt

using System.Net;
using System.Net.Sockets;
using System.Security.Cryptography;
using Cysharp.Threading.Tasks;

namespace ShrinkNetwork.ServerHost;

using ShrinkNetwork;

public sealed class KcpServerTransport : IShrinkNetworkAsyncTransport, IShrinkNetworkSessionControlTransport
{
    private sealed class KcpSessionState : IDisposable
    {
        public KcpSessionState(long sessionId, IPEndPoint remoteEndPoint, ShrinkKcpPeer peer)
        {
            SessionId = sessionId;
            RemoteEndPoint = remoteEndPoint;
            Peer = peer;
        }

        public long SessionId { get; }
        public IPEndPoint RemoteEndPoint { get; }
        public ShrinkKcpPeer Peer { get; }
        public string RemoteAddress => RemoteEndPoint.ToString();

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

    private readonly object _syncRoot = new();
    private readonly Dictionary<long, KcpSessionState> _sessionsById = new();
    private readonly Dictionary<uint, KcpSessionState> _sessionsByConversationId = new();
    private readonly UdpClient _udpClient;
    private readonly ShrinkKcpTransportOptions _options;

    private CancellationTokenSource? _cts;
    private long _sessionIdGenerator;

    public KcpServerTransport(IPAddress ipAddress, int port, ShrinkKcpTransportOptions? options = null)
    {
        _udpClient = new UdpClient(new IPEndPoint(ipAddress, port));
        _options = (options ?? new ShrinkKcpTransportOptions()).Clone();
        _options.Validate();
        _udpClient.Client.ReceiveBufferSize = _options.ReceiveBufferSize;
    }

    public bool IsStarted { get; private set; }

    public event Action<ShrinkNetworkTransportEvent>? OnEvent;

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

        IsStarted = true;
        _cts = new CancellationTokenSource();
        _ = ReceiveLoopAsync(_cts.Token);
        _ = UpdateLoopAsync(_cts.Token);
    }

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

        IsStarted = false;
        _cts?.Cancel();

        List<KcpSessionState> sessions;
        lock (_syncRoot)
        {
            sessions = _sessionsById.Values.ToList();
            _sessionsById.Clear();
            _sessionsByConversationId.Clear();
        }

        foreach (var session in sessions)
        {
            TrySendDisconnect(session);
            session.Dispose();
            OnEvent?.Invoke(ShrinkNetworkTransportEvent.Disconnected(session.SessionId, session.RemoteAddress));
        }

        _udpClient.Close();
    }

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

    public bool DisconnectSession(long sessionId, string? reason = null)
    {
        KcpSessionState? session;
        lock (_syncRoot)
        {
            _sessionsById.TryGetValue(sessionId, out session);
        }

        if (session == null)
            return false;

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

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

        session.Peer.Send(packetData ?? Array.Empty<byte>());
        return UniTask.CompletedTask;
    }

    private async Task ReceiveLoopAsync(CancellationToken cancellationToken)
    {
        while (!cancellationToken.IsCancellationRequested)
        {
            UdpReceiveResult result;
            try
            {
                result = await _udpClient.ReceiveAsync(cancellationToken);
            }
            catch (OperationCanceledException)
            {
                break;
            }
            catch (ObjectDisposedException)
            {
                break;
            }
            catch
            {
                if (!IsStarted)
                    break;
                continue;
            }

            try
            {
                HandleDatagram(result.Buffer, result.RemoteEndPoint);
            }
            catch (Exception ex)
            {
                ShrinkNetworkLogger.Exception(ex);
                ShrinkNetworkLogger.Warn($"[ShrinkNetwork] KCP server datagram handling failed: {ex.Message}");
            }
        }
    }

    private async Task UpdateLoopAsync(CancellationToken cancellationToken)
    {
        while (!cancellationToken.IsCancellationRequested)
        {
            List<KcpSessionState> snapshot;
            lock (_syncRoot)
            {
                snapshot = _sessionsById.Values.ToList();
            }

            foreach (var session in snapshot)
            {
                try
                {
                    session.Peer.Tick(packet => OnEvent?.Invoke(ShrinkNetworkTransportEvent.Packet(session.SessionId, packet)));

                    if (DateTime.UtcNow.Ticks - session.Peer.LastReceiveUtcTicks >
                        TimeSpan.FromMilliseconds(_options.IdleTimeoutMs).Ticks)
                    {
                        DisconnectSession(session.SessionId, sendRemoteNotice: true, "idle timeout");
                    }
                }
                catch (Exception ex)
                {
                    ShrinkNetworkLogger.Exception(ex);
                    DisconnectSession(session.SessionId, sendRemoteNotice: true, $"peer tick failed: {ex.GetType().Name}: {ex.Message}");
                }
            }

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

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

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

        if (ShrinkKcpTransportProtocol.TryReadDisconnect(datagram, out var disconnectedConversationId))
        {
            HandleDisconnect(remoteEndPoint, disconnectedConversationId);
            return;
        }

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

        KcpSessionState? session;
        lock (_syncRoot)
        {
            _sessionsByConversationId.TryGetValue(conversationId, out session);
        }

        if (session == null)
            return;
        if (!Equals(session.RemoteEndPoint, remoteEndPoint))
            return;

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

    private void HandleConnectRequest(IPEndPoint remoteEndPoint, long nonce, uint requestedConversationId)
    {
        KcpSessionState? existingSession;
        lock (_syncRoot)
        {
            existingSession = _sessionsById.Values.FirstOrDefault(x => Equals(x.RemoteEndPoint, remoteEndPoint));
        }

        if (existingSession != null)
        {
            ShrinkNetworkLogger.Info(
                $"[ShrinkNetwork][KCP-Server] Reusing session {existingSession.SessionId} for endpoint {existingSession.RemoteAddress}.");
            SendEnvelope(existingSession.RemoteEndPoint,
                ShrinkKcpTransportProtocol.CreateConnectAccept(nonce, existingSession.Peer.ConversationId));
            return;
        }

        var conversationId = AllocateConversationId(requestedConversationId);
        var sessionId = Interlocked.Increment(ref _sessionIdGenerator);
        var peer = new ShrinkKcpPeer(conversationId, _options,
            payload => SendEnvelope(remoteEndPoint, ShrinkKcpTransportProtocol.CreateDataPacket(conversationId, payload)));
        var session = new KcpSessionState(sessionId, remoteEndPoint, peer);

        lock (_syncRoot)
        {
            _sessionsById[sessionId] = session;
            _sessionsByConversationId[conversationId] = session;
        }

        SendEnvelope(remoteEndPoint, ShrinkKcpTransportProtocol.CreateConnectAccept(nonce, conversationId));
        OnEvent?.Invoke(ShrinkNetworkTransportEvent.Connected(sessionId, remoteEndPoint.ToString()));
    }

    private void HandleDisconnect(IPEndPoint remoteEndPoint, uint conversationId)
    {
        KcpSessionState? session;
        lock (_syncRoot)
        {
            _sessionsByConversationId.TryGetValue(conversationId, out session);
        }

        if (session == null)
            return;
        if (!Equals(session.RemoteEndPoint, remoteEndPoint))
            return;

        DisconnectSession(session.SessionId, sendRemoteNotice: false, "remote disconnect packet");
    }

    private void DisconnectSession(long sessionId, bool sendRemoteNotice, string reason)
    {
        KcpSessionState? session;
        lock (_syncRoot)
        {
            if (!_sessionsById.Remove(sessionId, out session))
                return;

            _sessionsByConversationId.Remove(session.Peer.ConversationId);
        }

        if (sendRemoteNotice)
            TrySendDisconnect(session);

        ShrinkNetworkLogger.Warn(
            $"[ShrinkNetwork][KCP-Server] Session {sessionId} {session.RemoteAddress} disconnected: {reason}");
        session.Dispose();
        OnEvent?.Invoke(ShrinkNetworkTransportEvent.Disconnected(sessionId, session.RemoteAddress));
    }

    private uint AllocateConversationId(uint requestedConversationId)
    {
        lock (_syncRoot)
        {
            if (requestedConversationId != 0 && !_sessionsByConversationId.ContainsKey(requestedConversationId))
                return requestedConversationId;
        }

        while (true)
        {
            var conversationId = unchecked((uint)RandomNumberGenerator.GetInt32(1, int.MaxValue));
            lock (_syncRoot)
            {
                if (!_sessionsByConversationId.ContainsKey(conversationId))
                    return conversationId;
            }
        }
    }

    private void TrySendDisconnect(KcpSessionState session)
    {
        try
        {
            SendEnvelope(session.RemoteEndPoint, ShrinkKcpTransportProtocol.CreateDisconnect(session.Peer.ConversationId));
        }
        catch
        {
        }
    }

    private void SendEnvelope(IPEndPoint remoteEndPoint, byte[] datagram)
    {
        _udpClient.Send(datagram, datagram.Length, remoteEndPoint);
    }
}