XFESystem

using System.Net;
using System.Net.Security;
using System.Net.Sockets;
using System.Security.Authentication;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using System.Text;
using XFESystem.Tls;

namespace XFESystem.NativeTls.Tests;

public sealed class TlsHandshakeTests
{
    [Theory]
    [InlineData(SslProtocols.Tls12, "TLS 1.2")]
    [InlineData(SslProtocols.Tls13, "TLS 1.3")]
    public async Task ManagedClientReallyNegotiatesTlsAndTransfersHttpOverLoopback(SslProtocols version, string label)
    {
        using var fixture = new CertificateFixtures(); using var certificate = fixture.Leaf();
        // Windows Schannel does not accept the ephemeral CNG handle returned by
        // CopyWithPrivateKey. Import only this generated test key with the
        // default temporary key lifetime; disposal removes the temporary key.
        using var serverCertificate = X509CertificateLoader.LoadPkcs12(certificate.Export(X509ContentType.Pkcs12), null);
        using var listener = new TcpListener(IPAddress.Loopback, 0); listener.Start();
        using var deadline = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        var serving = Task.Run(async () =>
        {
            using var accepted = await listener.AcceptTcpClientAsync(deadline.Token);
            using var server = new SslStream(accepted.GetStream(), false);
            await server.AuthenticateAsServerAsync(new SslServerAuthenticationOptions { ServerCertificate = serverCertificate, EnabledSslProtocols = version,
                ApplicationProtocols = [SslApplicationProtocol.Http11] }, deadline.Token);
            byte[] request = new byte[4096]; int count = await server.ReadAsync(request, deadline.Token);
            Assert.Contains("GET / HTTP/1.1", Encoding.ASCII.GetString(request, 0, count));
            await server.WriteAsync("HTTP/1.1 200 OK\r\nContent-Length: 7\r\nConnection: close\r\n\r\nXFE TLS"u8.ToArray(), deadline.Token);
            await server.FlushAsync(deadline.Token);
        }, deadline.Token);
        using var socket = new TcpClient(); await socket.ConnectAsync(IPAddress.Loopback, ((IPEndPoint)listener.LocalEndpoint).Port, deadline.Token);
        socket.ReceiveTimeout = socket.SendTimeout = 5000;
        TlsClientConnection connection;
        try { connection = new TlsClientConnection(socket.GetStream(), "test.xfe.invalid", new TlsCertificateValidator([fixture.RootBc]), new HostSecurity()); }
        catch (Exception clientFailure)
        {
            socket.Dispose();
            try { await serving; }
            catch (Exception serverFailure) { throw new AggregateException("Controlled TLS handshake failed on both endpoints.", clientFailure, serverFailure); }
            throw;
        }
        using var closeConnection = connection;
        Assert.True(connection.CertificateValidated); Assert.Equal(label, connection.Protocol);
        connection.Stream.Write("GET / HTTP/1.1\r\nHost: test.xfe.invalid\r\nConnection: close\r\n\r\n"u8);
        connection.Stream.Flush();
        var response = new byte[4096]; int received = connection.Stream.Read(response);
        Assert.Contains("XFE TLS", Encoding.ASCII.GetString(response, 0, received));
        await serving;
    }

    [Theory]
    [InlineData(true, false)]
    [InlineData(false, true)]
    public void MissingTimeOrRandomRejectsBeforeAnyHandshakeWrite(bool badTime, bool badRandom)
    {
        using var fixture = new CertificateFixtures(); using var wire = new CountingStream();
        Assert.Throws<InvalidOperationException>(() => new TlsClientConnection(wire, "test.xfe.invalid",
            new TlsCertificateValidator([fixture.RootBc]), new HostSecurity { BadTime = badTime, BadRandom = badRandom }));
        Assert.Equal(0, wire.BytesWritten);
    }

    [Theory]
    [InlineData("wrong.xfe.invalid", false)]
    [InlineData("test.xfe.invalid", true)]
    public async Task ActualHandshakeRejectsWrongIdentityAndUnknownCa(string host, bool unknownRoot)
    {
        using var fixture = new CertificateFixtures(); using var otherRoot = new CertificateFixtures(); using var certificate = fixture.Leaf();
        using var serverCertificate = X509CertificateLoader.LoadPkcs12(certificate.Export(X509ContentType.Pkcs12), null);
        using var listener = new TcpListener(IPAddress.Loopback, 0); listener.Start();
        using var deadline = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        var serving = Task.Run(async () =>
        {
            using var accepted = await listener.AcceptTcpClientAsync(deadline.Token);
            using var server = new SslStream(accepted.GetStream(), false);
            try
            {
                await server.AuthenticateAsServerAsync(new SslServerAuthenticationOptions { ServerCertificate = serverCertificate,
                    EnabledSslProtocols = SslProtocols.Tls12 }, deadline.Token);
            }
            catch (AuthenticationException) { }
            catch (IOException) { }
        }, deadline.Token);
        using var socket = new TcpClient(); await socket.ConnectAsync(IPAddress.Loopback, ((IPEndPoint)listener.LocalEndpoint).Port, deadline.Token);
        socket.ReceiveTimeout = socket.SendTimeout = 5000;
        Assert.ThrowsAny<Exception>(() => new TlsClientConnection(socket.GetStream(), host,
            new TlsCertificateValidator([unknownRoot ? otherRoot.RootBc : fixture.RootBc]), new HostSecurity()));
        socket.Dispose();
        await serving;
    }

    private sealed class HostSecurity : ITlsSecurityContext
    {
        public bool BadTime { get; init; }
        public bool BadRandom { get; init; }
        public DateTime UtcNow => BadTime ? DateTime.UnixEpoch : DateTime.UtcNow;
        public void FillRandom(Span<byte> destination)
        {
            if (BadRandom) throw new InvalidOperationException("Controlled absent entropy source.");
            RandomNumberGenerator.Fill(destination);
        }
    }

    private sealed class CountingStream : MemoryStream
    {
        public int BytesWritten { get; private set; }
        public override void Write(byte[] buffer, int offset, int count) { BytesWritten += count; base.Write(buffer, offset, count); }
        public override void Write(ReadOnlySpan<byte> buffer) { BytesWritten += buffer.Length; base.Write(buffer); }
    }
}