yuanshen/Yuanshen.Gameserver/Network/Kcp/KcpGateway.cs

using System.Buffers;
using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using Yuanshen.Gameserver.Options;
using Yuanshen.Kcp;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;

namespace Yuanshen.Gameserver.Network.Kcp;
internal sealed class KcpGateway : IGateway
{
    private readonly Random _random;
    private readonly ILogger _logger;
    private readonly IOptions<GatewayOptions> _options;
    private readonly NetSessionManager _sessionManager;

    private readonly ConcurrentDictionary<string, (uint ConvId, uint Token)> _handshakes = new();

    private uint _sessionCounter;
    private IKcpTransport<IKcpMultiplexConnection>? _transport;

    public KcpGateway(ILogger<KcpGateway> logger, IOptions<GatewayOptions> options, NetSessionManager sessionManager)
    {
        _logger = logger;
        _options = options;
        _sessionManager = sessionManager;

        _random = new();
    }

    public Task Start()
    {
        IPEndPoint bindEndPoint = _options.Value.EndPoint;

        _transport = KcpSocketTransport.CreateMultiplexConnection(new(bindEndPoint), 1400);
        _transport.SetCallbacks(20, HandleKcpHandshake);
        _transport.Start();

        _logger.LogInformation("KCP Gateway is up at {endPoint}", bindEndPoint);
        return Task.CompletedTask;
    }

    private async ValueTask HandleKcpHandshake(UdpReceiveResult receiveResult)
    {
        KcpHandshake handshake = KcpHandshake.ReadFrom(receiveResult.Buffer);

        _logger.LogDebug("KCP control datagram from {endPoint}: len={len} head=0x{head:X} tail=0x{tail:X}",
            receiveResult.RemoteEndPoint, receiveResult.Buffer.Length, handshake.Head, handshake.Tail);

        switch ((handshake.Head, handshake.Tail))
        {
            case (KcpHandshake.StartConversationHead, KcpHandshake.StartConversationTail):
                await OnStartConversationRequest(receiveResult.RemoteEndPoint);
                break;
        }
    }

    private async ValueTask OnStartConversationRequest(IPEndPoint clientEndPoint)
    {
        try
        {
            string key = clientEndPoint.ToString();

            if (_handshakes.TryGetValue(key, out (uint ConvId, uint Token) existing))
            {
                await SendConversationCreatedPacket(clientEndPoint, existing.ConvId, existing.Token);
                _logger.LogDebug("StartConversation retransmit from {endPoint}, reusing conv={conv}",
                    clientEndPoint, (long)existing.ConvId << 32 | existing.Token);
                return;
            }

            uint convId = Interlocked.Increment(ref _sessionCounter);
            uint token  = (uint)_random.Next();

            if (!_handshakes.TryAdd(key, (convId, token)))
            {
                (uint ConvId, uint Token) winner = _handshakes[key];
                await SendConversationCreatedPacket(clientEndPoint, winner.ConvId, winner.Token);
                return;
            }

            long convId64 = (long)convId << 32 | token;
            KcpConversation conversation = _transport!.Connection.CreateConversation(convId64, clientEndPoint);
            _ = RunAndReleaseAsync(convId64, conversation, clientEndPoint, key);

            await SendConversationCreatedPacket(clientEndPoint, convId, token);
            _logger.LogInformation("StartConversation accepted from {endPoint}, conv={conv}", clientEndPoint, convId64);
        }
        catch (Exception exception)
        {
            _logger.LogError("StartConversation from {endPoint} failed: {exception}", clientEndPoint, exception);
        }
    }

    private async Task RunAndReleaseAsync(long convId64, KcpConversation conversation, IPEndPoint clientEndPoint, string key)
    {
        try
        {
            await _sessionManager.RunSessionAsync(convId64, new KcpNetworkUnit(conversation, clientEndPoint));
        }
        finally
        {
            _handshakes.TryRemove(key, out _);
        }
    }

    private async ValueTask SendConversationCreatedPacket(IPEndPoint clientEndPoint, uint convId, uint token)
    {
        KcpHandshake handshakeResponse = new()
        {
            Head   = KcpHandshake.ConversationCreatedHead,
            Param1 = convId,
            Param2 = token,
            Data   = 1234567890,
            Tail   = KcpHandshake.ConversationCreatedTail
        };

        byte[] buffer = ArrayPool<byte>.Shared.Rent(20);
        try
        {
            Memory<byte> bufferMemory = buffer.AsMemory();

            handshakeResponse.WriteTo(buffer);
            await _transport!.SendPacketAsync(bufferMemory[..20], clientEndPoint, CancellationToken.None);
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(buffer);
        }
    }

    public Task Stop()
    {
        return Task.CompletedTask;
    }
}