Ananta-ps/Ananta.SDK/Network/TcpSession.cs

using System.Net.Sockets;
using System.Security.Cryptography;
using Ananta.SDK.Logging;
using Ananta.SDK.Rpc;

namespace Ananta.SDK.Network;

public sealed class TcpSession
{
    private readonly SemaphoreSlim _sendLock = new(1, 1);
    private readonly CancellationTokenSource _closed = new();
    /// <summary>会话关闭时触发。逐字对齐捐赠端 <c>TcpSession.ClosedToken</c>（<c>Ananta.SDK/Network/TcpSession.cs:13</c>）。
    /// ⛑ 主线原本已有 <c>_closed</c> 且 <c>Close()</c> 会 Cancel，只是没公开 → 这里<b>只加属性，不改生命周期</b>。</summary>
    public CancellationToken ClosedToken => _closed.Token;
    private Task? _heartbeat;
    private int _isClosed;
    private UxChaCha8? _readCryptor;
    private UxChaCha8? _writeCryptor;

    /// <summary>服务端 → 客户端传输层 keepalive 间隔（毫秒）。<b>0 = 关闭</b>（沿用旧行为）。</summary>
    public int HeartbeatMs { get; }
    /// <summary>单次加密写的上限（毫秒）；超时会退役这条连接。<b>0 = 不设上限</b>（沿用旧行为）。</summary>
    public int WriteTimeoutMs { get; }

    public string Id { get; } = Guid.NewGuid().ToString("N")[..6];
    public TcpClient Client { get; }
    public NetworkStream Stream { get; }
    public ServerLogger Log { get; }
    public Dictionary<string, object> Items { get; } = new();
    public bool CryptReady => _readCryptor is not null && _writeCryptor is not null;

    public TcpSession(TcpClient client, string scope, int heartbeatMs = 0, int writeTimeoutMs = 15000)
    {
        Client = client;
        Stream = client.GetStream();
        Log = new ServerLogger($"{scope}:{Id}");
        HeartbeatMs = heartbeatMs;
        WriteTimeoutMs = writeTimeoutMs;
        // Small real-time RPCs should not wait for Nagle/delayed-ACK batching.
        Client.NoDelay = true;
    }

    public Task<Frame?> ReadFrameAsync(CancellationToken token)
        => FrameCodec.ReadFrameAsync(Stream, _readCryptor, token);

    public async Task SendHandshakeAsync(ReadOnlyMemory<byte> clientHandshake, CancellationToken token)
    {
        if (clientHandshake.Length < 20)
            throw new InvalidDataException($"4229938 C2S handshake is only {clientHandshake.Length} bytes");

        // C2S handshake layout: 4229938 SendHandShakeFrontEnd writes
        // magic(4) + the plain 16-byte UXAES.Generate() key at [4..20)
        // (its UXRSA.Encrypt call is patched out client-side), padded to 260.
        byte[] aesKey = clientHandshake.Slice(4, 16).ToArray();
        Log.Info($"HS raw[0..32]={Convert.ToHexString(clientHandshake.Slice(0, 32).ToArray())}");

        // Decrypted S2C handshake blob expected by TrySetupChaChaOnHandshake:
        //   [0..32)  = client's write key  => server READ key
        //   [32..64) = client's read key   => server WRITE key
        //   [64..76) = shared 96-bit nonce
        // Empirically verified against 4229938: the client encrypts its frames
        // with blob[0..32) (FrameCodec raw-head capture + keystream match).
        byte[] clientWriteKey = RandomNumberGenerator.GetBytes(32);
        byte[] clientReadKey = RandomNumberGenerator.GetBytes(32);
        byte[] streamNonce = RandomNumberGenerator.GetBytes(12);
        byte[] chaChaHead = new byte[76];
        clientWriteKey.CopyTo(chaChaHead, 0);
        clientReadKey.CopyTo(chaChaHead, 32);
        streamNonce.CopyTo(chaChaHead, 64);

        byte[] aesNonce = RandomNumberGenerator.GetBytes(12);
        byte[] cipher = new byte[76];
        byte[] tag = new byte[16];
        using (var aes = new AesGcm(aesKey, 16))
            aes.Encrypt(aesNonce, chaChaHead, cipher, tag);

        byte[] payload = new byte[368];
        BitConverter.GetBytes(1).CopyTo(payload, 0);
        BitConverter.GetBytes(10).CopyTo(payload, 4);
        aesNonce.CopyTo(payload, 8);
        cipher.CopyTo(payload, 20);
        tag.CopyTo(payload, 96);
        // payload[112..368] is the RSA-2048 signature field. The client's
        // UXRSA.Verify result branch is patched out to accept the local
        // unsigned handshake.

        // The handshake reply itself is plaintext. Encryption becomes active
        // only for the next frame in each direction — which is why the keys are
        // published *after* this write instead of before it.
        await WriteRawAsync(0x01, payload, token);
        _readCryptor = new UxChaCha8(clientWriteKey, streamNonce);
        _writeCryptor = new UxChaCha8(clientReadKey, streamNonce);

        Log.Info($"HS keys clientWrite={Convert.ToHexString(clientWriteKey)} clientRead={Convert.ToHexString(clientReadKey)} nonce={Convert.ToHexString(streamNonce)} aes={Convert.ToHexString(aesKey)}");

        Log.Info("HS ok -> UXChaCha8 enabled");

        if (HeartbeatMs > 0) _heartbeat ??= SendHeartbeatsAsync();
    }

    private async Task SendHeartbeatsAsync()
    {
        // Keep an inbound liveness signal on every socket, including idle gates,
        // independently of the client's RPC and heartbeat scheduling.
        using var timer = new PeriodicTimer(TimeSpan.FromMilliseconds(HeartbeatMs));
        try
        {
            while (await timer.WaitForNextTickAsync(_closed.Token))
                await SendFrameAsync(0x03, BitConverter.GetBytes(Environment.TickCount64), _closed.Token);
        }
        catch (OperationCanceledException) when (_closed.IsCancellationRequested) { }
        catch (Exception ex)
        {
            Log.Warn($"heartbeat send stopped: {ex.Message}");
            Close();
        }
    }

    public Task SendFrameAsync(byte mode, byte[] body, CancellationToken token)
        => WriteEncryptedFrameAsync(mode, body, token);

    public Task NotifyAsync(uint methodId, byte[] body, CancellationToken token)
    {
        RuntimeLogs.OutgoingNotify(Log.Scope, methodId, body);
        return WriteEncryptedFrameAsync(0x09, RpcPacketWriter.WriteNotify(methodId, body), token);
    }

    public Task ReturnAsync(uint methodId, int invokeId, int err, byte[] body, CancellationToken token)
    {
        Log.Info($"<- ret {RpcMethodNames.Display(methodId)} #{invokeId} e={err} {body.Length}b");
        RuntimeLogs.OutgoingReturn(Log.Scope, methodId, invokeId, err, body);
        return WriteEncryptedFrameAsync(0x09, RpcPacketWriter.WriteReturn(methodId, invokeId, err, body), token);
    }

    private Task WriteRawAsync(byte mode, ReadOnlyMemory<byte> body, CancellationToken token)
        => WriteFrameAsync(mode, body, null, token);

    private async Task WriteFrameAsync(byte mode, ReadOnlyMemory<byte> body, UxChaCha8? cryptor, CancellationToken token)
    {
        using var deadline = WriteTimeoutMs > 0
            ? CancellationTokenSource.CreateLinkedTokenSource(token, _closed.Token)
            : null;
        if (deadline is not null) deadline.CancelAfter(WriteTimeoutMs);
        var linked = deadline?.Token ?? token;
        var acquired = false;
        try
        {
            await _sendLock.WaitAsync(linked);
            acquired = true;
            await FrameCodec.WriteFrameAsync(Stream, mode, body, cryptor, linked);
        }
        catch (OperationCanceledException) when (!_closed.IsCancellationRequested && !token.IsCancellationRequested)
        {
            // A partially-written encrypted frame cannot be retried: the stream cipher already
            // advanced. Retire only this socket so everyone else keeps being served.
            Close();
            throw new IOException($"TCP send cancelled or exceeded the {(WriteTimeoutMs > 0 ? WriteTimeoutMs : 0)}-ms write deadline");
        }
        catch (Exception ex) when (ex is IOException or SocketException)
        {
            Close();
            throw;
        }
        finally
        {
            if (acquired) _sendLock.Release();
        }
    }

    private Task WriteEncryptedFrameAsync(byte mode, ReadOnlyMemory<byte> body, CancellationToken token)
        => WriteFrameAsync(mode, body, _writeCryptor, token);

    public void Close()
    {
        if (Interlocked.Exchange(ref _isClosed, 1) != 0) return;
        _closed.Cancel();
        try { Client.Close(); } catch { }
        // Pending async sends may still own or await this semaphore. Disposing it
        // here races their finally/Release and hides the actual socket closure.
        // It owns no wait handle; let it be collected with this session.
    }
}