MoShiro/backend/shared/tests/MoShiro.Measurement.Tests/MsQuicConnectionTests.cs

using System.Net;
using System.Collections.Concurrent;
using System.Runtime.InteropServices;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using Fleet.Core;

namespace MoShiro.Measurement.Tests;

public sealed class MsQuicConnectionTests
{
    [Theory]
    [InlineData(1200, true)]
    [InlineData(1280, true)]
    [InlineData(1472, false)]
    public void DatagramState_UsesTheNativeUint16Alignment(int maximum, bool enabled)
    {
        nint nativeEvent = Marshal.AllocHGlobal(IntPtr.Size + 4);
        try
        {
            Marshal.WriteByte(nativeEvent, IntPtr.Size, enabled ? (byte)1 : (byte)0);
            Marshal.WriteByte(nativeEvent, IntPtr.Size + 1, 0xA5); // padding is not part of MaxSendLength
            Marshal.WriteInt16(nativeEvent, IntPtr.Size + 2, (short)maximum);
            Assert.Equal((enabled, maximum), MsQuicConnection.ReadDatagramState(nativeEvent));
        }
        finally
        {
            Marshal.FreeHGlobal(nativeEvent);
        }
    }

    [Fact]
    public async Task DatagramConnection_CanSendCloseAndReconnect_WithoutBlockingNativeWorkers()
    {
        var log = new CaptureLog();
        var counters = new DataPlaneCounters();
        var router = new MeasurementDatagramRouter(new SingleSession(new DataSession(1, new byte[32])), null,
            new MonoClock(TimeProvider.System), counters, log);
        using ECDsa key = ECDsa.Create(ECCurve.NamedCurves.nistP256);
        var request = new CertificateRequest("CN=localhost", key, HashAlgorithmName.SHA256);
        using X509Certificate2 certificate = request.CreateSelfSigned(DateTimeOffset.UnixEpoch, DateTimeOffset.UnixEpoch.AddYears(100));
        await using var listener = new MsQuicListener(IPAddress.Loopback, [0], certificate,
            (connection, data) =>
            {
                // Bad measurement payloads must be counted without flooding logs.
                for (int i = 0; i < 100; i++)
                {
                    router.Dispatch(connection, data);
                }
                connection.Send(data);
            }, log, null);
        listener.Start();

        for (int attempt = 0; attempt < 3; attempt++)
        {
            var reply = new TaskCompletionSource<byte[]>(TaskCreationOptions.RunContinuationsAsynchronously);
            using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(10));
            MsQuicConnection client = await MsQuicConnection.ConnectAsync(IPAddress.Loopback, listener.BoundPorts[0], 0,
                key.ExportSubjectPublicKeyInfo(), (_, data) => reply.TrySetResult(data.ToArray()), log, timeout.Token);
            try
            {
                while (!client.IsReady)
                {
                    await Task.Delay(10, timeout.Token);
                }

                Assert.InRange(client.MaxDatagramBytes, 64, 1500);
                byte[] payload = Enumerable.Range(0, 64).Select(i => (byte)i).ToArray();
                // DATAGRAM is unreliable, and the peer can receive the first
                // packet before its handshake-complete notification permits replies.
                while (!reply.Task.IsCompleted)
                {
                    Assert.True(client.Send(payload));
                    await Task.WhenAny(reply.Task, Task.Delay(20, timeout.Token));
                    timeout.Token.ThrowIfCancellationRequested();
                }
                Assert.Equal(payload, await reply.Task.WaitAsync(timeout.Token));
                Assert.NotNull(client.TransportSnapshot);
            }
            finally
            {
                // The old callback called RegistrationClose, deadlocking its own
                // worker. A second connection then hung in SetParam/GetParam.
                await client.DisposeAsync().AsTask().WaitAsync(TimeSpan.FromSeconds(3));
                await client.WaitForShutdownAsync().WaitAsync(TimeSpan.FromSeconds(3));
            }
        }
        Assert.True(counters.DroppedInvalid >= 300);
        var dropped = Assert.Single(log.Items, x => x.Message == "measurement datagram dropped");
        Assert.Equal("invalid_header", dropped.Fields["reason"]);
        Assert.NotNull(dropped.Fields["remote"]);
        Assert.NotNull(dropped.Fields["quic_id"]);
        var connected = log.Items.Where(x => x.Message == "MsQuic QUIC handshake completed").ToArray();
        Assert.Equal(6, connected.Select(x => x.Fields["quic_id"]).Distinct().Count());
        Assert.All(connected, entry => Assert.NotNull(entry.Fields["remote"]));
        Assert.DoesNotContain(log.Items.SelectMany(x => x.Fields.Keys), key => key is "key" or "session_id" or "mac" or "signature");
    }

    private sealed class CaptureLog : INodeLog
    {
        public ConcurrentQueue<(string Message, Dictionary<string, object?> Fields)> Items { get; } = new();

        public void Write(NodeLogLevel level, string message, params (string Key, object? Value)[] fields) =>
            Items.Enqueue((message, fields.ToDictionary(x => x.Key, x => x.Value)));
    }
}