MoShiro/backend/agent/src/MoShiro.Agent/DataPlane/Registries.cs

using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using Fleet.Core;
using Fleet.Link;
using Fleet.Rpc;
using MoShiro.Measurement;
using MoShiro.Protocol;

namespace MoShiro.Agent;

/// <summary>进程内计数（只有计数，没有主机信息）：用于心跳和面板上的节点状态。数据面的包计数由共享的 <see cref="DataPlaneCounters"/> 提供。</summary>
public sealed class AgentCounters
{
    private long _fragments;
    private long _fragmentsDiscarded;

    public DataPlaneCounters Data { get; } = new();

    public long PacketsIn => Data.PacketsIn;

    public long PacketsOut => Data.PacketsOut;

    public long DroppedInvalid => Data.DroppedInvalid;

    public long DroppedRateLimited => Data.DroppedRateLimited;

    public long Fragments => Interlocked.Read(ref _fragments);

    public long FragmentsDiscarded => Interlocked.Read(ref _fragmentsDiscarded);

    public void FragmentAccepted() => Interlocked.Increment(ref _fragments);

    public void FragmentDiscarded() => Interlocked.Increment(ref _fragmentsDiscarded);
}

public sealed record PeerContext(DialerSession Session, MonoClock Clock);

/// <summary>
/// 一个已经认证的拨号方会话：数据面的 <c>session_id</c> 和密钥（由点对点握手里的 ECDH 派生）、最近一次的 UDP 源地址、
/// 每个流的重放窗口。数据面的包不绑定五元组，NAT 映射变化之后自动跟随新的源地址。
/// </summary>
public sealed class DialerSession : DataSession
{
    private readonly ConcurrentDictionary<uint, ReplayWindow> _windows = new();
    private readonly ConcurrentDictionary<ulong, byte> _nonces = new();
    private SocketAddress? _datagramSource;

    public DialerSession(PeerSessionInfo info)
        : base(info.Credentials.SessionIdValue, info.Credentials.Key)
    {
        Info = info;
    }

    public PeerSessionInfo Info { get; }

    public string DialerId => Info.PeerId;

    public override string DiagnosticPeerId => DialerId;

    public string? CloseReason { get; set; }

    public string IpVersion => Info.Remote is IPEndPoint endpoint && endpoint.Address.AddressFamily == AddressFamily.InterNetworkV6 && !endpoint.Address.IsIPv4MappedToIPv6 ? IpVersions.V6 : IpVersions.V4;

    public string PathKey => DialerId + "~" + IpVersion;

    /// <summary>拨号方 的程序版本（握手里自报，只用于展示）。</summary>
    public string? Version => Info.Version;

    /// <summary>拨号方 的远端地址（经过 NAT 之后的公网 IP）。</summary>
    public string? RemoteAddress => (Info.Remote as IPEndPoint)?.Address.ToString();

    public long ConnectedAtUtcMs { get; init; }

    public RpcSession<PeerContext>? Rpc { get; set; }

    /// <summary>最近一次认证通过的 UDP 包的源地址（Agent 把单向数据发往这里）。</summary>
    public override IPEndPoint? DatagramRemoteEndPoint
    {
        get
        {
            SocketAddress? source = Volatile.Read(ref _datagramSource);
            return source is null ? null : (IPEndPoint)new IPEndPoint(source.Family == AddressFamily.InterNetworkV6 ? IPAddress.IPv6Any : IPAddress.Any, 0).Create(source);
        }
    }

    /// <summary>记录 UDP 源地址；NAT 映射变化后地址变了返回 true（上报 endpoint_changed）。第一次看到地址不算变化。</summary>
    public bool UpdateSource(SocketAddress source)
    {
        SocketAddress? previous = Volatile.Read(ref _datagramSource);
        if (previous is not null && previous.Equals(source))
        {
            return false;
        }

        // SocketAddress 的缓冲区会被接收循环复用，所以必须复制
        var copy = new SocketAddress(source.Family, source.Size);
        for (int i = 0; i < source.Size; i++)
        {
            copy[i] = source[i];
        }

        Volatile.Write(ref _datagramSource, copy);
        return previous is not null;
    }

    public bool AcceptSequence(uint streamId, uint seq)
    {
        ReplayWindow window = _windows.GetOrAdd(streamId, _ => new ReplayWindow());
        lock (window)
        {
            return window.TryAccept(seq);
        }
    }

    /// <summary>TCP 首帧的随机数只能用一次，防止被截获之后重放。</summary>
    public override bool UseNonce(ReadOnlySpan<byte> nonce)
    {
        ulong hash = 14695981039346656037UL; // FNV-1a：随机数本身已经是 128 位随机，这里只需要一个去重用的指纹
        foreach (byte b in nonce)
        {
            hash = (hash ^ b) * 1099511628211UL;
        }

        if (_nonces.Count > 4096)
        {
            _nonces.Clear();
        }

        return _nonces.TryAdd(hash, 0);
    }

    public void ForgetStream(uint streamId) => _windows.TryRemove(streamId, out _);
}

/// <summary>每个拨号方的 IPv4、IPv6 会话独立；重连只替换同地址族的旧会话。</summary>
public sealed class SessionRegistry : IDataSessions
{
    private readonly object _gate = new();
    private readonly Dictionary<string, DialerSession> _byDialer = new(StringComparer.Ordinal);
    private readonly Dictionary<long, DialerSession> _bySession = [];

    public IReadOnlyList<DialerSession> All
    {
        get
        {
            lock (_gate)
            {
                return [.. _byDialer.Values];
            }
        }
    }

    public DialerSession? Get(string dialerId, string ipVersion = IpVersions.V4)
    {
        lock (_gate)
        {
            return _byDialer.GetValueOrDefault(dialerId + "~" + ipVersion);
        }
    }

    /// <summary>登记新会话；返回同一个拨号方的旧会话（如果有）。</summary>
    public DialerSession? Replace(DialerSession session)
    {
        lock (_gate)
        {
            _byDialer.TryGetValue(session.PathKey, out DialerSession? old);
            if (old is not null)
            {
                _bySession.Remove(old.SessionId);
            }

            _byDialer[session.PathKey] = session;
            _bySession[session.SessionId] = session;
            return old;
        }
    }

    public void Remove(DialerSession session)
    {
        lock (_gate)
        {
            if (_byDialer.TryGetValue(session.PathKey, out DialerSession? current) && ReferenceEquals(current, session))
            {
                _byDialer.Remove(session.PathKey);
            }

            if (_bySession.TryGetValue(session.SessionId, out DialerSession? bySession) && ReferenceEquals(bySession, session))
            {
                _bySession.Remove(session.SessionId);
            }
        }
    }

    public DataSession? Find(long sessionId)
    {
        lock (_gate)
        {
            return _bySession.GetValueOrDefault(sessionId);
        }
    }
}

/// <summary>数据面上一个流（一个线路任务，或测试的一个阶段）的元数据：准入策略据此限速，合并据此找到积木的记录。</summary>
public sealed class StreamInfo
{
    private readonly object _bucketGate = new();
    private int _fragmentsDone;
    private double _tokens = -1;
    private long _lastRefillUs;

    public StreamInfo(uint id, string kind, ParamValues parameters)
    {
        Id = id;
        Kind = kind;
        Params = parameters;
        KindSpec spec = Kinds.Find(kind)!;
        DeclaredPps = kind switch
        {
            "udp_echo" or "tcp_echo" => parameters.GetInt("rate_hz"),
            "udp_train" => parameters.GetNumber("mbps") * 1_000_000 / 8 / Math.Max(1, parameters.GetInt("size")),
            _ => 20,
        };
        Port = spec.Params.Any(p => p.Name == "port") ? (int)parameters.GetInt("port") : 0;
    }

    public uint Id { get; }

    public string Kind { get; }

    public ParamValues Params { get; }

    public TaskId? TaskId { get; init; }

    public TestId? TestId { get; init; }

    public int? Phase { get; init; }

    /// <summary>这个流属于哪个拨号方会话：别的会话的包即使带着这个流编号也不会被接受。</summary>
    public DialerSession? Owner { get; init; }

    public int Port { get; }

    public double DeclaredPps { get; }

    /// <summary>运行这个流的监听方积木（合并时从它取记录）。</summary>
    public RunningBlock? Running { get; set; }

    /// <summary>该流已经处理（交给 Backend）的列块数：测试阶段据此等待全部列块到齐。</summary>
    public int FragmentsDone => _fragmentsDone;

    public void FragmentDone() => Interlocked.Increment(ref _fragmentsDone);

    /// <summary>令牌桶：实际速率超过声明速率的 1.2 倍（外加一点突发余量）就丢弃。</summary>
    public bool TryAdmit(long nowUs)
    {
        double rate = DeclaredPps * 1.2;
        double burst = Math.Max(10, rate * 0.25);
        lock (_bucketGate)
        {
            if (_tokens < 0)
            {
                _tokens = burst;
                _lastRefillUs = nowUs;
            }

            _tokens = Math.Min(burst, _tokens + ((nowUs - _lastRefillUs) / 1_000_000.0 * rate));
            _lastRefillUs = nowUs;
            if (_tokens < 1)
            {
                return false;
            }

            _tokens -= 1;
            return true;
        }
    }
}

public sealed class StreamRegistry : IRecordSource
{
    private readonly ConcurrentDictionary<uint, StreamInfo> _streams = new();
    private int _next;

    public uint Allocate() => (uint)Interlocked.Increment(ref _next);

    public void Add(StreamInfo stream) => _streams[stream.Id] = stream;

    public void Remove(uint id) => _streams.TryRemove(id, out _);

    public bool TryGet(uint id, out StreamInfo stream) => _streams.TryGetValue(id, out stream!);

    public StreamInfo? FindByTask(TaskId task, DialerSession owner) => _streams.Values.FirstOrDefault(s => s.TaskId == task && ReferenceEquals(s.Owner, owner));

    public StreamInfo? FindByTest(TestId test, int phase) => _streams.Values.FirstOrDefault(s => s.TestId == test && s.Phase == phase);

    public IReadOnlyCollection<StreamInfo> All => [.. _streams.Values];

    public RowStore? Records(uint streamId, string streamName) =>
        _streams.TryGetValue(streamId, out StreamInfo? info) ? info.Running?.Block.Records(streamName) : null;
}

/// <summary>
/// Agent 的准入策略（安全边界）：重放窗口 → 流的令牌桶（声明速率的 1.2 倍）→ 源地址学习（NAT 重绑定自动跟随）。
/// 积木本身不含任何策略；这里的任何一步失败，数据面都是静默丢弃并计数、不回复。
/// </summary>
public sealed class AgentAdmission : IAdmission
{
    private readonly StreamRegistry _streams;

    public AgentAdmission(StreamRegistry streams)
    {
        _streams = streams;
    }

    /// <summary>NAT 映射变化，源地址变了：上报 <c>endpoint_changed</c>。</summary>
    public Action<DialerSession>? OnEndpointChanged { get; set; }

    public Admission AdmitDatagram(DataSession session, in MeasurementDatagramHeader header, SocketAddress source, long nowUs)
    {
        if (session is not DialerSession dialer || !_streams.TryGet(header.StreamId, out StreamInfo? stream) || !ReferenceEquals(stream.Owner, dialer) || !dialer.AcceptSequence(header.StreamId, header.Seq))
        {
            return Admission.Invalid;
        }

        if (!stream.TryAdmit(nowUs))
        {
            return Admission.RateLimited;
        }

        if (dialer.UpdateSource(source))
        {
            OnEndpointChanged?.Invoke(dialer);
        }

        return Admission.Accepted;
    }

    public bool AdmitMessage(DataSession session, uint streamId, long nowUs) =>
        _streams.TryGet(streamId, out StreamInfo? stream) && ReferenceEquals(stream.Owner, session) && stream.TryAdmit(nowUs);
}