MoShiro/backend/shared/MoShiro.Measurement/Execution/BudgetGuard.cs

using System.Net;
using System.Text.Json;
using MoShiro.Protocol;

namespace MoShiro.Measurement;

/// <summary>节点的本地硬上限（写在本机配置里，Backend 无法修改）。</summary>
public sealed class NodeLimits
{
    public double MaxTotalMbps { get; init; } = 1200;

    public long MaxTotalPps { get; init; } = 20_000;

    /// <summary>目标任务允许访问的网段；默认为空，即不开放任何目标任务。</summary>
    public IReadOnlyList<string> AllowedTargetCidrs { get; init; } = [];

    public int MaxConcurrentTests { get; init; } = 1;

    /// <summary>任务参数里的 <c>port</c> 只能取这些端口（本机配置的 listen 列表）。</summary>
    public IReadOnlyList<int> AllowedQuicPorts { get; init; } = [7443];

    public IReadOnlyList<int> AllowedTcpPorts { get; init; } = [7444];
}

/// <summary>被拒绝的原因，用于 <c>rejected[]</c>：<c>unsupported</c>、<c>invalid_params</c>、<c>budget_exceeded</c>、<c>policy_denied</c>。</summary>
public static class RejectReasons
{
    public const string Unsupported = "unsupported";
    public const string InvalidParams = "invalid_params";
    public const string BudgetExceeded = "budget_exceeded";
    public const string PolicyDenied = "policy_denied";
}

public sealed record TaskDecision(bool Accepted, string? Reason, ParamValues? Resolved, KindSpec? Kind)
{
    public static TaskDecision Reject(string reason) => new(false, reason, null, null);
}

/// <summary>
/// 本地硬上限检查：Backend 下发的任何参数都要先经过 <see cref="ParamResolver"/> 校验，再经过这里的本地上限检查，两道都不能省略。
/// 持续任务按顺序累计码率和包速率，超出总预算的任务被拒绝（并给出原因），其余照常运行。
/// </summary>
public sealed class BudgetGuard
{
    private readonly NodeLimits _limits;
    private readonly IReadOnlyCollection<string> _supportedKinds;
    private readonly List<(IPAddress Network, int Prefix)> _cidrs = [];

    public BudgetGuard(NodeLimits limits, IReadOnlyCollection<string> supportedKinds)
    {
        _limits = limits;
        _supportedKinds = supportedKinds;
        foreach (string cidr in limits.AllowedTargetCidrs)
        {
            if (TryParseCidr(cidr, out IPAddress? network, out int prefix))
            {
                _cidrs.Add((network, prefix));
            }
        }
    }

    /// <summary>检查一个任务；<paramref name="usedMbps"/> 和 <paramref name="usedPps"/> 是已经接受的任务占用的预算，接受之后会累加。</summary>
    public TaskDecision Check(TaskSpec task, ref double usedMbps, ref long usedPps)
    {
        KindSpec? kind = Kinds.Find(task.Kind);
        if (kind is null || !_supportedKinds.Contains(task.Kind))
        {
            return TaskDecision.Reject(RejectReasons.Unsupported);
        }

        ParamResolution resolved = ParamResolver.Resolve(kind.Params, task.Params);
        if (!resolved.IsValid)
        {
            return TaskDecision.Reject(RejectReasons.InvalidParams);
        }

        if (task.Scope == "target")
        {
            if (!TargetAllowed(task.Target))
            {
                return TaskDecision.Reject(RejectReasons.PolicyDenied);
            }
        }
        else if (!PortAllowed(kind, resolved.Values!))
        {
            return TaskDecision.Reject(RejectReasons.PolicyDenied);
        }

        double mbps = kind.BudgetBitsPerSecond(resolved.Values!) / 1_000_000.0;
        long pps = EstimatePps(kind.Id, resolved.Values!);
        if (usedMbps + mbps > _limits.MaxTotalMbps || usedPps + pps > _limits.MaxTotalPps)
        {
            return TaskDecision.Reject(RejectReasons.BudgetExceeded);
        }

        usedMbps += mbps;
        usedPps += pps;
        return new TaskDecision(true, null, resolved.Values, kind);
    }

    /// <summary>按需测试的一个阶段是否在本地上限之内（单独运行，不和持续任务累计）。</summary>
    public TaskDecision CheckPhase(string kindId, IReadOnlyDictionary<string, JsonElement> parameters)
    {
        double mbps = 0;
        long pps = 0;
        return Check(new TaskSpec { Id = new TaskId("tsk_phase"), Scope = "line", Kind = kindId, Params = new Dictionary<string, JsonElement>(parameters) }, ref mbps, ref pps);
    }

    public bool TargetAllowed(string? target)
    {
        if (target is null)
        {
            return false;
        }

        string host = target;
        if (IPEndPoint.TryParse(target, out IPEndPoint? endpoint))
        {
            host = endpoint.Address.ToString();
        }

        if (!IPAddress.TryParse(host, out IPAddress? address))
        {
            return false;
        }

        foreach ((IPAddress network, int prefix) in _cidrs)
        {
            if (InCidr(address, network, prefix))
            {
                return true;
            }
        }

        return false;
    }

    private bool PortAllowed(KindSpec kind, ParamValues values)
    {
        if (!kind.Params.Any(p => p.Name == "port"))
        {
            return true;
        }

        int port = (int)values.GetInt("port");
        return kind.Blocks.Transport == "quic_datagram" ? _limits.AllowedQuicPorts.Contains(port) : _limits.AllowedTcpPorts.Contains(port);
    }

    private static long EstimatePps(string kindId, ParamValues values) => kindId switch
    {
        "udp_echo" or "tcp_echo" => values.GetInt("rate_hz") * 2,
        "icmp_echo" => values.GetInt("rate_hz") * 2,
        "udp_train" => (long)(values.GetNumber("mbps") * 1_000_000 / 8 / Math.Max(1, values.GetInt("size"))),
        "connect" => values.GetInt("per_minute") * 6 / 60,
        _ => 1,
    };

    private static bool TryParseCidr(string text, out IPAddress network, out int prefix)
    {
        network = IPAddress.None;
        prefix = 0;
        int slash = text.IndexOf('/', StringComparison.Ordinal);
        if (slash < 0 || !IPAddress.TryParse(text.AsSpan(0, slash), out IPAddress? parsed) || !int.TryParse(text.AsSpan(slash + 1), out prefix))
        {
            return false;
        }

        network = parsed;
        return prefix >= 0 && prefix <= parsed.GetAddressBytes().Length * 8;
    }

    private static bool InCidr(IPAddress address, IPAddress network, int prefix)
    {
        byte[] a = address.GetAddressBytes();
        byte[] n = network.GetAddressBytes();
        if (a.Length != n.Length)
        {
            return false;
        }

        int bits = prefix;
        for (int i = 0; i < a.Length && bits > 0; i++, bits -= 8)
        {
            int mask = bits >= 8 ? 0xFF : (0xFF << (8 - bits)) & 0xFF;
            if ((a[i] & mask) != (n[i] & mask))
            {
                return false;
            }
        }

        return true;
    }
}