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[] TlsPublicKey => _identity.Certificate.PublicKey.ExportSubjectPublicKeyInfo();

    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());

        // 测试里的 Agent 只在回环上监听：手填 public_host（手动的优先于自动识别，Agent 上报本机真实地址也不会覆盖）
        using HttpResponseMessage patched = await backend.Admin.PatchAsync("/api/admin/agents/" + enrolled.Value.NodeId, BackendFixture.Json(new JsonObject { ["public_host"] = "127.0.0.1" }));
        Assert.True(patched.IsSuccessStatusCode, await patched.Content.ReadAsStringAsync());

        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()
    {
        using EnrolledProbe probe = await Backend.CreateProbeAsync();
        // 假 Probe 要在注册目录被释放之后继续用这把密钥：导出再导入一份
        NodeIdentity key = NodeIdentity.FromPkcs8(probe.Identity.PeerIdentity.ExportPkcs8());
        string publicKey = Convert.ToBase64String(key.PublicKeySpki);

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

    public async Task<ProbeSimulator> ConnectProbeAsync()
    {
        ProbeSimulator probe = await CreateProbeAsync();
        await probe.ConnectAsync("127.0.0.1", App.ControlPort, TlsPublicKey, 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");
        // 先处理 connect（关掉），再改 udp_echo / tcp_echo 的速率：否则中途的带宽预算会多出 connect 那一份
        foreach (JsonNode? task in tasks["items"]!.AsArray().OrderBy(t => t!["kind"]!.GetValue<string>() == "connect" ? 0 : 1))
        {
            if (task!["scope"]!.GetValue<string>() != "line" || task["agent_id"] is not null || task["kind"]!.GetValue<string>() is not ("udp_echo" or "tcp_echo" or "connect"))
            {
                continue;
            }

            string kind = task["kind"]!.GetValue<string>();
            var parameters = new JsonObject { ["port"] = kind == "udp_echo" ? App.UdpPorts[0] : App.TcpPorts[0] };
            // 默认的 connect 模板在这些用例里不需要（也占预算）：对这个 Agent 关掉它
            bool enabled = kind != "connect";
            if (rateHz is { } hz && kind != "connect")
            {
                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,
                ["enabled"] = enabled,
            }));
            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)
        {
            // 临时目录清理失败不影响测试结果
        }
    }
}