MoShiro/backend/tests/MoShiro.TestSupport/AgentSimulator.cs

using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;
using Fleet.Core;
using Fleet.Identity;
using Fleet.Rpc;
using MoShiro.Protocol;

namespace MoShiro.TestSupport;

/// <summary>
/// 一个会说 Fleet 协议的假 Agent：用加入令牌注册 → 连接 → 签名认证 → hello → 接收期望状态 → 上传原始块并等待 ACK。
/// 只用于测试和离线工具，不是真正的节点。
/// </summary>
public sealed class AgentSimulator : IAsyncDisposable
{
    private readonly NodeIdentity _identity;
    private readonly ClientWebSocket _socket = new();
    private readonly ConcurrentDictionary<Guid, TaskCompletionSource<bool>> _acks = new();
    private readonly ConcurrentDictionary<long, TaskCompletionSource<RpcMessage>> _responses = new();
    private readonly CancellationTokenSource _stop = new();
    private readonly SemaphoreSlim _sendLock = new(1, 1);
    private long _nextId;
    private int _disposed;
    private Task? _loop;

    private AgentSimulator(NodeIdentity identity, NodeId nodeId, string status)
    {
        _identity = identity;
        ExportedKey = identity.ExportPkcs8();
        NodeId = nodeId;
        EnrollmentStatus = status;
    }

    public NodeId NodeId { get; }

    /// <summary>身份私钥（PKCS#8）：释放之后仍可用它以同一个身份重新连接。</summary>
    public byte[] ExportedKey { get; }

    public string EnrollmentStatus { get; }

    public string AuthStatus { get; private set; } = "";

    public ConcurrentQueue<TaskSet> TaskSets { get; } = new();

    public ConcurrentQueue<TestRunParams> TestRuns { get; } = new();

    /// <summary>面板发来的发布清单（<c>moshiro.update.offer</c>）和更新请求（<c>moshiro.update.apply</c>）。</summary>
    public ConcurrentQueue<UpdateOfferParams> UpdateOffers { get; } = new();

    private int _updateApplies;

    public int UpdateApplies => Volatile.Read(ref _updateApplies);

    public ConcurrentQueue<TestId> TestCancels { get; } = new();

    public IReadOnlyList<RejectedItem> RejectNextState { get; set; } = [];

    public bool Revoked { get; private set; }

    public string? CloseReason { get; private set; }

    /// <summary>收到 moshiro.test.run 时的应答：默认接受。</summary>
    public Func<TestRunParams, (bool Accepted, string? ErrorCode)> TestRunHandler { get; set; } = _ => (true, null);

    public byte[] PublicKey => _identity.PublicKeySpki;

    /// <summary>用这个节点的私钥为认证随机数签名（可以冒充别的节点 ID：用来测试服务端的校验）。</summary>
    public byte[] SignAuth(NodeId claimedNode, string nonce) => _identity.Sign(SignaturePurpose.Auth, claimedNode.Value, nonce);

    public static async Task<AgentSimulator> EnrollAsync(HttpClient http, string joinToken, string name, string[]? kinds = null, string publicHost = "203.0.113.10")
    {
        NodeIdentity identity = NodeIdentity.Generate();
        string pubkey = Convert.ToBase64String(identity.PublicKeySpki);
        byte[] tokenHash = SHA256.HashData(Encoding.UTF8.GetBytes(joinToken));
        string proof = Convert.ToBase64String(identity.Sign(SignaturePurpose.Enroll, Convert.ToHexStringLower(tokenHash), pubkey));
        var body = new JsonObject
        {
            ["join_token"] = joinToken,
            ["name"] = name,
            ["pubkey"] = pubkey,
            ["tls_spki"] = Convert.ToBase64String(RandomNumberGenerator.GetBytes(32)),
            ["proof"] = proof,
            ["node_version"] = "0.0.1-sim",
            ["meta"] = new JsonObject
            {
                ["probe_tcp_port"] = 7443,
                ["probe_udp_ports"] = new JsonArray(7443),
                ["public_host"] = publicHost,
                ["kinds"] = new JsonArray((kinds ?? ["udp_echo", "tcp_echo", "udp_train", "connect", "hold", "tcp_bulk", "icmp_echo"]).Select(k => (JsonNode?)JsonValue.Create(k)).ToArray()),
            },
        };

        using var content = new StringContent(body.ToJsonString(), Encoding.UTF8, "application/json");
        using HttpResponseMessage response = await http.PostAsync("/api/agent/v1/enroll", content);
        string text = await response.Content.ReadAsStringAsync();
        if (!response.IsSuccessStatusCode)
        {
            identity.Dispose();
            throw new InvalidOperationException($"enroll failed: {(int)response.StatusCode} {text}");
        }

        JsonObject result = JsonNode.Parse(text)!.AsObject();
        return new AgentSimulator(identity, new NodeId(result["node_id"]!.GetValue<string>()), result["status"]!.GetValue<string>());
    }

    /// <summary>用已有的身份恢复一个模拟器（例如后端重启之后同一个节点重新连接）。</summary>
    public static AgentSimulator Restore(byte[] pkcs8, NodeId nodeId) => new(NodeIdentity.FromPkcs8(pkcs8), nodeId, "active");

    /// <summary>连接并完成认证与 hello。认证失败时抛出异常。</summary>
    public async Task ConnectAsync(Uri baseAddress, bool sendHello = true, string[]? kinds = null)
    {
        var uri = new UriBuilder(baseAddress) { Scheme = baseAddress.Scheme == "https" ? "wss" : "ws", Path = "/api/agent/v1/rpc" }.Uri;
        await _socket.ConnectAsync(uri, _stop.Token);

        RpcMessage challenge = await ReceiveMessageAsync(_stop.Token) ?? throw new InvalidOperationException("no challenge");
        ChallengeParams nonce = challenge.Params!.Value.Deserialize(FleetRpcContext.Default.ChallengeParams)!;
        string signature = Convert.ToBase64String(_identity.Sign(SignaturePurpose.Auth, NodeId.Value, nonce.Nonce));
        await SendAsync(RpcMessage.Request(NextId(), FleetMethods.Auth, new AuthParams { NodeId = NodeId.Value, Signature = signature }, FleetRpcContext.Default.AuthParams));
        RpcMessage? authReply = await ReceiveMessageAsync(_stop.Token);
        if (authReply?.Error is { } error)
        {
            throw new InvalidOperationException($"auth failed: {error.Data?.Code} {error.Message}");
        }

        AuthStatus = authReply!.Result!.Value.Deserialize(FleetRpcContext.Default.AuthResult)!.Status;
        _loop = Task.Run(ReceiveLoopAsync);
        if (sendHello)
        {
            await HelloAsync(kinds);
        }
    }

    public async Task HelloAsync(string[]? kinds = null)
    {
        var meta = JsonSerializer.SerializeToElement(
            new AgentHelloMeta
            {
                TlsSpki = Convert.ToBase64String(RandomNumberGenerator.GetBytes(32)),
                ProbeTcpPort = 7443,
                ProbeUdpPorts = [7443],
                LocalLimits = new LocalLimits { MaxTotalMbps = 200, MaxTotalPps = 20_000, MaxConcurrentTests = 1 },
            },
            MoShiroJsonContext.Default.AgentHelloMeta);
        var hello = new HelloParams
        {
            NodeVersion = "0.0.1-sim",
            Capabilities = [MoShiroCapabilities.Base, MoShiroCapabilities.KindsPrefix + string.Join(',', kinds ?? ["udp_echo", "tcp_echo", "udp_train", "connect", "hold", "tcp_bulk", "icmp_echo"])],
            NowUtcMs = TimeProvider.System.GetUtcNow().ToUnixTimeMilliseconds(),
            Meta = meta,
        };
        await RequestAsync(FleetMethods.Hello, JsonSerializer.SerializeToElement(hello, FleetRpcContext.Default.HelloParams));
    }

    /// <summary>上传一个原始块并等待服务端 ACK（持久化之后才会确认）。</summary>
    /// <summary>设置之后，这个假 Agent 发出的每个块都归属于这个 Probe（夹具的默认 Probe）。</summary>
    public ProbeId? ProbeForChunks { get; set; }

    public async Task<bool> SendChunkAsync(RawChunk chunk, TimeSpan? timeout = null)
    {
        if (ProbeForChunks is { } probe)
        {
            chunk.Header.ProbeId = probe;
        }

        byte[] bytes = RawChunkCodec.Encode(chunk);
        Guid id = chunk.Header.ChunkId.Value;
        var ack = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
        _acks[id] = ack;
        await SendBinaryAsync(OutboxFrame.Build(MoShiroChannels.RawChunk, id, bytes));
        return await ack.Task.WaitAsync(timeout ?? TimeSpan.FromSeconds(10));
    }

    /// <summary>发送一个任意的文本帧（用来测试服务端对非法输入的处理）。</summary>
    public Task SendTextAsync(string text) => SendFrameAsync(Encoding.UTF8.GetBytes(text), WebSocketMessageType.Text);

    /// <summary>发送一个任意的二进制帧（用来测试服务端对非法输入的处理）。</summary>
    public Task SendBinaryAsync(byte[] frame) => SendFrameAsync(frame, WebSocketMessageType.Binary);

    public Task SendEventAsync(string kind, string? detail = null) =>
        SendAsync(RpcMessage.Notification(
            MoShiroMethods.Event,
            new AgentEventParams { Kind = kind, AtUtcMs = TimeProvider.System.GetUtcNow().ToUnixTimeMilliseconds(), Detail = detail },
            MoShiroJsonContext.Default.AgentEventParams));

    public Task SendHostAsync(HostSnapshot snapshot) =>
        SendAsync(RpcMessage.Notification(MoShiroMethods.Host, snapshot, MoShiroJsonContext.Default.HostSnapshot));

    public Task SendTestStatusAsync(TestId test, string status, int? phase = null, string? errorCode = null) =>
        SendAsync(RpcMessage.Notification(
            MoShiroMethods.TestStatus,
            new TestStatusParams { TestId = test, Status = status, Phase = phase, ErrorCode = errorCode },
            MoShiroJsonContext.Default.TestStatusParams));

    /// <summary>像节点的更新器报告的那样，随心跳上报更新状态。</summary>
    public Task SendUpdateStatusAsync(UpdateStatus status) =>
        SendAsync(RpcMessage.Notification(
            FleetMethods.Heartbeat,
            new HeartbeatParams { Extra = JsonSerializer.SerializeToElement(new AgentHeartbeatExtra { Update = status }, MoShiroJsonContext.Default.AgentHeartbeatExtra) },
            FleetRpcContext.Default.HeartbeatParams));

    public async Task<RpcMessage> RequestRawAsync(string method, JsonElement? parameters = null) => await RequestAsync(method, parameters);

    public async Task<TaskSet> WaitForTaskSetAsync(Func<TaskSet, bool> predicate, TimeSpan? timeout = null)
    {
        long started = TimeProvider.System.GetTimestamp();
        TimeSpan limit = timeout ?? TimeSpan.FromSeconds(10);
        while (TimeProvider.System.GetElapsedTime(started) < limit)
        {
            foreach (TaskSet set in TaskSets.ToArray().Reverse())
            {
                if (predicate(set))
                {
                    return set;
                }
            }

            await Task.Delay(25);
        }

        throw new TimeoutException("the expected task set never arrived");
    }

    public async Task WaitUntilAsync(Func<bool> condition, TimeSpan? timeout = null)
    {
        long started = TimeProvider.System.GetTimestamp();
        TimeSpan limit = timeout ?? TimeSpan.FromSeconds(10);
        while (TimeProvider.System.GetElapsedTime(started) < limit)
        {
            if (condition())
            {
                return;
            }

            await Task.Delay(25);
        }

        throw new TimeoutException("condition was not met in time");
    }

    public async ValueTask DisposeAsync()
    {
        if (Interlocked.Exchange(ref _disposed, 1) != 0)
        {
            return;
        }

        await _stop.CancelAsync();
        if (_socket.State == WebSocketState.Open)
        {
            try
            {
                await _socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "bye", CancellationToken.None);
            }
            catch (WebSocketException)
            {
                // 已经断开
            }
        }

        if (_loop is not null)
        {
            try
            {
                await _loop;
            }
            catch (OperationCanceledException)
            {
                // 正常结束
            }
        }

        _socket.Dispose();
        _identity.Dispose();
        _stop.Dispose();
        _sendLock.Dispose();
    }

    private long NextId() => Interlocked.Increment(ref _nextId);

    private async Task<RpcMessage> RequestAsync(string method, JsonElement? parameters)
    {
        long id = NextId();
        var waiter = new TaskCompletionSource<RpcMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
        _responses[id] = waiter;
        await SendAsync(new RpcMessage { Id = id, Method = method, Params = parameters });
        return await waiter.Task.WaitAsync(TimeSpan.FromSeconds(10));
    }

    private Task SendAsync(RpcMessage message) =>
        SendFrameAsync(JsonSerializer.SerializeToUtf8Bytes(message, FleetRpcContext.Default.RpcMessage), WebSocketMessageType.Text);

    private async Task SendFrameAsync(byte[] bytes, WebSocketMessageType type)
    {
        await _sendLock.WaitAsync(_stop.Token);
        try
        {
            await _socket.SendAsync(bytes, type, endOfMessage: true, _stop.Token);
        }
        finally
        {
            _sendLock.Release();
        }
    }

    private async Task<RpcMessage?> ReceiveMessageAsync(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)
            {
                CloseReason = result.CloseStatusDescription;
                if (_socket.State == WebSocketState.CloseReceived)
                {
                    // 按协议确认对端发起的关闭
                    try
                    {
                        await _socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "ack", CancellationToken.None);
                    }
                    catch (WebSocketException)
                    {
                        // 对端已经断开
                    }
                }

                return null;
            }

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

    private async Task ReceiveLoopAsync()
    {
        try
        {
            while (!_stop.IsCancellationRequested)
            {
                RpcMessage? message = await ReceiveMessageAsync(_stop.Token);
                if (message is null)
                {
                    return;
                }

                await HandleAsync(message);
            }
        }
        catch (OperationCanceledException)
        {
            // 正常结束
        }
        catch (WebSocketException)
        {
            // 连接断开
        }
    }

    private async Task HandleAsync(RpcMessage message)
    {
        if (message.IsResponse)
        {
            if (_responses.TryRemove(message.Id!.Value, out TaskCompletionSource<RpcMessage>? waiter))
            {
                waiter.TrySetResult(message);
            }

            return;
        }

        switch (message.Method)
        {
            case FleetMethods.StateSync:
                StateSyncParams sync = message.Params!.Value.Deserialize(FleetRpcContext.Default.StateSyncParams)!;
                if (sync.Type == TaskSet.StateType)
                {
                    TaskSets.Enqueue(sync.Doc.Deserialize(MoShiroJsonContext.Default.TaskSet)!);
                }

                var result = new StateSyncResult { AppliedVersion = sync.Version, Rejected = [.. RejectNextState] };
                await SendAsync(RpcMessage.Success(message.Id!.Value, result, FleetRpcContext.Default.StateSyncResult));
                break;
            case FleetMethods.OutboxAck:
                OutboxAckParams ack = message.Params!.Value.Deserialize(FleetRpcContext.Default.OutboxAckParams)!;
                foreach (string id in ack.Ids.Concat(ack.Duplicates))
                {
                    if (_acks.TryRemove(Guid.Parse(id), out TaskCompletionSource<bool>? waiter))
                    {
                        waiter.TrySetResult(ack.Duplicates.Contains(id) ? false : true);
                    }
                }

                break;
            case FleetMethods.Revoked:
                Revoked = true;
                break;
            case MoShiroMethods.UpdateOffer:
                UpdateOffers.Enqueue(message.Params!.Value.Deserialize(MoShiroJsonContext.Default.UpdateOfferParams)!);
                break;
            case MoShiroMethods.UpdateApply:
                Interlocked.Increment(ref _updateApplies);
                break;
            case MoShiroMethods.TestRun:
                TestRunParams run = message.Params!.Value.Deserialize(MoShiroJsonContext.Default.TestRunParams)!;
                TestRuns.Enqueue(run);
                (bool accepted, string? errorCode) = TestRunHandler(run);
                if (accepted)
                {
                    await SendAsync(RpcMessage.Success(message.Id!.Value, new TestRunResult { Accepted = true }, MoShiroJsonContext.Default.TestRunResult));
                }
                else
                {
                    await SendAsync(RpcMessage.Failure(message.Id, new FleetError(new ErrorCode(errorCode ?? "test_busy"), "refused")));
                }

                break;
            case MoShiroMethods.TestCancel:
                TestCancelParams cancel = message.Params!.Value.Deserialize(MoShiroJsonContext.Default.TestCancelParams)!;
                TestCancels.Enqueue(cancel.TestId);
                await SendAsync(RpcMessage.Success(message.Id!.Value, new TestRunResult { Accepted = true }, MoShiroJsonContext.Default.TestRunResult));
                break;
        }
    }
}