squad-rain-ops-mini/Helpers/RCON/RconClient.cs

using System;
using System.Buffers;
using System.Collections.Concurrent;
using System.IO;
using System.Net.Sockets;
using System.Text;
using System.Threading;
using System.Threading.Channels;
using System.Threading.Tasks;

namespace RainOpsMini
{
    public enum RconPacketType : int
    {
        SERVERDATA_RESPONSE_VALUE = 0,
        SERVERDATA_EXECCOMMAND = 2,
        SERVERDATA_AUTH_RESPONSE = 2,
        SERVERDATA_AUTH = 3
    }

    public class RconPacket
    {
        public int Id { get; private set; }
        public int Type { get; private set; }
        public string Body { get; private set; }

        public RconPacket(int id, int type, string body)
        {
            Id = id;
            Type = type;
            Body = body;
        }

        public byte[] ToBytes()
        {
            // Body + Null + Null (Empty String)
            byte[] bodyBytes = Encoding.UTF8.GetBytes(Body);
            // Size = Id(4) + Type(4) + Body(N) + 2
            int packetSize = 4 + 4 + bodyBytes.Length + 2;

            using (MemoryStream ms = new MemoryStream())
            using (BinaryWriter writer = new BinaryWriter(ms))
            {
                writer.Write(packetSize);
                writer.Write(Id);
                writer.Write(Type);
                writer.Write(bodyBytes);
                writer.Write((byte)0); // Null for Body
                writer.Write((byte)0); // Null for Empty String
                return ms.ToArray();
            }
        }
    }

    public class RconClient : IDisposable
    {
        private TcpClient? _tcpClient;
        private NetworkStream? _stream;
        private readonly string _ip;
        private readonly int _port;
        private readonly string _password;
        private int _sequenceId = 0;
        private volatile bool _connected = false;
        private CancellationTokenSource? _connectionCts;

        // 接收通道
        private readonly Channel<RconPacket> _incomingChannel;
        // 路由表：RequestID -> ResponseChannel
        private readonly ConcurrentDictionary<int, Channel<RconPacket>> _commandChannels = new();
        
        // 发送锁，确保命令串行执行
        private readonly SemaphoreSlim _sendLock = new SemaphoreSlim(1, 1);

        public event Action<string>? OnLog;
        public event Action<string>? OnMessageReceived;
        public bool IsConnected => _connected;

        public RconClient(string ip, int port, string password)
        {
            _ip = ip;
            _port = port;
            _password = password;
            // 无界通道，确保高吞吐时不阻塞读取
            _incomingChannel = Channel.CreateUnbounded<RconPacket>();
        }

        public async Task ConnectAsync()
        {
            _tcpClient = new TcpClient();
            await _tcpClient.ConnectAsync(_ip, _port);
            _stream = _tcpClient.GetStream();
            _connected = true;
            _connectionCts = new CancellationTokenSource();

            // 启动接收循环
            _ = ReceiveLoopAsync(_connectionCts.Token);
            // 启动包路由循环
            _ = PacketRouterAsync(_connectionCts.Token);
            // 启动保活循环
            _ = KeepAliveLoopAsync(_connectionCts.Token);

            await AuthenticateAsync();
        }

        private async Task PacketRouterAsync(CancellationToken token)
        {
            try
            {
                while (await _incomingChannel.Reader.WaitToReadAsync(token))
                {
                    while (_incomingChannel.Reader.TryRead(out var packet))
                    {
                        // DEBUG: Log every packet
                        OnLog?.Invoke($"[PacketRouter] Received Packet: Id={packet.Id}, Type={packet.Type}, BodyLen={packet.Body.Length}");

                        if (_commandChannels.TryGetValue(packet.Id, out var channel))
                        {
                            await channel.Writer.WriteAsync(packet, token);
                        }
                        else
                        {
                            // 所有的非命令响应包（或ID不匹配的包）都视为异步消息/日志
                            if (!string.IsNullOrEmpty(packet.Body))
                            {
                                OnMessageReceived?.Invoke(packet.Body);
                            }
                            else
                            {
                                // Log empty packets that don't match commands?
                                // OnLog?.Invoke($"[PacketRouter] Ignored empty packet Id={packet.Id}");
                            }
                        }
                    }
                }
            }
            catch (OperationCanceledException)
            {
                // Normal shutdown
            }
            catch (Exception ex)
            {
                OnLog?.Invoke($"[PacketRouter] Error: {ex.Message}");
            }
        }

        private async Task AuthenticateAsync()
        {
            int authId = Interlocked.Increment(ref _sequenceId);
            var packet = new RconPacket(authId, (int)RconPacketType.SERVERDATA_AUTH, _password);
            
            // 准备接收通道
            var responseChannel = Channel.CreateUnbounded<RconPacket>();
            _commandChannels[authId] = responseChannel;

            try
            {
                await SendPacketAsync(packet);

                // 等待认证响应
                using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
                try
                {
                    while (true)
                    {
                        var response = await responseChannel.Reader.ReadAsync(cts.Token);
                        if (response.Type == (int)RconPacketType.SERVERDATA_AUTH_RESPONSE)
                        {
                            if (response.Id == -1) throw new Exception("RCON Authentication Failed");
                            if (response.Id == authId) return; // Success
                        }
                        // 认证阶段的 Log 可以在这里处理，也可以忽略（PacketRouter 不会转发不匹配 ID 的包到这里，
                        // 但如果服务器发送了 ID 匹配但类型不是 AUTH_RESPONSE 的包... RCON 协议通常不会这样）
                        // 注意：Source RCON 中，如果认证失败，ID 会是 -1。
                        // 如果 ID 不匹配，PacketRouter 已经分发给 OnMessageReceived 了。
                        // 所以这里读到的肯定是 ID == authId 的包。
                        
                        // 不过，AUTH_RESPONSE 的 ID 应该回传 authId。
                        // 唯一例外是认证失败回传 -1。
                        // 可是如果回传 -1，PacketRouter 找不到 -1 的 channel，就会当做 MessageReceived 处理吗？
                        // 这是一个问题。
                        // 认证是一个特殊情况。
                        // 我们应该也监听 -1 吗？或者让 PacketRouter 特殊处理 AUTH_RESPONSE？
                        
                        // 修正：如果认证失败，服务器返回 ID=-1。
                        // 我们需要注册 ID=-1 吗？不，因为 -1 是通用的。
                        // 简单做法：PacketRouter 看到 SERVERDATA_AUTH_RESPONSE 且 ID=-1 时，
                        // 应该尝试通知当前的认证等待者？
                        // 或者，我们在 AuthenticateAsync 里不做这么复杂，
                        // 认证阶段通常没有并发命令。
                        // 我们可以临时注册 -1 到 responseChannel？
                    }
                }
                catch (OperationCanceledException)
                {
                    throw new TimeoutException("RCON Authentication Timed Out");
                }
            }
            finally
            {
                _commandChannels.TryRemove(authId, out _);
            }
        }

        public async Task<string> SendCommandAsync(string command, int timeoutSeconds = 10)
        {
            if (!_connected) throw new InvalidOperationException("Not connected");

            await _sendLock.WaitAsync();
            try
            {
                int cmdId = Interlocked.Increment(ref _sequenceId);
                var packet = new RconPacket(cmdId, (int)RconPacketType.SERVERDATA_EXECCOMMAND, command);

                // 1. 注册响应通道
                var responseChannel = Channel.CreateUnbounded<RconPacket>();
                _commandChannels[cmdId] = responseChannel;

                try
                {
                    // 2. 发送命令
                    await SendPacketAsync(packet);

                    // 3. 收集响应
                    StringBuilder result = new StringBuilder();
                    bool hasStartedReceiving = false;

                    // 总超时
                    using var totalCts = new CancellationTokenSource(TimeSpan.FromSeconds(timeoutSeconds));

                    try
                    {
                        while (!totalCts.Token.IsCancellationRequested)
                        {
                            // 动态等待策略
                            int waitMs = hasStartedReceiving ? 300 : 3000;

                            using var packetCts = CancellationTokenSource.CreateLinkedTokenSource(totalCts.Token);
                            packetCts.CancelAfter(waitMs);

                            try
                            {
                                // 只读取属于该 ID 的包
                                var response = await responseChannel.Reader.ReadAsync(packetCts.Token);

                                hasStartedReceiving = true;
                                result.Append(response.Body);
                            }
                            catch (OperationCanceledException)
                            {
                                if (packetCts.Token.IsCancellationRequested && !totalCts.Token.IsCancellationRequested)
                                {
                                    // waitMs 超时触发
                                    if (hasStartedReceiving)
                                    {
                                        return result.ToString();
                                    }
                                    else
                                    {
                                        return string.Empty;
                                    }
                                }
                                throw; // 总超时
                            }
                        }
                    }
                    catch (OperationCanceledException)
                    {
                        // 总超时
                        if (result.Length > 0) return result.ToString();
                        return "Error: Timeout waiting for response";
                    }

                    return result.ToString();
                }
                finally
                {
                    _commandChannels.TryRemove(cmdId, out _);
                }
            }
            finally
            {
                _sendLock.Release();
            }
        }

        private async Task SendPacketAsync(RconPacket packet)
        {
            if (_stream == null) return;
            byte[] data = packet.ToBytes();
            await _stream.WriteAsync(data, 0, data.Length);
            await _stream.FlushAsync();
        }

        private async Task ReceiveLoopAsync(CancellationToken token)
        {
            try
            {
                byte[] sizeBuffer = new byte[4];

                while (!token.IsCancellationRequested && _stream != null)
                {
                    // 1. Read Size (4 bytes)
                    if (!await ReadExactAsync(_stream, sizeBuffer, 4, token)) break;
                    int size = BitConverter.ToInt32(sizeBuffer, 0);

                    if (size < 0 || size > 10 * 1024 * 1024) // Sanity check (10MB limit)
                    {
                        Log($"Invalid packet size: {size}. Disconnecting.");
                        break;
                    }

                    // 2. Read Payload (Size bytes)
                    byte[] payload = new byte[size];
                    if (!await ReadExactAsync(_stream, payload, size, token)) break;

                    // 3. Parse
                    var packet = ParsePacket(payload);

                    // 4. Dispatch to Channel
                    await _incomingChannel.Writer.WriteAsync(packet, token);
                }
            }
            catch (Exception ex) when (!(ex is OperationCanceledException))
            {
                Log($"Receive Error: {ex.Message}");
            }
            finally
            {
                _connected = false;
                _incomingChannel.Writer.TryComplete();
            }
        }

        private async Task<bool> ReadExactAsync(NetworkStream stream, byte[] buffer, int count, CancellationToken token)
        {
            int totalRead = 0;
            while (totalRead < count)
            {
                int read = await stream.ReadAsync(buffer, totalRead, count - totalRead, token);
                if (read == 0) return false;
                totalRead += read;
            }
            return true;
        }

        private RconPacket ParsePacket(byte[] payload)
        {
            using var ms = new MemoryStream(payload);
            using var reader = new BinaryReader(ms);

            int id = reader.ReadInt32();
            int type = reader.ReadInt32();

            // 剩余部分是 Body + 2个空字节 (Null Terminator)
            // 我们需要安全地读取 Body 字符串
            int bodyLen = (int)(ms.Length - ms.Position);
            byte[] bodyBytes = reader.ReadBytes(bodyLen);

            // 从尾部开始扫描，去掉所有的 0
            int realLen = bodyLen;
            while (realLen > 0 && bodyBytes[realLen - 1] == 0)
            {
                realLen--;
            }

            string body = Encoding.UTF8.GetString(bodyBytes, 0, realLen);
            return new RconPacket(id, type, body);
        }

        private async Task KeepAliveLoopAsync(CancellationToken token)
        {
            // 保持连接活跃，如果长时间没有命令交互
            while (!token.IsCancellationRequested)
            {
                await Task.Delay(TimeSpan.FromSeconds(30), token);
                // 可以在这里实现心跳，目前仅依靠 TCP KeepAlive
            }
        }

        private void Log(string msg) => OnLog?.Invoke($"[{DateTime.Now:HH:mm:ss}] {msg}");

        public void Dispose()
        {
            _connectionCts?.Cancel();
            _tcpClient?.Dispose();
            _sendLock.Dispose();
        }
    }
}