ShrinkNetwork/Runtime/Transport/Kcp/ShrinkKcpPeer.cs

using System;
using System.Runtime.InteropServices;
using kcp;

namespace ShrinkNetwork
{
    internal sealed unsafe class ShrinkKcpPeer : IDisposable
    {
        private readonly object _syncRoot = new();
        private readonly Action<byte[]> _sendDatagram;
        private readonly byte[] _receiveBuffer;
        private readonly GCHandle _selfHandle;
        private IKCPCB* _kcp;
        private bool _disposed;
        private uint _nextUpdateTime;

        public ShrinkKcpPeer(uint conversationId, ShrinkKcpTransportOptions options, Action<byte[]> sendDatagram)
        {
            if (conversationId == 0)
                throw new ArgumentOutOfRangeException(nameof(conversationId));
            if (options == null)
                throw new ArgumentNullException(nameof(options));
            if (sendDatagram == null)
                throw new ArgumentNullException(nameof(sendDatagram));

            options.Validate();

            ConversationId = conversationId;
            _sendDatagram = sendDatagram;
            _receiveBuffer = new byte[options.MaxMessageSize];
            _selfHandle = GCHandle.Alloc(this);
            _kcp = KCP_INTERFACE.ikcp_create(conversationId, (void*)GCHandle.ToIntPtr(_selfHandle));
            if (_kcp == null)
                throw new InvalidOperationException("KCP create failed.");

            KCP_INTERFACE.ikcp_setmtu(_kcp, options.Mtu);
            KCP_INTERFACE.ikcp_wndsize(_kcp, options.SendWindow, options.ReceiveWindow);
            KCP_INTERFACE.ikcp_nodelay(_kcp, options.NoDelay ? 1 : 0, options.Interval, options.Resend,
                options.DisableCongestionControl ? 1 : 0);
            KCP_INTERFACE.ikcp_setoutput(_kcp, &HandleOutput);

            LastReceiveUtcTicks = DateTime.UtcNow.Ticks;
            _nextUpdateTime = GetNowMs();
        }

        public uint ConversationId { get; }
        public long LastReceiveUtcTicks { get; private set; }

        public void Input(byte[] datagram, int offset, int length, Action<byte[]> onPacket)
        {
            if (datagram == null)
                throw new ArgumentNullException(nameof(datagram));
            if (offset < 0 || offset > datagram.Length)
                throw new ArgumentOutOfRangeException(nameof(offset));
            if (length < 0 || offset + length > datagram.Length)
                throw new ArgumentOutOfRangeException(nameof(length));
            if (length == 0)
                return;
            if (onPacket == null)
                throw new ArgumentNullException(nameof(onPacket));

            lock (_syncRoot)
            {
                EnsureNotDisposed();

                fixed (byte* basePtr = datagram)
                {
                    var result = KCP_INTERFACE.ikcp_input(_kcp, basePtr + offset, length);
                    if (result < 0)
                        throw new InvalidOperationException($"KCP input failed. Result={result}");
                }

                LastReceiveUtcTicks = DateTime.UtcNow.Ticks;
                UpdateInternal(GetNowMs());
                DrainReceiveQueue(onPacket);
            }
        }

        public void Send(byte[] payload)
        {
            if (payload == null)
                throw new ArgumentNullException(nameof(payload));
            if (payload.Length == 0)
                return;

            lock (_syncRoot)
            {
                EnsureNotDisposed();

                fixed (byte* payloadPtr = payload)
                {
                    var result = KCP_INTERFACE.ikcp_send(_kcp, payloadPtr, payload.Length);
                    if (result < 0)
                        throw new InvalidOperationException($"KCP send failed. Result={result}");
                }

                UpdateInternal(GetNowMs());
            }
        }

        public void Tick(Action<byte[]> onPacket)
        {
            if (onPacket == null)
                throw new ArgumentNullException(nameof(onPacket));

            lock (_syncRoot)
            {
                if (_disposed)
                    return;

                var now = GetNowMs();
                if (now < _nextUpdateTime)
                    return;

                UpdateInternal(now);
                DrainReceiveQueue(onPacket);
            }
        }

        public void Dispose()
        {
            lock (_syncRoot)
            {
                if (_disposed)
                    return;

                _disposed = true;
                if (_kcp != null)
                {
                    KCP_INTERFACE.ikcp_release(_kcp);
                    _kcp = null;
                }

                if (_selfHandle.IsAllocated)
                    _selfHandle.Free();
            }
        }

        private void EnsureNotDisposed()
        {
            if (_disposed || _kcp == null)
                throw new ObjectDisposedException(nameof(ShrinkKcpPeer));
        }

        private void UpdateInternal(uint now)
        {
            KCP_INTERFACE.ikcp_update(_kcp, now);
            _nextUpdateTime = KCP_INTERFACE.ikcp_check(_kcp, now);
        }

        private void DrainReceiveQueue(Action<byte[]> onPacket)
        {
            while (true)
            {
                var peekSize = KCP_INTERFACE.ikcp_peeksize(_kcp);
                if (peekSize < 0)
                    return;

                if (peekSize > _receiveBuffer.Length)
                    throw new InvalidOperationException(
                        $"KCP message too large. Size={peekSize}, Buffer={_receiveBuffer.Length}");

                var received = 0;
                fixed (byte* receivePtr = _receiveBuffer)
                {
                    received = KCP_INTERFACE.ikcp_recv(_kcp, receivePtr, peekSize);
                }

                if (received <= 0)
                    return;

                LastReceiveUtcTicks = DateTime.UtcNow.Ticks;
                var packet = new byte[received];
                Buffer.BlockCopy(_receiveBuffer, 0, packet, 0, received);
                onPacket(packet);
            }
        }

        private static uint GetNowMs()
        {
            // KCP expects a wrapping uint32 millisecond clock; Environment.TickCount
            // matches that contract and still works on Unity's .NET Framework target.
            return unchecked((uint)Environment.TickCount);
        }

        private int HandleOutputInternal(byte* buffer, int length)
        {
            if (length <= 0)
                return 0;

            var datagram = new byte[length];
            Marshal.Copy((IntPtr)buffer, datagram, 0, length);
            _sendDatagram(datagram);
            return 0;
        }

        private static int HandleOutput(byte* buffer, int length, IKCPCB* kcp, void* user)
        {
            if (user == null)
                return -1;

            var handle = GCHandle.FromIntPtr((IntPtr)user);
            if (handle.Target is not ShrinkKcpPeer peer)
                return -1;

            return peer.HandleOutputInternal(buffer, length);
        }
    }
}