MoShiro/backend/src/MoShiro.Backend.Application/Agents/AgentService.cs

using System.Net;
using System.Net.Sockets;
using System.Text.RegularExpressions;
using Fleet.Core;
using Fleet.Server;
using MoShiro.Protocol;

namespace MoShiro.Backend.Application;

public sealed record AgentView(AgentRecord Agent, bool Online, IReadOnlyList<ConnectedProbeInfo> Probes, AgentHeartbeatExtra? Heartbeat, int EndpointChanges, long? HeartbeatAtMs = null, HostSnapshot? Host = null)
{
    /// <summary>至少有一个 Probe 连着。</summary>
    public bool ProbeConnected => Probes.Count > 0;
}

/// <summary>Agent 的管理：列表、修改、审批、启停、接管、删除，以及 Probe 需要的 Agent 列表。</summary>
public sealed partial class AgentService
{
    private readonly IAgentStore _agents;
    private readonly AgentRuntime _runtime;
    private readonly IAgentGateway _gateway;
    private readonly TaskPlanner _planner;
    private readonly ISettingsStore _settings;

    private readonly NodeInfoService _nodes;
    private readonly LinkService _links;

    public AgentService(IAgentStore agents, AgentRuntime runtime, IAgentGateway gateway, TaskPlanner planner, ISettingsStore settings, NodeInfoService nodes, LinkService links)
    {
        _links = links;
        _nodes = nodes;
        _agents = agents;
        _runtime = runtime;
        _gateway = gateway;
        _planner = planner;
        _settings = settings;
    }

    public async ValueTask<IReadOnlyList<AgentView>> ListAsync(CancellationToken ct)
    {
        IReadOnlyList<AgentRecord> agents = await _agents.ListAsync(ct);
        return [.. agents.Select(View)];
    }

    public async ValueTask<AgentView?> GetAsync(NodeId id, CancellationToken ct) =>
        await _agents.GetAsync(id, ct) is { } agent ? View(agent) : null;

    public AgentView View(AgentRecord agent)
    {
        _runtime.TryGet(agent.Id, out AgentRuntimeState? state);
        return new AgentView(agent, _gateway.IsOnline(agent.Id), state?.Probes ?? [], state?.LastHeartbeat, state?.EndpointChanges ?? 0, state?.LastHeartbeatAtMs, state?.Host);
    }

    private static bool ValidHost(string host, string ipVersion)
    {
        if (IPAddress.TryParse(host, out IPAddress? address))
        {
            return !address.IsIPv4MappedToIPv6 && address.AddressFamily == (ipVersion == IpVersions.V6 ? AddressFamily.InterNetworkV6 : AddressFamily.InterNetwork);
        }

        return host.Length <= 253 && Uri.CheckHostName(host) == UriHostNameType.Dns;
    }

    public async ValueTask<Result<AgentView>> PatchAsync(NodeId id, AgentPatch patch, long? expectedRowVersion, CancellationToken ct)
    {
        var errors = new Dictionary<string, string>();
        if (patch.Name is not null && !NamePattern().IsMatch(patch.Name))
        {
            errors["name"] = "must be 1-64 characters: letters (any language), digits, '.', '_' or '-'";
        }

        if (patch.PublicHost is not null && patch.PublicHostAuto != true && !ValidHost(patch.PublicHost, IpVersions.V4))
        {
            errors["public_host"] = "must be a host name or IP address";
        }

        if (patch.Ipv6Host is not null && patch.Ipv6HostAuto != true && patch.Ipv6Host.Length > 0 && !ValidHost(patch.Ipv6Host, IpVersions.V6))
        {
            errors["ipv6_host"] = "must be an IPv6 address or host name; empty disables IPv6 collection";
        }

        if (patch.Labels is not null && (patch.Labels.Count > 16 || patch.Labels.Any(l => l.Length is 0 or > 32)))
        {
            errors["labels"] = "at most 16 labels of 1-32 characters";
        }

        if (errors.Count > 0)
        {
            return ValidationError.Fail<AgentView>(errors);
        }

        Result<AgentRecord> result = await _agents.PatchAsync(id, patch, expectedRowVersion, ct);
        if (!result.TryGet(out AgentRecord? agent, out FleetError? error))
        {
            return Result<AgentView>.Fail(error);
        }

        if (patch.PublicHostAuto == true || patch.Ipv6HostAuto == true)
        {
            // 撤销手填：马上按已经上报的地址重新识别
            (string? host, string? v6) = AgentAddressPolicy.Choose(agent, agent.ReportedAddresses ?? [], null);
            await _agents.ApplyAutoAddressesAsync(id, host, v6, ct);
            agent = await _agents.GetAsync(id, ct) ?? agent;
        }

        if (patch.Name is not null || patch.PublicHost is not null || patch.Ipv6Host is not null || patch.PublicHostAuto == true || patch.Ipv6HostAuto == true)
        {
            await AfterAgentSetChangedAsync(ct);
        }

        return Result<AgentView>.Ok(View(agent));
    }

    public ValueTask<Result<AgentView>> ApproveAsync(NodeId id, CancellationToken ct) => ChangeStatusAsync(id, NodeStatus.Active, ct);

    public ValueTask<Result<AgentView>> EnableAsync(NodeId id, CancellationToken ct) => ChangeStatusAsync(id, NodeStatus.Active, ct);

    public async ValueTask<Result<AgentView>> DisableAsync(NodeId id, CancellationToken ct)
    {
        Result<AgentView> result = await ChangeStatusAsync(id, NodeStatus.Disabled, ct);
        if (result.IsSuccess)
        {
            await _gateway.DisconnectAsync(id, "disabled", ct);
        }

        return result;
    }

    public async ValueTask<Result<AgentView>> TakeoverAsync(NodeId pendingId, CancellationToken ct)
    {
        Result<AgentRecord> result = await _agents.TakeoverAsync(pendingId, ct);
        if (!result.TryGet(out AgentRecord? agent, out FleetError? error))
        {
            return Result<AgentView>.Fail(error);
        }

        // 公钥换了，旧连接使用的是旧密钥
        await _gateway.DisconnectAsync(agent.Id, "key replaced", ct);
        await AfterAgentSetChangedAsync(ct);
        return Result<AgentView>.Ok(View(agent));
    }

    public async ValueTask DeleteAsync(NodeId id, CancellationToken ct)
    {
        await _gateway.DisconnectAsync(id, "deleted", ct);
        await _agents.DeleteAsync(id, ct);
        await _nodes.ForgetAsync(id.Value, ct);
        await _links.ForgetNodeAsync(id.Value, ct);
        _runtime.Remove(id);
        await AfterAgentSetChangedAsync(ct);
    }

    /// <summary>这个 Probe 要连接的 Agent 列表（已激活、有公开地址，并且线路是开着的）。版本号在列表变化时递增。</summary>
    public async ValueTask<ProbeAgentList> ProbeAgentsAsync(string probeId, CancellationToken ct)
    {
        LinkSnapshot links = await _links.SnapshotAsync(ct);
        IReadOnlyList<AgentRecord> agents = await _agents.ListAsync(ct);
        long version = long.TryParse(await _settings.GetAsync(TaskPlanner.ProbeAgentsVersionKey, ct), out long v) ? v : 0;
        var result = new ProbeAgentList { Version = version };
        foreach (AgentRecord agent in agents.Where(a => a.Status == NodeStatus.Active && a.PublicHost.Length > 0 && links.IsLinked(probeId, a.Id.Value)))
        {
            result.Agents.Add(new AgentEndpoint
            {
                AgentId = agent.Id,
                Name = agent.Name,
                Host = agent.PublicHost,
                Ipv6Host = agent.Ipv6Supported && agent.Ipv6Host.Length > 0 ? agent.Ipv6Host : null,
                TcpPort = agent.ProbeTcpPort,
                UdpPorts = [.. agent.ProbeUdpPorts],
                Spki = Convert.ToBase64String(agent.TlsSpki),
            });
        }

        return result;
    }

    /// <summary>Agent 集合变了：递增列表版本，重算各 Agent 的期望状态，并把新的 Agent 列表推给 Probe（都是期望状态，一起通知）。</summary>
    public async ValueTask AfterAgentSetChangedAsync(CancellationToken ct)
    {
        await _planner.BumpProbeAgentsVersionAsync(ct);
        _gateway.NotifyAllTaskSetsChanged();
    }

    private async ValueTask<Result<AgentView>> ChangeStatusAsync(NodeId id, NodeStatus status, CancellationToken ct)
    {
        Result<AgentRecord> result = await _agents.SetStatusAsync(id, status, ct);
        if (!result.TryGet(out AgentRecord? agent, out FleetError? error))
        {
            return Result<AgentView>.Fail(error);
        }

        await AfterAgentSetChangedAsync(ct);
        _gateway.NotifyTaskSetChanged(id);
        return Result<AgentView>.Ok(View(agent));
    }

    [GeneratedRegex(@"^[\p{L}\p{M}\p{N}._-]{1,64}$")]
    private static partial Regex NamePattern();
}