zcbot/windows-node/Zcbot.WindowsNode/Protocol/NodeProtocolClient.cs

using System.Net;
using System.Net.WebSockets;
using System.Text.Json;

namespace Zcbot.WindowsNode;

internal sealed class NodeProtocolClient : IAsyncDisposable
{
    private const int MaximumMessageBytes = 1024 * 1024;
    private readonly WebSocket socket;
    private readonly SemaphoreSlim sendLock = new(1, 1);

    internal NodeProtocolClient(WebSocket socket)
    {
        this.socket = socket;
    }

    internal bool IsOpen => socket.State == WebSocketState.Open;

    internal static async Task<NodeProtocolClient> ConnectAsync(
        NodeConfig config,
        CancellationToken cancellationToken)
    {
        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);
            return new NodeProtocolClient(socket);
        }
        catch (WebSocketException exception) when (
            socket.HttpStatusCode is HttpStatusCode.Unauthorized or HttpStatusCode.Forbidden)
        {
            socket.Dispose();
            throw new NodeEndpointException(
                $"WebSocket handshake returned HTTP {(int?)socket.HttpStatusCode}. "
                + "Verify that the server includes the current Windows Node routes and that "
                + $"the reverse proxy forwards WebSocket Upgrade for {endpoint.AbsolutePath}. "
                + exception.Message);
        }
        catch
        {
            socket.Dispose();
            throw;
        }
    }

    internal async Task SendAsync(
        string type,
        object payload,
        CancellationToken cancellationToken)
    {
        var envelope = NodeProtocolCodec.Encode(type, payload);
        await sendLock.WaitAsync(cancellationToken);
        try
        {
            await socket.SendAsync(
                new ArraySegment<byte>(envelope),
                WebSocketMessageType.Text,
                endOfMessage: true,
                cancellationToken);
        }
        finally
        {
            sendLock.Release();
        }
    }

    internal async Task<NodeProtocolMessage?> ReceiveAsync(CancellationToken cancellationToken)
    {
        var buffer = new byte[16 * 1024];
        using var message = new MemoryStream();
        while (socket.State == WebSocketState.Open)
        {
            var result = await socket.ReceiveAsync(
                new ArraySegment<byte>(buffer), cancellationToken);
            if (result.MessageType == WebSocketMessageType.Close)
            {
                ThrowIfRejected(result);
                return null;
            }
            if (result.MessageType != WebSocketMessageType.Text)
            {
                throw new JsonException("Only text WebSocket messages are supported.");
            }

            message.Write(buffer, 0, result.Count);
            if (message.Length > MaximumMessageBytes)
            {
                throw new JsonException("Server message exceeded 1 MiB.");
            }
            if (result.EndOfMessage)
            {
                return NodeProtocolCodec.Decode(message.ToArray());
            }
        }
        return null;
    }

    public ValueTask DisposeAsync()
    {
        socket.Dispose();
        sendLock.Dispose();
        return ValueTask.CompletedTask;
    }

    private static void ThrowIfRejected(WebSocketReceiveResult result)
    {
        if ((int?)result.CloseStatus == 4003)
        {
            throw new NodeConfigurationException(
                "节点身份已被服务端拒绝。请在管理后台确认该 Node ID 未被禁用或删除；"
                + "若记录不存在或身份已撤销，请清除本机身份并使用新注册码重新注册。"
                + $" 服务端信息：{result.CloseStatusDescription}");
        }
        if (result.CloseStatus == WebSocketCloseStatus.PolicyViolation)
        {
            throw new NodeConfigurationException(
                $"节点上报被服务端拒绝：{result.CloseStatusDescription}");
        }
    }
}