MoShiro/backend/fleet/src/Fleet.Link/FleetLink.cs

using System.Net.WebSockets;
using System.Text.Json;
using System.Text.Json.Serialization.Metadata;
using Fleet.Core;
using Fleet.Identity;
using Fleet.Rpc;

namespace Fleet.Link;

public sealed record LinkContext(FleetLink Link);

/// <summary>期望状态的处理者：应用完整快照，返回已应用的版本和被拒绝的条目（带结构化原因）。</summary>
public interface IStateHandler
{
    string Type { get; }

    ValueTask<StateSyncResult> ApplyAsync(long version, JsonElement document, CancellationToken ct);
}

public sealed class StateHandler<TDoc> : IStateHandler
{
    private readonly JsonTypeInfo<TDoc> _info;
    private readonly Func<long, TDoc, CancellationToken, ValueTask<StateSyncResult>> _apply;

    public StateHandler(string type, JsonTypeInfo<TDoc> info, Func<long, TDoc, CancellationToken, ValueTask<StateSyncResult>> apply)
    {
        Type = type;
        _info = info;
        _apply = apply;
    }

    public string Type { get; }

    public async ValueTask<StateSyncResult> ApplyAsync(long version, JsonElement document, CancellationToken ct)
    {
        TDoc? doc = document.Deserialize(_info);
        return doc is null ? new StateSyncResult { AppliedVersion = 0 } : await _apply(version, doc, ct);
    }
}

public sealed class FleetLinkOptions
{
    /// <summary>面板的地址（http 或 https），连接路径固定为 <c>/api/agent/v1/rpc</c>。</summary>
    public required Uri Server { get; init; }

    public required NodeIdentity Identity { get; init; }

    public required NodeId NodeId { get; init; }

    public string NodeVersion { get; init; } = "0.0.0";

    /// <summary>框架之外的能力，例如 <c>moshiro.v1</c> 与 <c>moshiro.kinds:udp_echo,tcp_echo</c>。</summary>
    public IReadOnlyList<string> Capabilities { get; init; } = [];

    /// <summary>附加到 <c>fleet.hello</c> 的产品元数据。</summary>
    public Func<JsonElement?>? HelloMeta { get; init; }

    /// <summary>附加到 <c>fleet.heartbeat</c> 的进程内计数（不能包含主机信息）。</summary>
    public Func<JsonElement?>? HeartbeatExtra { get; init; }

    public TimeSpan HeartbeatInterval { get; init; } = TimeSpan.FromSeconds(10);

    public int MaxFrameBytes { get; init; } = 16 << 20;

    /// <summary>发件箱在途字节上限（服务端慢时不会把内存撑爆）。</summary>
    public long MaxInflightBytes { get; init; } = 4 << 20;

    public Backoff Backoff { get; init; } = new();
}

/// <summary>
/// 节点到服务端的连接：挑战 → 签名认证 → hello → 状态调和 / 可靠发件箱 / 心跳，断线之后指数退避重连。
/// 被服务端撤销（<c>fleet.revoked</c>）或禁用时停止重试。
/// </summary>
public sealed class FleetLink
{
    private const string RpcPath = "/api/agent/v1/rpc";

    private readonly FleetLinkOptions _options;
    private readonly RpcRouter<LinkContext> _router;
    private readonly Dictionary<string, IStateHandler> _handlers;
    private readonly TimeProvider _time;
    private readonly INodeLog _log;
    private volatile RpcSession<LinkContext>? _session;
    private volatile bool _stopped;

    public FleetLink(FleetLinkOptions options, RpcRouter<LinkContext> router, IEnumerable<IStateHandler> handlers, Outbox outbox, TimeProvider time, INodeLog log)
    {
        _options = options;
        _router = router;
        _handlers = handlers.ToDictionary(h => h.Type, StringComparer.Ordinal);
        Outbox = outbox;
        _time = time;
        _log = log;

        _router
            .On(FleetMethods.StateSync, FleetRpcContext.Default.StateSyncParams, FleetRpcContext.Default.StateSyncResult, HandleStateSyncAsync)
            .OnNotify(FleetMethods.OutboxAck, FleetRpcContext.Default.OutboxAckParams, HandleAckAsync)
            .OnNotify(FleetMethods.Revoked, FleetRpcContext.Default.RevokedParams, HandleRevokedAsync);
    }

    public Outbox Outbox { get; }

    public bool IsConnected => _session is not null;

    public RpcSession<LinkContext>? CurrentSession => _session;

    /// <summary>入站方法的路由器（只读，测试用它核对声明表和实际注册的处理函数）。</summary>
    public RpcRouter<LinkContext> Router => _router;

    /// <summary>被服务端撤销或禁用：不会再重连。</summary>
    public bool Revoked => _stopped;

    /// <summary>每次连接并完成 hello 之后调用。</summary>
    public Func<RpcSession<LinkContext>, CancellationToken, ValueTask>? OnConnected { get; set; }

    public Action<string>? OnDisconnected { get; set; }

    public async Task RunAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested && !_stopped)
        {
            string reason;
            try
            {
                reason = await RunOnceAsync(ct);
                _options.Backoff.Reset();
            }
            catch (OperationCanceledException) when (ct.IsCancellationRequested)
            {
                return;
            }
            catch (Exception ex) when (ex is WebSocketException or IOException or HttpRequestException or InvalidOperationException or TimeoutException)
            {
                reason = ex.Message;
            }

            OnDisconnected?.Invoke(reason);
            if (ct.IsCancellationRequested)
            {
                return;
            }

            if (_stopped)
            {
                _log.Warn("link stopped", ("err", reason));
                return;
            }

            TimeSpan delay = _options.Backoff.Next();
            _log.Warn("link down, retrying", ("err", reason), ("retry_ms", (long)delay.TotalMilliseconds));
            await Task.Delay(delay, _time, ct);
        }
    }

    private async Task<string> RunOnceAsync(CancellationToken ct)
    {
        using var socket = new ClientWebSocket();
        socket.Options.KeepAliveInterval = TimeSpan.FromSeconds(15);
        socket.Options.KeepAliveTimeout = TimeSpan.FromSeconds(20); // 对端 20 秒没有回 Pong 就判定连接已死
        var builder = new UriBuilder(_options.Server) { Scheme = _options.Server.Scheme == "https" ? "wss" : "ws", Path = RpcPath };
        await socket.ConnectAsync(builder.Uri, ct);
        await using var transport = new WsTransport(socket, _options.MaxFrameBytes);

        string? failure = await AuthenticateAsync(transport, ct);
        if (failure is not null)
        {
            return failure;
        }

        var session = new RpcSession<LinkContext>(transport, _router, new LinkContext(this), _time);
        using CancellationTokenSource sessionCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        Task run = session.RunAsync(sessionCts.Token);
        try
        {
            var hello = new HelloParams
            {
                NodeVersion = _options.NodeVersion,
                Capabilities = [.. _options.Capabilities],
                NowUtcMs = _time.GetUtcNow().ToUnixTimeMilliseconds(),
                Meta = _options.HelloMeta?.Invoke(),
            };
            Result<HelloResult> helloResult = await session.RequestAsync(FleetMethods.Hello, hello, FleetRpcContext.Default.HelloParams, FleetRpcContext.Default.HelloResult, sessionCts.Token);
            if (!helloResult.IsSuccess)
            {
                return "hello failed: " + helloResult.Error.Code;
            }

            Outbox.ResetSent();
            _session = session;
            _log.Info("link up", ("result", "connected"));
            Task pump = PumpAsync(session, sessionCts.Token);
            Task heartbeat = HeartbeatAsync(session, sessionCts.Token);
            if (OnConnected is not null)
            {
                await OnConnected(session, sessionCts.Token);
            }

            Task finished = await Task.WhenAny(run, pump, heartbeat);
            await sessionCts.CancelAsync();
            await Task.WhenAll(Quiet(run), Quiet(pump), Quiet(heartbeat));
            return finished == run ? "closed by server" : "link task failed";
        }
        finally
        {
            _session = null;
            await sessionCts.CancelAsync();
            await Quiet(run);
            await session.DisposeAsync();
        }
    }

    /// <summary>挑战 → 签名 → 认证结果。失败返回原因；服务端禁用该节点时不再重试。</summary>
    private async Task<string?> AuthenticateAsync(WsTransport transport, CancellationToken ct)
    {
        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(15), _time);
        using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(ct, timeout.Token);
        RpcMessage? challenge = await ReceiveMessageAsync(transport, linked.Token);
        if (challenge?.Method != FleetMethods.Challenge || challenge.Params is null)
        {
            return "no challenge from the server";
        }

        ChallengeParams nonce = challenge.Params.Value.Deserialize(FleetRpcContext.Default.ChallengeParams)!;
        string signature = Convert.ToBase64String(_options.Identity.Sign(SignaturePurpose.Auth, _options.NodeId.Value, nonce.Nonce));
        RpcMessage request = RpcMessage.Request(1, FleetMethods.Auth, new AuthParams { NodeId = _options.NodeId.Value, Signature = signature }, FleetRpcContext.Default.AuthParams);
        await transport.SendAsync(new RpcFrame(FrameKind.Text, JsonSerializer.SerializeToUtf8Bytes(request, FleetRpcContext.Default.RpcMessage)), linked.Token);

        while (true)
        {
            RpcMessage? reply = await ReceiveMessageAsync(transport, linked.Token);
            if (reply is null)
            {
                return "connection closed during authentication";
            }

            if (reply.Method == FleetMethods.Revoked)
            {
                _stopped = true;
                return "revoked";
            }

            if (!reply.IsResponse)
            {
                continue;
            }

            if (reply.Error is { } error)
            {
                if (error.Data?.Code is "fleet_node_disabled")
                {
                    _stopped = true;
                }

                return "authentication failed: " + error.Data?.Code;
            }

            return null;
        }
    }

    private static async Task<RpcMessage?> ReceiveMessageAsync(WsTransport transport, CancellationToken ct)
    {
        RpcFrame? frame = await transport.ReceiveAsync(ct);
        return frame is { Kind: FrameKind.Text } f ? JsonSerializer.Deserialize(f.Data.Span, FleetRpcContext.Default.RpcMessage) : null;
    }

    private async Task PumpAsync(RpcSession<LinkContext> session, CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            byte[] frame = await Outbox.TakeNextAsync(_options.MaxInflightBytes, ct);
            await session.SendBinaryAsync(frame, ct);
        }
    }

    private async Task HeartbeatAsync(RpcSession<LinkContext> session, CancellationToken ct)
    {
        using var timer = new PeriodicTimer(_options.HeartbeatInterval, _time);
        while (await timer.WaitForNextTickAsync(ct))
        {
            await session.NotifyAsync(FleetMethods.Heartbeat, new HeartbeatParams { Extra = _options.HeartbeatExtra?.Invoke() }, FleetRpcContext.Default.HeartbeatParams, ct);
        }
    }

    private async ValueTask<Result<StateSyncResult>> HandleStateSyncAsync(LinkContext context, StateSyncParams parameters, CancellationToken ct)
    {
        if (!_handlers.TryGetValue(parameters.Type, out IStateHandler? handler))
        {
            return Result<StateSyncResult>.Fail(FleetErrors.Unsupported, $"unknown state type {parameters.Type}");
        }

        return Result<StateSyncResult>.Ok(await handler.ApplyAsync(parameters.Version, parameters.Doc, ct));
    }

    private ValueTask HandleAckAsync(LinkContext context, OutboxAckParams parameters, CancellationToken ct)
    {
        Outbox.Acknowledge(parameters.Ids.Concat(parameters.Duplicates).Select(id => Guid.TryParse(id, out Guid g) ? g : Guid.Empty));
        return ValueTask.CompletedTask;
    }

    private ValueTask HandleRevokedAsync(LinkContext context, RevokedParams parameters, CancellationToken ct)
    {
        _stopped = true;
        _log.Warn("revoked by the server", ("err", parameters.Reason));
        return ValueTask.CompletedTask;
    }

    private static async Task Quiet(Task task)
    {
        try
        {
            await task;
        }
        catch (Exception ex) when (ex is OperationCanceledException or WebSocketException or IOException or ObjectDisposedException)
        {
            // 会话结束时的正常异常
        }
    }
}