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

using System.Text.RegularExpressions;
using Fleet.Core;
using Fleet.Identity;
using Fleet.Server;
using MoShiro.Protocol;

namespace MoShiro.Backend.Application;

/// <summary>本地硬上限：写在 Probe 本机的配置里，面板无法修改（在路由器上建议调低）。</summary>
public sealed record ProbeLimits(double MaxTestMbps, long MaxTotalPps);

/// <summary><c>probe.json</c>：只在生成（或轮换密钥）时返回一次，含私钥；<c>Server</c> 是 Probe 要连的面板地址。</summary>
public sealed record ProbeIdentityFile(ProbeId ProbeId, string PrivateKey, string Server, ProbeLimits Limits);

/// <summary>某个 Probe 现在连着的一个 Agent（Agent 上报）。</summary>
public sealed record ProbeConnection(NodeId AgentId, string AgentName, bool AgentOnline, string? Version, string? Remote, long SinceMs, string IpVersion = IpVersions.V4);

public sealed record ProbeView(ProbeRecord Probe, bool Online, IReadOnlyList<ProbeConnection> Connections)
{
    public bool Revoked => Probe.RevokedAtMs is not null;
}

/// <summary>
/// Probe 身份（Backend 生成，私钥只返回一次）：可以有多个，每个都能改名、轮换密钥、吊销、删除。
/// 所有未吊销的 Probe 的公钥都会写进每个 Agent 的期望状态（<c>peers</c>）；吊销、删除、轮换密钥之后，Agent 立刻断开对应会话，Probe 到面板的连接也会被拒绝。
/// 一条线路是一对（Probe，Agent）。
/// </summary>
public sealed partial class ProbeService
{
    private readonly IProbeStore _probes;
    private readonly IAgentStore _agents;
    private readonly AgentRuntime _runtime;
    private readonly IAgentGateway _gateway;
    private readonly TimeProvider _time;
    private readonly NodeInfoService _nodes;
    private readonly LinkService _links;

    public ProbeService(IProbeStore probes, IAgentStore agents, AgentRuntime runtime, IAgentGateway gateway, TimeProvider time, NodeInfoService nodes, LinkService links)
    {
        _links = links;
        _nodes = nodes;
        _probes = probes;
        _agents = agents;
        _runtime = runtime;
        _gateway = gateway;
        _time = time;
    }

    public async ValueTask<Result<ProbeIdentityFile>> CreateAsync(string? name, string server, CancellationToken ct)
    {
        string finalName = string.IsNullOrWhiteSpace(name) ? "probe" : name.Trim();
        if (!NamePattern().IsMatch(finalName))
        {
            return ValidationError.Fail<ProbeIdentityFile>("name", "must be 1-64 characters: letters (any language), digits, '.', '_' or '-'");
        }

        if ((await _probes.ListAsync(ct)).Any(p => string.Equals(p.Name, finalName, StringComparison.OrdinalIgnoreCase)))
        {
            return Result<ProbeIdentityFile>.Fail(MoShiroErrors.NameConflict, $"a probe named {finalName} already exists");
        }

        using NodeIdentity identity = NodeIdentity.Generate();
        var probe = new ProbeRecord(Ulid.NewPrefixed(ProbeId.Prefix, _time), finalName, identity.PublicKeySpki, _time.GetUtcNow().ToUnixTimeMilliseconds(), null);
        await _probes.SaveAsync(probe, ct);
        _gateway.NotifyAllTaskSetsChanged();
        return Result<ProbeIdentityFile>.Ok(BuildFile(probe, identity, server));
    }

    public async ValueTask<IReadOnlyList<ProbeView>> ListAsync(CancellationToken ct)
    {
        IReadOnlyList<ProbeRecord> probes = await _probes.ListAsync(ct);
        IReadOnlyList<AgentRecord> agents = await _agents.ListAsync(ct);
        return [.. probes.Select(p => View(p, agents))];
    }

    public async ValueTask<ProbeView?> GetAsync(string id, CancellationToken ct) =>
        await _probes.GetAsync(id, ct) is { } probe ? View(probe, await _agents.ListAsync(ct)) : null;

    public async ValueTask<Result<ProbeView>> RenameAsync(string id, string name, CancellationToken ct)
    {
        if (await _probes.GetAsync(id, ct) is not { } probe)
        {
            return Result<ProbeView>.Fail(MoShiroErrors.ProbeNotFound, "probe not found");
        }

        string trimmed = name.Trim();
        if (!NamePattern().IsMatch(trimmed))
        {
            return ValidationError.Fail<ProbeView>("name", "must be 1-64 characters: letters (any language), digits, '.', '_' or '-'");
        }

        if ((await _probes.ListAsync(ct)).Any(p => p.Id != id && string.Equals(p.Name, trimmed, StringComparison.OrdinalIgnoreCase)))
        {
            return Result<ProbeView>.Fail(MoShiroErrors.NameConflict, $"a probe named {trimmed} already exists");
        }

        await _probes.SaveAsync(probe with { Name = trimmed }, ct);
        return Result<ProbeView>.Ok((await GetAsync(id, ct))!);
    }

    /// <summary>换一把新密钥（id 不变）：旧私钥立即失效，所有 Agent 上这个 Probe 的会话被断开。返回新的 probe.json（只此一次）。</summary>
    public async ValueTask<Result<ProbeIdentityFile>> RotateAsync(string id, string server, CancellationToken ct)
    {
        if (await _probes.GetAsync(id, ct) is not { } probe)
        {
            return Result<ProbeIdentityFile>.Fail(MoShiroErrors.ProbeNotFound, "probe not found");
        }

        if (probe.RevokedAtMs is not null)
        {
            return Result<ProbeIdentityFile>.Fail(MoShiroErrors.ProbeRevoked, "a revoked probe cannot get a new key; create a new probe instead");
        }

        using NodeIdentity identity = NodeIdentity.Generate();
        ProbeRecord updated = probe with { PublicKey = identity.PublicKeySpki };
        await _probes.SaveAsync(updated, ct);
        _gateway.NotifyAllTaskSetsChanged();
        await DisconnectEverywhereAsync(id, ct);
        return Result<ProbeIdentityFile>.Ok(BuildFile(updated, identity, server));
    }

    /// <summary>吊销：保留记录和历史数据，公钥从所有 Agent 的名单里移除，Agent 立即断开它的会话。</summary>
    public async ValueTask<Result<ProbeView>> RevokeAsync(string id, CancellationToken ct)
    {
        if (await _probes.GetAsync(id, ct) is null)
        {
            return Result<ProbeView>.Fail(MoShiroErrors.ProbeNotFound, "probe not found");
        }

        await _probes.RevokeAsync(id, _time.GetUtcNow().ToUnixTimeMilliseconds(), ct);
        // 吊销会阻止未来认证；也断开当前直连 Backend 的长会话，不等待 Probe 心跳超时。
        await _gateway.DisconnectAsync(new NodeId(id), "probe revoked", ct);
        _gateway.NotifyAllTaskSetsChanged();
        return Result<ProbeView>.Ok((await GetAsync(id, ct))!);
    }

    /// <summary>删除：Probe 身份和它的线路从面板消失；已经入库的数据按保留策略自然过期。</summary>
    public async ValueTask<bool> DeleteAsync(string id, CancellationToken ct)
    {
        bool deleted = await _probes.DeleteAsync(id, ct);
        if (deleted)
        {
            await _nodes.ForgetAsync(id, ct);
            await _links.ForgetNodeAsync(id, ct);
            _gateway.NotifyAllTaskSetsChanged();
        }

        return deleted;
    }

    /// <summary>让某个 Agent 断开这个 Probe（Probe 会按退避重连，除非它已经被吊销）。</summary>
    public async ValueTask<Result<bool>> DisconnectAsync(string probeId, NodeId agentId, CancellationToken ct)
    {
        if (await _probes.GetAsync(probeId, ct) is null)
        {
            return Result<bool>.Fail(MoShiroErrors.ProbeNotFound, "probe not found");
        }

        if (await _agents.GetAsync(agentId, ct) is null)
        {
            return Result<bool>.Fail(MoShiroErrors.AgentNotFound, "agent not found");
        }

        return await _gateway.DisconnectProbeAsync(agentId, new ProbeId(probeId), ct);
    }

    private async ValueTask DisconnectEverywhereAsync(string probeId, CancellationToken ct)
    {
        foreach ((NodeId agent, _) in _runtime.ConnectionsOf(probeId))
        {
            await _gateway.DisconnectProbeAsync(agent, new ProbeId(probeId), ct);
        }
    }

    private ProbeView View(ProbeRecord probe, IReadOnlyList<AgentRecord> agents)
    {
        var connections = new List<ProbeConnection>();
        foreach ((NodeId agentId, ConnectedProbeInfo info) in _runtime.ConnectionsOf(probe.Id))
        {
            AgentRecord? agent = agents.FirstOrDefault(a => a.Id == agentId);
            if (agent is not null)
            {
                connections.Add(new ProbeConnection(agentId, agent.Name, _gateway.IsOnline(agentId), info.Version, info.Remote, info.SinceMs, info.IpVersion));
            }
        }

        return new ProbeView(probe, _gateway.IsOnline(new NodeId(probe.Id)), [.. connections.OrderBy(c => c.AgentName, StringComparer.Ordinal)]);
    }

    private static ProbeIdentityFile BuildFile(ProbeRecord probe, NodeIdentity identity, string server) =>
        new(new ProbeId(probe.Id), Convert.ToBase64String(identity.ExportPkcs8()), server, new ProbeLimits(MaxTestMbps: 100, MaxTotalPps: 20_000));

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