MoShiro/backend/tests/MoShiro.TestSupport/AgentWire.cs

using System.Net.WebSockets;
using System.Text.Json;
using Fleet.Core;
using Fleet.Rpc;

namespace MoShiro.TestSupport;

public sealed record AuthAttempt(bool Ok, string? Status, string? ErrorCode, string? CloseReason);

/// <summary>只做 Fleet 的挑战/认证握手：用来测试被拒绝的连接（冒名、被禁用的节点等），不保持连接。</summary>
public static class AgentWire
{
    public static async Task<AuthAttempt> AuthenticateAsync(Uri baseAddress, NodeId claimedNode, Func<string, byte[]> signNonce)
    {
        using var socket = new ClientWebSocket();
        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        var uri = new UriBuilder(baseAddress) { Scheme = baseAddress.Scheme == "https" ? "wss" : "ws", Path = "/api/agent/v1/rpc" }.Uri;
        await socket.ConnectAsync(uri, timeout.Token);

        RpcMessage? challenge = await ReceiveAsync(socket, timeout.Token);
        ChallengeParams nonce = challenge!.Params!.Value.Deserialize(FleetRpcContext.Default.ChallengeParams)!;
        var auth = new AuthParams { NodeId = claimedNode.Value, Signature = Convert.ToBase64String(signNonce(nonce.Nonce)) };
        byte[] request = JsonSerializer.SerializeToUtf8Bytes(RpcMessage.Request(1, FleetMethods.Auth, auth, FleetRpcContext.Default.AuthParams), FleetRpcContext.Default.RpcMessage);
        await socket.SendAsync(request, WebSocketMessageType.Text, endOfMessage: true, timeout.Token);

        string? closeReason = null;
        while (true)
        {
            RpcMessage? reply = await ReceiveAsync(socket, timeout.Token);
            if (reply is null)
            {
                closeReason = socket.CloseStatusDescription;
                return new AuthAttempt(false, null, null, closeReason);
            }

            if (reply.IsResponse)
            {
                return reply.Error is { } error
                    ? new AuthAttempt(false, null, error.Data?.Code, closeReason)
                    : new AuthAttempt(true, reply.Result!.Value.Deserialize(FleetRpcContext.Default.AuthResult)!.Status, null, null);
            }
        }
    }

    private static async Task<RpcMessage?> ReceiveAsync(ClientWebSocket socket, CancellationToken ct)
    {
        var buffer = new byte[16 * 1024];
        using var message = new MemoryStream();
        while (true)
        {
            WebSocketReceiveResult result = await socket.ReceiveAsync(buffer, ct);
            if (result.MessageType == WebSocketMessageType.Close)
            {
                return null;
            }

            message.Write(buffer, 0, result.Count);
            if (result.EndOfMessage)
            {
                return JsonSerializer.Deserialize(message.ToArray(), FleetRpcContext.Default.RpcMessage);
            }
        }
    }
}