MoShiro/backend/fleet/src/Fleet.Server.Link/NodeConnection.cs

using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Threading.Channels;
using Fleet.Core;
using Fleet.Identity;
using Fleet.Rpc;

namespace Fleet.Server;

/// <summary>
/// 一个节点的 WebSocket 会话：挑战 → 签名认证 → hello → 状态推送 / 收件箱 / 产品 RPC。
/// 收到的消息按顺序处理，所以产品的处理函数不能等待“向同一个节点发起的请求”，否则会死锁。
/// </summary>
internal sealed partial class NodeConnection : IAsyncDisposable
{
    private static readonly TimeSpan s_authTimeout = TimeSpan.FromSeconds(15);
    private static readonly TimeSpan s_presenceRefresh = TimeSpan.FromSeconds(60);
    private static readonly TimeSpan s_statePushRetry = TimeSpan.FromSeconds(5);
    private static readonly TimeSpan s_closeTimeout = TimeSpan.FromSeconds(2);

    private readonly FleetLinkHost _host;
    private readonly WebSocket _socket;
    private readonly SemaphoreSlim _sendLock = new(1, 1);
    private readonly ConcurrentDictionary<long, TaskCompletionSource<RpcMessage>> _pending = new();
    private readonly Channel<bool> _stateSignal = Channel.CreateBounded<bool>(new BoundedChannelOptions(1) { FullMode = BoundedChannelFullMode.DropWrite });
    private readonly Dictionary<string, long> _lastPushedVersion = new(StringComparer.Ordinal);
    private readonly CancellationTokenSource _lifetime = new();
    private LinkCipher? _cipher;
    private long _nextRequestId;
    private long _lastPresenceTicks;
    private NodeRecord? _node;
    private NodeSessionInfo? _session;
    private string? _closeReason;
    private bool _disposed;
    private bool _helloDone;
    private int _violations;

    private readonly string? _remoteIp;

    public NodeConnection(FleetLinkHost host, WebSocket socket, string? remoteIp = null)
    {
        _host = host;
        _socket = socket;
        _remoteIp = remoteIp;
    }

    public NodeId NodeId => _node!.Id;

    public NodeSessionInfo Session => _session!;

    public async Task RunAsync(CancellationToken ct)
    {
        using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(ct, _lifetime.Token);
        CancellationToken token = linked.Token;
        bool registered = false;
        Task? pusher = null;
        try
        {
            if (!await AuthenticateAsync(token))
            {
                return;
            }

            NodeConnection? old = _host.Register(this);
            registered = true;
            if (old is not null)
            {
                await old.CloseAsync("replaced");
            }

            string status = StatusName(_node!.Status);
            FleetLog.Connected(_host.Logger, _node.Id.Value, status);
            await UpdatePresenceAsync(online: true, token);
            await FireAsync(NodeEventKind.Online, token);

            pusher = Task.Run(() => StatePushLoopAsync(token), CancellationToken.None);
            await ReceiveLoopAsync(token);
        }
        catch (OperationCanceledException)
        {
            // 关闭原因由触发取消的一方设置
        }
        catch (WebSocketException)
        {
            _closeReason = "connection lost";
        }
        finally
        {
            await _lifetime.CancelAsync();
            if (pusher is not null)
            {
                try
                {
                    await pusher;
                }
                catch (OperationCanceledException)
                {
                    // 正常结束
                }
            }

            foreach (TaskCompletionSource<RpcMessage> waiter in _pending.Values)
            {
                waiter.TrySetCanceled(CancellationToken.None);
            }

            if (registered)
            {
                _host.Unregister(this);
                FleetLog.Disconnected(_host.Logger, _node!.Id.Value, _closeReason ?? "closed");

                // 被新连接替换时，在线状态属于新连接。
                if (!_host.IsOnline(_node.Id))
                {
                    await UpdatePresenceAsync(online: false, CancellationToken.None);
                    await FireAsync(NodeEventKind.Offline, CancellationToken.None);
                }
            }

            if (_socket.State is WebSocketState.Open or WebSocketState.CloseReceived)
            {
                // 完整的关闭握手：发送 Close 并等待对端确认（有上限），这样对端能收到关闭原因
                try
                {
                    using var closeTimeout = new CancellationTokenSource(s_closeTimeout, _host.Time);
                    await _socket.CloseAsync(WebSocketCloseStatus.NormalClosure, _closeReason ?? "closed", closeTimeout.Token);
                }
                catch (WebSocketException)
                {
                    // 对端已经断开
                }
                catch (OperationCanceledException)
                {
                    // 对端没有在时限内确认：不再等待
                }
            }

        }
    }

    public async ValueTask DisposeAsync()
    {
        await _sendLock.WaitAsync();
        try
        {
            _disposed = true;
            _cipher?.Dispose();
            _lifetime.Dispose();
        }
        finally
        {
            // API callers may already be waiting to send; let them observe _disposed.
            _sendLock.Release();
        }
    }
}