zcbot/windows-node/Zcbot.WindowsNode/NodeConnectionLoop.cs

using System.Net;
using System.Net.WebSockets;
using System.Reflection;
using System.Runtime.InteropServices;
using System.Text;
using System.Text.Json;

namespace Zcbot.WindowsNode;

internal sealed class NodeConnectionLoop(NodeConfig config, Action<NodeStatus>? statusChanged = null)
{
    private static readonly TimeSpan[] Backoff =
    [
        TimeSpan.FromSeconds(1),
        TimeSpan.FromSeconds(2),
        TimeSpan.FromSeconds(5),
        TimeSpan.FromSeconds(10),
        TimeSpan.FromSeconds(30),
        TimeSpan.FromSeconds(60),
    ];

    internal async Task RunAsync(CancellationToken cancellationToken)
    {
        var attempt = 0;
        while (!cancellationToken.IsCancellationRequested)
        {
            try
            {
                Report(NodeState.Connecting, "正在连接 zcbot…");
                await ConnectOnceAsync(cancellationToken);
                attempt = 0;
            }
            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
            {
                throw;
            }
            catch (WebSocketException exception)
            {
                Report(NodeState.Offline, "连接中断，等待重连");
                Console.Error.WriteLine($"[WARN] WebSocket disconnected: {exception.Message}");
            }
            catch (IOException exception)
            {
                Report(NodeState.Offline, "网络不可用，等待重连");
                Console.Error.WriteLine($"[WARN] Connection I/O failed: {exception.Message}");
            }
            catch (JsonException exception)
            {
                Report(NodeState.Offline, "服务端消息无效，等待重连");
                Console.Error.WriteLine($"[WARN] Invalid server message: {exception.Message}");
            }

            var baseDelay = Backoff[Math.Min(attempt, Backoff.Length - 1)];
            attempt++;
            var jitter = TimeSpan.FromMilliseconds(Random.Shared.Next(0, 750));
            var delay = baseDelay + jitter;
            Console.WriteLine($"[INFO] Reconnecting in {delay.TotalSeconds:F1}s.");
            await Task.Delay(delay, cancellationToken);
        }
    }

    private async Task ConnectOnceAsync(CancellationToken cancellationToken)
    {
        using var socket = new ClientWebSocket();
        socket.Options.SetRequestHeader("Authorization", $"Bearer {config.NodeToken}");
        socket.Options.SetRequestHeader("X-Node-Id", config.NodeId.ToString());
        socket.Options.KeepAliveInterval = TimeSpan.FromSeconds(config.HeartbeatSeconds);
        socket.Options.CollectHttpResponseDetails = true;

        var endpoint = NodeUri.WebSocketEndpoint(config.ServerUrl);
        Console.WriteLine($"[INFO] Connecting to {endpoint.GetLeftPart(UriPartial.Path)}.");
        try
        {
            await socket.ConnectAsync(endpoint, cancellationToken);
        }
        catch (WebSocketException) when (
            socket.HttpStatusCode is HttpStatusCode.Unauthorized or HttpStatusCode.Forbidden)
        {
            Report(NodeState.AuthenticationRequired, "身份失效，需要重新注册");
            throw new NodeConfigurationException(
                "Node credentials were rejected. Ask an administrator to re-register this node.");
        }
        Console.WriteLine("[OK] Node connected.");
        Report(NodeState.Online, "已连接");

        await SendAsync(socket, "hello", RuntimePayload(), cancellationToken);
        using var heartbeatStop = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        var heartbeat = HeartbeatLoopAsync(socket, heartbeatStop.Token);
        try
        {
            await ReceiveLoopAsync(socket, cancellationToken);
        }
        finally
        {
            heartbeatStop.Cancel();
            try
            {
                await heartbeat;
            }
            catch (OperationCanceledException) when (heartbeatStop.IsCancellationRequested)
            {
            }
        }
    }

    private async Task HeartbeatLoopAsync(ClientWebSocket socket, CancellationToken cancellationToken)
    {
        using var timer = new PeriodicTimer(TimeSpan.FromSeconds(config.HeartbeatSeconds));
        while (await timer.WaitForNextTickAsync(cancellationToken))
        {
            await SendAsync(socket, "heartbeat", RuntimePayload(), cancellationToken);
        }
    }

    private static async Task ReceiveLoopAsync(
        ClientWebSocket socket, CancellationToken cancellationToken)
    {
        var buffer = new byte[16 * 1024];
        using var message = new MemoryStream();
        while (socket.State == WebSocketState.Open)
        {
            var result = await socket.ReceiveAsync(buffer, cancellationToken);
            if (result.MessageType == WebSocketMessageType.Close)
            {
                if (result.CloseStatus == WebSocketCloseStatus.PolicyViolation)
                {
                    throw new NodeConfigurationException(
                        $"Node credentials or install identity were rejected: {result.CloseStatusDescription}");
                }
                return;
            }
            if (result.MessageType != WebSocketMessageType.Text)
            {
                throw new JsonException("Only text WebSocket messages are supported.");
            }
            message.Write(buffer, 0, result.Count);
            if (!result.EndOfMessage)
            {
                if (message.Length > 1024 * 1024)
                {
                    throw new JsonException("Server message exceeded 1 MiB.");
                }
                continue;
            }
            using var document = JsonDocument.Parse(message.ToArray());
            if (document.RootElement.TryGetProperty("type", out var type))
            {
                Console.WriteLine($"[INFO] Server message: {type.GetString()}.");
            }
            message.SetLength(0);
        }
    }

    private void Report(NodeState state, string message) =>
        statusChanged?.Invoke(NodeStatus.Create(state, message));

    private async Task SendAsync(
        ClientWebSocket socket, string type, object payload, CancellationToken cancellationToken)
    {
        var envelope = JsonSerializer.SerializeToUtf8Bytes(new
        {
            protocol_version = 1,
            message_id = Guid.NewGuid(),
            type,
            sent_at = DateTimeOffset.UtcNow,
            payload,
        });
        await socket.SendAsync(
            envelope, WebSocketMessageType.Text, endOfMessage: true, cancellationToken);
    }

    private object RuntimePayload()
    {
        var root = Path.GetPathRoot(Environment.SystemDirectory) ?? "C:\\";
        return new
        {
            install_id = config.InstallId,
            node_name = config.NodeName,
            node_version = Assembly.GetExecutingAssembly().GetName().Version?.ToString(3) ?? "0.1.0",
            os_version = RuntimeInformation.OSDescription,
            capabilities = config.Capabilities,
            available_slots = 1,
            disk_free_bytes = new DriveInfo(root).AvailableFreeSpace,
            desktop_session = Environment.UserInteractive,
        };
    }
}