MoShiro/backend/shared/MoShiro.Measurement/Net/UdpMux.cs

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

namespace MoShiro.Measurement;

/// <summary>一个通过校验的 UDP 包，交给按 <c>stream_id</c> 登记的处理函数。处理函数在接收循环里同步执行，不能阻塞。</summary>
public readonly ref struct UdpDelivery
{
    private readonly Span<byte> _packet;
    private readonly Socket _socket;
    private readonly SocketAddress _source;
    private readonly DataPlaneCounters _counters;

    internal UdpDelivery(UdpHeader header, Span<byte> packet, long receiveUs, DataSession session, Socket socket, SocketAddress source, DataPlaneCounters counters)
    {
        Header = header;
        _packet = packet;
        ReceiveUs = receiveUs;
        Session = session;
        _socket = socket;
        _source = source;
        _counters = counters;
    }

    public UdpHeader Header { get; }

    public ReadOnlySpan<byte> Packet => _packet;

    /// <summary>包到达的时刻（本节点的单调时钟，微秒）：从接收返回的那一刻取，不含校验的耗时。</summary>
    public long ReceiveUs { get; }

    public DataSession Session { get; }

    /// <summary>
    /// 原地回复：同样的长度，只把 type 改成 <paramref name="type"/> 并重新计算 MAC，发回请求的实际源地址。
    /// 回复与请求等长，所以监听方不能被当成放大器。发送失败（源地址已经不可达）不抛出。
    /// </summary>
    public void ReplyAs(UdpPacketType type)
    {
        _packet[3] = (byte)type;
        UdpPacket.Seal(_packet, Session.Key);
        try
        {
            _socket.SendTo(_packet, SocketFlags.None, _source);
            _counters.PacketOut();
        }
        catch (SocketException)
        {
            // 回复丢失由对端的缺失体现
        }
    }
}

public delegate void UdpHandler(in UdpDelivery delivery);

/// <summary>
/// UDP 数据面：一个或多个套接字，按 <c>stream_id</c> 把包分发给积木。校验顺序：长度、magic、version、type →
/// 流和会话都存在 → MAC → （监听方）<see cref="IAdmission"/>（重放窗口、限速）。任何一步失败都静默丢弃并计数，不回复。
/// 接收循环只做校验和分发，不分配对象。两端共用这一份代码：Agent 监听多个端口并带准入策略，Probe 每个 Agent 一个临时端口。
/// </summary>
public sealed class UdpMux : IDisposable
{
    private const int WindowsUdpConnReset = -1744830452; // SIO_UDP_CONNRESET

    private readonly IDataSessions _sessions;
    private readonly IAdmission? _admission;
    private readonly MonoClock _clock;
    private readonly DataPlaneCounters _counters;
    private readonly INodeLog _log;
    private readonly IPAddress _bind;
    private readonly IReadOnlyList<int> _ports;
    private readonly List<Socket> _sockets = [];
    private readonly ConcurrentDictionary<uint, UdpHandler> _handlers = new();

    /// <param name="ports">要监听的端口（0 表示由系统分配）；空表示一个临时端口。</param>
    public UdpMux(IDataSessions sessions, IAdmission? admission, MonoClock clock, DataPlaneCounters counters, INodeLog log, IPAddress bind, IReadOnlyList<int> ports)
    {
        _sessions = sessions;
        _admission = admission;
        _clock = clock;
        _counters = counters;
        _log = log;
        _bind = bind;
        _ports = ports.Count == 0 ? [0] : ports;
    }

    public IReadOnlyList<int> BoundPorts { get; private set; } = [];

    public DataPlaneCounters Counters => _counters;

    /// <summary>绑定套接字（端口被占用等错误在这里抛出）。</summary>
    public void Bind()
    {
        var bound = new List<int>();
        foreach (int port in _ports)
        {
            var socket = new Socket(_bind.AddressFamily, SocketType.Dgram, ProtocolType.Udp) { ReceiveBufferSize = 4 << 20, SendBufferSize = 4 << 20 };
            if (_bind.AddressFamily == AddressFamily.InterNetworkV6)
            {
                socket.DualMode = true;
            }

            if (OperatingSystem.IsWindows())
            {
                // Windows 上对一个已关闭的端口发送之后，下一次接收会抛 ConnectionReset：关掉这个行为
                socket.IOControl(WindowsUdpConnReset, [0], null);
            }

            socket.Bind(new IPEndPoint(_bind, port));
            bound.Add(((IPEndPoint)socket.LocalEndPoint!).Port);
            _sockets.Add(socket);
        }

        BoundPorts = bound;
    }

    public void Register(uint streamId, UdpHandler handler) => _handlers[streamId] = handler;

    public void Unregister(uint streamId) => _handlers.TryRemove(streamId, out _);

    /// <summary>从第一个套接字发出一个已经封好的包。发送失败（没有路由、缓冲区满）不抛出：丢包本身就是测量结果。</summary>
    public void Send(EndPoint to, ReadOnlySpan<byte> packet)
    {
        try
        {
            _sockets[0].SendTo(packet, SocketFlags.None, to);
            _counters.PacketOut();
        }
        catch (SocketException)
        {
            // 记为丢失
        }
    }

    public async Task RunAsync(CancellationToken ct)
    {
        try
        {
            await Task.WhenAll(_sockets.Select(s => ReceiveLoopAsync(s, ct)));
        }
        catch (ObjectDisposedException)
        {
            // 套接字已释放：停止
        }
    }

    public void Dispose()
    {
        foreach (Socket socket in _sockets)
        {
            socket.Dispose();
        }
    }

    private async Task ReceiveLoopAsync(Socket socket, CancellationToken ct)
    {
        byte[] buffer = new byte[UdpPacket.MaxBytes + 64];
        var source = new SocketAddress(socket.AddressFamily);
        while (!ct.IsCancellationRequested)
        {
            int length;
            try
            {
                length = await socket.ReceiveFromAsync(buffer.AsMemory(), SocketFlags.None, source, ct);
            }
            catch (OperationCanceledException)
            {
                return;
            }
            catch (ObjectDisposedException)
            {
                return;
            }
            catch (SocketException ex)
            {
                // 个别错误（例如对端的 ICMP 不可达）不能让接收循环退出
                _log.Warn("udp receive failed", ("err", ex.SocketErrorCode));
                continue;
            }

            Dispatch(socket, buffer, length, source, _clock.NowUs());
        }
    }

    private void Dispatch(Socket socket, byte[] buffer, int length, SocketAddress source, long receiveUs)
    {
        _counters.PacketIn();
        Span<byte> packet = buffer.AsSpan(0, length);
        if (!UdpPacket.TryParseHeader(packet, out UdpHeader header)
            || !_handlers.TryGetValue(header.StreamId, out UdpHandler? handler)
            || _sessions.Find(header.SessionId) is not { } session
            || !UdpPacket.VerifyMac(packet, session.Key))
        {
            _counters.DropInvalid();
            return;
        }

        if (_admission is not null)
        {
            switch (_admission.AdmitPacket(session, header, source, receiveUs))
            {
                case Admission.Invalid:
                    _counters.DropInvalid();
                    return;
                case Admission.RateLimited:
                    _counters.DropRateLimited();
                    return;
            }
        }

        handler(new UdpDelivery(header, packet, receiveUs, session, socket, source, _counters));
    }
}