MoShiro/backend/tests/MoShiro.Backend.Tests/AgentHarness.cs

using System.Text.Json.Nodes;
using Fleet.Core;
using Fleet.Enrollment.Client;
using Fleet.Hosting;
using Fleet.Identity;
using MoShiro.Agent;
using MoShiro.Protocol;
using MoShiro.TestSupport;

namespace MoShiro.Backend.Tests;

/// <summary>
/// 进程内的真实 Agent：用 <see cref="EnrollmentClient"/> 注册（和 <c>fleet-enroll</c> 同一份代码）、写出身份目录、
/// 在随机端口上启动 <see cref="AgentApp"/> 并连接真实的后端。
/// </summary>
public sealed class AgentHarness : IAsyncDisposable
{
    private readonly CancellationTokenSource _stop = new();
    private readonly string _directory;
    private readonly LoadedNodeIdentity _identity;
    private readonly Task _run;

    private AgentHarness(BackendFixture backend, string directory, LoadedNodeIdentity identity, AgentApp app)
    {
        Backend = backend;
        _directory = directory;
        _identity = identity;
        App = app;
        _run = Task.Run(() => app.RunAsync(_stop.Token), CancellationToken.None);
    }

    public BackendFixture Backend { get; }

    public AgentApp App { get; }

    public NodeId NodeId => _identity.NodeId;

    public string Id => _identity.NodeId.Value;

    public string DataDirectory => _directory;

    public byte[] TlsSpki => Convert.FromBase64String(_identity.File.TlsSpki);

    public static async Task<AgentHarness> StartAsync(BackendFixture backend, Action<AgentConfig>? configure = null, string? name = null)
    {
        string directory = Path.Combine(Path.GetTempPath(), "moshiro-agent-" + Guid.NewGuid().ToString("N"));
        string token = await backend.CreateJoinTokenAsync();
        using var http = new HttpClient();
        Result<EnrollmentResult> enrolled = await EnrollmentClient.EnrollAsync(http, backend.BaseAddress, token, name ?? backend.UniqueName("agent"), directory, AgentApp.Version, null, TimeProvider.System, CancellationToken.None);
        Assert.True(enrolled.IsSuccess, enrolled.IsSuccess ? "" : enrolled.Error.ToString());

        var config = new AgentConfig
        {
            Fleet = { Server = backend.BaseAddress.ToString(), IdentityDir = directory },
            Listen = { Bind = "127.0.0.1", ControlPort = 0, UdpPorts = [0], TcpPorts = [0] },
        };
        configure?.Invoke(config);
        LoadedNodeIdentity identity = LoadedNodeIdentity.Load(directory);
        var app = new AgentApp(config, identity, directory, TimeProvider.System, new ConsoleNodeLog(TimeProvider.System, Console.Out));
        var harness = new AgentHarness(backend, directory, identity, app);
        await app.Ready.WaitAsync(TimeSpan.FromSeconds(10));
        await BackendFixture.EventuallyAsync(() => Task.FromResult(app.Link.IsConnected), because: "the agent should connect to the backend");
        return harness;
    }

    /// <summary>在后端生成 Probe 身份（私钥只返回一次），并构造连到这个 Agent 的假 Probe。</summary>
    public async Task<ProbeSimulator> CreateProbeAsync()
    {
        JsonNode identity = await Backend.CreateProbeAsync();
        NodeIdentity key = NodeIdentity.FromPkcs8(Convert.FromBase64String(identity["private_key"]!.GetValue<string>()));
        string publicKey = Convert.ToBase64String(key.PublicKeySpki);

        // 等 Agent 的期望状态里有了这把公钥（否则握手会被拒绝）
        await BackendFixture.EventuallyAsync(
            () => Task.FromResult(App.IsPeerKnown(identity["probe_id"]!.GetValue<string>(), publicKey)),
            because: "the agent should learn the probe's key");
        return new ProbeSimulator(identity["probe_id"]!.GetValue<string>(), key);
    }

    public async Task<ProbeSimulator> ConnectProbeAsync()
    {
        ProbeSimulator probe = await CreateProbeAsync();
        await probe.ConnectAsync("127.0.0.1", App.ControlPort, TlsSpki, App.UdpPorts[0], NodeId);
        return probe;
    }

    /// <summary>为这个 Agent 创建覆盖任务，让线路任务使用它实际监听的端口（默认模板里是 7443/7444，不在本地允许的端口里）。</summary>
    public async Task OverrideLineTasksAsync(int? rateHz = null)
    {
        JsonNode tasks = await Backend.GetJsonAsync(Backend.Admin, "/api/admin/tasks");
        foreach (JsonNode? task in tasks["items"]!.AsArray())
        {
            if (task!["scope"]!.GetValue<string>() != "line" || task["agent_id"] is not null || task["kind"]!.GetValue<string>() is not ("udp_echo" or "tcp_echo"))
            {
                continue;
            }

            string kind = task["kind"]!.GetValue<string>();
            var parameters = new JsonObject { ["port"] = kind == "udp_echo" ? App.UdpPorts[0] : App.TcpPorts[0] };
            if (rateHz is { } hz)
            {
                parameters["rate_hz"] = hz;
            }

            using HttpResponseMessage created = await Backend.Admin.PostAsync("/api/admin/tasks", BackendFixture.Json(new JsonObject
            {
                ["scope"] = "line",
                ["kind"] = kind,
                ["agent_id"] = Id,
                ["base_task_id"] = task["id"]!.GetValue<string>(),
                ["params"] = parameters,
            }));
            Assert.True(created.StatusCode == System.Net.HttpStatusCode.Created, await created.Content.ReadAsStringAsync());
        }
    }

    public async ValueTask DisposeAsync()
    {
        await _stop.CancelAsync();
        try
        {
            await _run.WaitAsync(TimeSpan.FromSeconds(10));
        }
        catch (Exception ex) when (ex is OperationCanceledException or TimeoutException)
        {
            // 停止中
        }

        App.Dispose();
        _identity.Dispose();
        _stop.Dispose();
        try
        {
            Directory.Delete(_directory, recursive: true);
        }
        catch (IOException)
        {
            // 临时目录清理失败不影响测试结果
        }
    }
}