QVPN/tests/QVPN.Tests/HealthProbeTests.cs

using System.Diagnostics;
using System.Net;
using System.Net.Security;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using System.Text;
using QVPN.Core;

public static class HealthProbeTests
{
    private static readonly Uri[] Endpoints =
    [
        new("https://first.example/generate_204"),
        new("https://second.example/generate_204"),
        new("https://third.example/generate_204")
    ];

    public static async Task RunAsync()
    {
        await SuccessCancelsPendingAsync();
        await SharedBudgetBoundsLatencyAsync();
        await CallerCancellationPropagatesAsync();
        await FailuresDoNotProduceSuccessAsync();
        await ThirdEndpointRecoversThrottledPeersAsync();
        await FreshConnectionsDoNotReuseSuccessfulSessionAsync();
        await TunRequiresBothPathsAsync();
        await TunPathsShareBudgetAsync();
        await ProxySuccessDoesNotMaskSystemTimeoutAsync();
        await TunCallerCancellationPropagatesAsync();
        await SocksModeSkipsSystemPathAsync();
    }

    public static async Task SuccessCancelsPendingAsync()
    {
        var peersStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        var peersCancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        int started = 0, cancelled = 0;
        using var handler = new StubHandler(async (request, token) =>
        {
            if (request.RequestUri!.Host == Endpoints[0].Host)
            {
                // Success requires both peers to be in flight: a sequential probe cannot pass.
                await peersStarted.Task.WaitAsync(token);
                return new HttpResponseMessage(HttpStatusCode.NoContent);
            }
            if (Interlocked.Increment(ref started) == 2) peersStarted.TrySetResult();
            try { await Task.Delay(Timeout.InfiniteTimeSpan, token); }
            catch (OperationCanceledException)
            {
                if (Interlocked.Increment(ref cancelled) == 2) peersCancelled.TrySetResult();
                throw;
            }
            throw new InvalidOperationException("A pending transport unexpectedly completed.");
        });
        using var http = new HttpClient(handler);
        using var probe = new HealthProbe(http, Endpoints, TimeSpan.FromSeconds(1));
        Check(await probe.CheckAsync(), "Parallel successful endpoint was not accepted.");
        await peersCancelled.Task.WaitAsync(TimeSpan.FromSeconds(1));
        Check(cancelled == 2, "Successful probe did not cancel both losing requests.");
    }

    public static async Task SharedBudgetBoundsLatencyAsync()
    {
        using var release = new CancellationTokenSource();
        var allCancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        int started = 0, cancelled = 0;
        using var handler = new StubHandler(async (_, token) =>
        {
            Interlocked.Increment(ref started);
            using var registration = token.Register(() =>
            {
                if (Interlocked.Increment(ref cancelled) == Endpoints.Length) allCancelled.TrySetResult();
            });
            // Model a transport that acknowledges cancellation but does not finish promptly.
            await Task.Delay(Timeout.InfiniteTimeSpan, release.Token);
            return new HttpResponseMessage(HttpStatusCode.NoContent);
        });
        using var http = new HttpClient(handler);
        using var probe = new HealthProbe(http, Endpoints, TimeSpan.FromMilliseconds(100));
        var watch = Stopwatch.StartNew();
        try
        {
            Check(!await probe.CheckAsync(), "An expired budget reported success.");
            watch.Stop();
            Check(watch.Elapsed < TimeSpan.FromSeconds(1), "Probe waited beyond its shared budget for losing transports.");
            await allCancelled.Task.WaitAsync(TimeSpan.FromSeconds(1));
            Check(started == 3 && cancelled == 3, "Budget did not start and cancel all parallel probes.");
        }
        finally { release.Cancel(); }
    }

    public static async Task CallerCancellationPropagatesAsync()
    {
        var allStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        int started = 0;
        using var handler = new StubHandler(async (_, token) =>
        {
            if (Interlocked.Increment(ref started) == Endpoints.Length) allStarted.TrySetResult();
            await Task.Delay(Timeout.InfiniteTimeSpan, token);
            throw new InvalidOperationException("A cancelled transport unexpectedly completed.");
        });
        using var http = new HttpClient(handler);
        using var probe = new HealthProbe(http, Endpoints);
        using var cancellation = new CancellationTokenSource();
        var running = probe.CheckAsync(cancellation.Token);
        await allStarted.Task.WaitAsync(TimeSpan.FromSeconds(1));
        cancellation.Cancel();
        try
        {
            await running.WaitAsync(TimeSpan.FromSeconds(1));
            throw new InvalidOperationException("Caller cancellation was converted into a health failure.");
        }
        catch (OperationCanceledException error)
        {
            Check(error.CancellationToken == cancellation.Token, "Cancellation did not preserve the caller's token.");
        }
    }

    public static async Task FailuresDoNotProduceSuccessAsync()
    {
        int requests = 0;
        using var handler = new StubHandler((request, _) =>
        {
            Interlocked.Increment(ref requests);
            if (request.RequestUri!.Host == Endpoints[0].Host)
                return Task.FromResult(new HttpResponseMessage(HttpStatusCode.OK) { Content = new StringContent("204") });
            if (request.RequestUri.Host == Endpoints[1].Host)
                return Task.FromResult(new HttpResponseMessage(HttpStatusCode.Redirect));
            return Task.FromException<HttpResponseMessage>(new HttpRequestException("Synthetic transport failure."));
        });
        using var http = new HttpClient(handler);
        using var probe = new HealthProbe(http, Endpoints);
        Check(!await probe.CheckAsync(), "A redirect, HTTP 200, or transport error was accepted as healthy.");
        Check(requests == Endpoints.Length, "Not all independent failure endpoints were checked.");
    }

    public static async Task ThirdEndpointRecoversThrottledPeersAsync()
    {
        using var handler = new StubHandler((request, _) => Task.FromResult(new HttpResponseMessage(
            request.RequestUri!.Host == Endpoints[2].Host ? HttpStatusCode.NoContent : HttpStatusCode.TooManyRequests)));
        using var http = new HttpClient(handler);
        using var probe = new HealthProbe(http, Endpoints);
        Check(await probe.CheckAsync(), "Throttled peers masked the third successful endpoint.");
    }

    public static async Task FreshConnectionsDoNotReuseSuccessfulSessionAsync()
    {
        using var key = RSA.Create(2048);
        var certificateRequest = new CertificateRequest("CN=localhost", key, HashAlgorithmName.SHA256, RSASignaturePadding.Pkcs1);
        using var generated = certificateRequest.CreateSelfSigned(DateTimeOffset.UtcNow.AddMinutes(-1), DateTimeOffset.UtcNow.AddMinutes(5));
        // Schannel cannot serve TLS with the ephemeral key returned above.
        // Import a temporary user key that is removed when the certificate is disposed.
        using var certificate = X509CertificateLoader.LoadPkcs12(generated.Export(X509ContentType.Pfx), null, X509KeyStorageFlags.UserKeySet);
        using var serverStop = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        var listener = new TcpListener(IPAddress.Loopback, 0);
        listener.Start();
        int connections = 0;
        var server = Task.Run(async () =>
        {
            var workers = new List<Task>();
            try
            {
                while (true)
                {
                    var tcp = await listener.AcceptTcpClientAsync(serverStop.Token);
                    Interlocked.Increment(ref connections);
                    workers.Add(ServeKeepAliveAsync(tcp, certificate, serverStop.Token));
                }
            }
            catch (OperationCanceledException) when (serverStop.IsCancellationRequested) { }
            finally { await Task.WhenAll(workers); }
        });
        using var handler = new SocketsHttpHandler
        {
            UseProxy = false,
            SslOptions = new SslClientAuthenticationOptions
            {
                // Trust only this synthetic loopback server's ephemeral certificate.
                RemoteCertificateValidationCallback = (_, remote, _, _) => remote?.GetCertHashString() == certificate.GetCertHashString()
            }
        };
        using var http = new HttpClient(handler);
        using var probe = new HealthProbe(http,
            [new Uri("https://127.0.0.1:" + ((IPEndPoint)listener.LocalEndpoint).Port + "/generate_204")]);
        try
        {
            Check(await probe.CheckAsync(), "First synthetic HTTPS check failed.");
            Check(await probe.CheckAsync(), "Second synthetic HTTPS check failed.");
            Check(Volatile.Read(ref connections) == 2, "Health reused an old successful HTTPS session instead of testing a new connection.");
        }
        finally
        {
            serverStop.Cancel();
            await server.WaitAsync(TimeSpan.FromSeconds(2));
            listener.Stop();
        }
    }

    private static async Task ServeKeepAliveAsync(TcpClient tcp, X509Certificate2 certificate, CancellationToken token)
    {
        using (tcp)
        using (var tls = new SslStream(tcp.GetStream()))
        {
            try
            {
                await tls.AuthenticateAsServerAsync(new SslServerAuthenticationOptions { ServerCertificate = certificate }, token);
                var buffer = new byte[4096];
                var headers = new StringBuilder();
                while (true)
                {
                    int count = await tls.ReadAsync(buffer, token);
                    if (count == 0) return;
                    headers.Append(Encoding.ASCII.GetString(buffer, 0, count));
                    if (!headers.ToString().Contains("\r\n\r\n", StringComparison.Ordinal)) continue;
                    bool close = headers.ToString().Contains("\r\nConnection: close\r\n", StringComparison.OrdinalIgnoreCase);
                    headers.Clear();
                    // Permit persistent connections unless the health request
                    // explicitly asks to close this session.
                    await tls.WriteAsync("HTTP/1.1 204 No Content\r\nContent-Length: 0\r\n\r\n"u8.ToArray(), token);
                    if (close) return;
                }
            }
            catch (OperationCanceledException) when (token.IsCancellationRequested) { }
            catch (System.Security.Authentication.AuthenticationException) when (token.IsCancellationRequested) { }
            catch (IOException) { }
        }
    }

    public static async Task TunRequiresBothPathsAsync()
    {
        foreach (var (proxyHealthy, systemHealthy) in new[] { (true, false), (false, true), (true, true) })
        {
            using var proxyHandler = new StubHandler((_, _) => Task.FromResult(new HttpResponseMessage(proxyHealthy ? HttpStatusCode.NoContent : HttpStatusCode.BadGateway)));
            using var systemHandler = new StubHandler(async (_, token) =>
            {
                // Proxy success must not cancel the still-running system check.
                await Task.Delay(25, token);
                return new HttpResponseMessage(systemHealthy ? HttpStatusCode.NoContent : HttpStatusCode.BadGateway);
            });
            using var proxy = new HttpClient(proxyHandler);
            using var system = new HttpClient(systemHandler);
            using var probe = new HealthProbe(proxy, Endpoints, TimeSpan.FromSeconds(1), system);
            var result = await probe.CheckPathsAsync(true);
            Check(result.ProxyAvailable == proxyHealthy && result.SystemAvailable == systemHealthy, "TUN check confused proxy and system reachability.");
            Check(result.Healthy == (proxyHealthy && systemHealthy), "TUN health accepted only one working path.");
            Check(result.SystemPathFailed == (proxyHealthy && !systemHealthy), "Local tunnel failure was attributed to the VPN endpoint.");
        }
    }

    public static async Task TunPathsShareBudgetAsync()
    {
        int started = 0, cancelled = 0;
        var allCancelled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        using var release = new CancellationTokenSource();
        using var handler = new StubHandler(async (_, token) =>
        {
            Interlocked.Increment(ref started);
            using var registration = token.Register(() => { if (Interlocked.Increment(ref cancelled) == Endpoints.Length * 2) allCancelled.TrySetResult(); });
            await Task.Delay(Timeout.InfiniteTimeSpan, release.Token);
            return new HttpResponseMessage(HttpStatusCode.NoContent);
        });
        using var proxy = new HttpClient(handler, false);
        using var system = new HttpClient(handler, false);
        using var probe = new HealthProbe(proxy, Endpoints, TimeSpan.FromMilliseconds(100), system);
        var watch = Stopwatch.StartNew();
        try
        {
            var result = await probe.CheckPathsAsync(true);
            Check(!result.Healthy && !result.ProxyAvailable && result.SystemAvailable == false, "Timed-out TUN paths reported health.");
            Check(watch.Elapsed < TimeSpan.FromSeconds(1), "TUN paths exceeded the shared deadline.");
            await allCancelled.Task.WaitAsync(TimeSpan.FromSeconds(1));
            Check(started == Endpoints.Length * 2 && cancelled == Endpoints.Length * 2, "Both TUN paths were not checked and cancelled in parallel.");
        }
        finally { release.Cancel(); }
    }

    public static async Task TunCallerCancellationPropagatesAsync()
    {
        int started = 0;
        var allStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        using var handler = new StubHandler(async (_, token) =>
        {
            if (Interlocked.Increment(ref started) == Endpoints.Length * 2) allStarted.TrySetResult();
            await Task.Delay(Timeout.InfiniteTimeSpan, token);
            return new HttpResponseMessage(HttpStatusCode.NoContent);
        });
        using var proxy = new HttpClient(handler, false);
        using var system = new HttpClient(handler, false);
        using var probe = new HealthProbe(proxy, Endpoints, systemHttp: system);
        using var cancellation = new CancellationTokenSource();
        var running = probe.CheckPathsAsync(true, cancellation.Token);
        await allStarted.Task.WaitAsync(TimeSpan.FromSeconds(1));
        cancellation.Cancel();
        try
        {
            await running.WaitAsync(TimeSpan.FromSeconds(1));
            throw new InvalidOperationException("TUN cancellation was reported as endpoint failure.");
        }
        catch (OperationCanceledException error) { Check(error.CancellationToken == cancellation.Token, "TUN cancellation lost the caller's token."); }
    }

    public static async Task ProxySuccessDoesNotMaskSystemTimeoutAsync()
    {
        using var release = new CancellationTokenSource();
        using var proxyHandler = new StubHandler((_, _) => Task.FromResult(new HttpResponseMessage(HttpStatusCode.NoContent)));
        using var systemHandler = new StubHandler(async (_, _) =>
        {
            // A Windows route/DNS transport need not honor cancellation promptly.
            await Task.Delay(Timeout.InfiniteTimeSpan, release.Token);
            return new HttpResponseMessage(HttpStatusCode.NoContent);
        });
        using var proxy = new HttpClient(proxyHandler);
        using var system = new HttpClient(systemHandler);
        using var probe = new HealthProbe(proxy, Endpoints, TimeSpan.FromMilliseconds(100), system);
        try
        {
            var result = await probe.CheckPathsAsync(true).WaitAsync(TimeSpan.FromSeconds(1));
            Check(result.SystemPathFailed && !result.Healthy, "Proxy success hid a timed-out Windows path or penalized the endpoint.");
        }
        finally { release.Cancel(); }
    }

    public static async Task SocksModeSkipsSystemPathAsync()
    {
        int systemRequests = 0;
        using var proxyHandler = new StubHandler((_, _) => Task.FromResult(new HttpResponseMessage(HttpStatusCode.NoContent)));
        using var systemHandler = new StubHandler((_, _) =>
        {
            Interlocked.Increment(ref systemRequests);
            return Task.FromResult(new HttpResponseMessage(HttpStatusCode.BadGateway));
        });
        using var proxy = new HttpClient(proxyHandler);
        using var system = new HttpClient(systemHandler);
        using var probe = new HealthProbe(proxy, Endpoints, systemHttp: system);
        var result = await probe.CheckPathsAsync(false);
        Check(result.Healthy && result.SystemAvailable is null && systemRequests == 0, "SOCKS diagnostics unexpectedly required the system tunnel.");
    }

    private static void Check(bool condition, string message)
    {
        if (!condition) throw new InvalidOperationException(message);
    }

    private sealed class StubHandler(Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> send) : HttpMessageHandler
    {
        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken token) => send(request, token);
    }
}