//! Client-side TLS over a `std.Io.net.Stream`, used for upstream DoT.
//!
//! `std.crypto.tls.Client` holds its `Io.Reader`/`Io.Writer` by value and its
//! `input`/`output` point at the stream reader/writer stored beside it, so
//! `TlsStream` is pinned: it MUST NOT be moved or copied after `init`.

const std = @import("std");
const tls = std.crypto.tls;
const net = std.Io.net;
const Certificate = std.crypto.Certificate;

const log = std.log.scoped(.tls_client);

pub const ErrorClass = enum { handshake, certificate, io, protocol };

pub const InitError = tls.Client.InitError || error{CertificateBundleLoadFailure};

/// Exhaustive over `std.crypto.tls.Client.InitError`. A new error in that stdlib set
/// breaks this switch at compile time, which is the point.
fn classifyInit(err: tls.Client.InitError) ErrorClass {
    return switch (err) {
        error.InsufficientEntropy,
        error.TlsAlert,
        error.TlsUnexpectedMessage,
        error.TlsIllegalParameter,
        error.TlsDecryptFailure,
        error.TlsDecryptError,
        error.TlsRecordOverflow,
        error.TlsBadRecordMac,
        => .handshake,

        error.CertificateFieldHasInvalidLength,
        error.CertificateHostMismatch,
        error.CertificatePublicKeyInvalid,
        error.CertificateExpired,
        error.CertificateFieldHasWrongDataType,
        error.CertificateIssuerMismatch,
        error.CertificateNotYetValid,
        error.CertificateSignatureAlgorithmMismatch,
        error.CertificateSignatureAlgorithmUnsupported,
        error.CertificateSignatureInvalid,
        error.CertificateSignatureInvalidLength,
        error.CertificateSignatureNamedCurveUnsupported,
        error.CertificateSignatureUnsupportedBitCount,
        error.CertificateTimeInvalid,
        error.CertificateHasUnrecognizedObjectId,
        error.CertificateHasInvalidBitString,
        error.UnsupportedCertificateVersion,
        error.TlsCertificateNotVerified,
        error.TlsBadSignatureScheme,
        error.TlsBadRsaSignatureBitCount,
        error.SignatureVerificationFailed,
        error.InvalidSignature,
        => .certificate,

        error.ReadFailed,
        error.WriteFailed,
        error.Canceled,
        error.DiskQuota,
        error.LockViolation,
        error.NotOpenForWriting,
        => .io,

        error.TlsConnectionTruncated,
        error.TlsDecodeError,
        error.InvalidEncoding,
        error.IdentityElement,
        error.MessageTooLong,
        error.NegativeIntoUnsigned,
        error.TargetTooSmall,
        error.BufferTooSmall,
        error.NotSquare,
        error.NonCanonical,
        error.WeakPublicKey,
        => .protocol,
    };
}

/// Errors reachable after the handshake, or added by this module.
fn classifyOther(err: anyerror) ErrorClass {
    return switch (err) {
        error.CertificateBundleLoadFailure,
        error.CertificateIssuerNotFound,
        => .certificate,

        error.TlsBadLength,
        error.TlsSequenceOverflow,
        => .protocol,

        error.EndOfStream,
        error.ConnectionResetByPeer,
        error.ConnectionRefused,
        error.ConnectionTimedOut,
        error.BrokenPipe,
        error.NetworkDown,
        error.NetworkUnreachable,
        error.HostUnreachable,
        error.Timeout,
        error.SocketUnconnected,
        error.SystemResources,
        error.AccessDenied,
        => .io,

        else => .protocol,
    };
}

pub fn classify(err: anyerror) ErrorClass {
    inline for (@typeInfo(tls.Client.InitError).error_set.?) |member| {
        if (err == @field(anyerror, member.name)) {
            return classifyInit(@field(tls.Client.InitError, member.name));
        }
    }
    return classifyOther(err);
}

pub const TlsStream = struct {
    stream_reader: net.Stream.Reader,
    stream_writer: net.Stream.Writer,
    client: tls.Client,

    pub const Options = struct {
        /// SNI name, and the name matched against the leaf certificate. Sent and
        /// matched under both `ca` settings.
        host: []const u8,
        /// `.insecure_skip_verify` drops only the chain of trust: the stdlib
        /// verifies the leaf certificate host name before it consults `ca`
        /// (crypto/tls/Client.zig:645), so `host` must still match the leaf
        /// certificate or the handshake fails with `error.CertificateHostMismatch`.
        /// Issuer authorization and certificate expiry are not checked.
        ca: enum { system, insecure_skip_verify },
        /// Plaintext read buffer.
        read_buffer: []u8,
        /// Plaintext write buffer.
        write_buffer: []u8,
        /// Ciphertext read buffer; must hold at least `tls.Client.min_buffer_len` bytes.
        stream_read_buffer: []u8,
        /// Ciphertext write buffer; must hold at least `tls.Client.min_buffer_len` bytes.
        stream_write_buffer: []u8,
    };

    /// In-place init. `self` must not move afterwards.
    ///
    /// With `.ca = .system`, an empty `bundle` is filled by
    /// `Certificate.Bundle.rescan` under `bundle_lock`. The caller owns the
    /// lifetime of `bundle`, `bundle_lock` and `stream`.
    pub fn init(
        self: *TlsStream,
        io: std.Io,
        stream: *net.Stream,
        bundle: *Certificate.Bundle,
        bundle_lock: *std.Io.RwLock,
        gpa: std.mem.Allocator,
        options: Options,
    ) InitError!void {
        std.debug.assert(options.stream_read_buffer.len >= tls.Client.min_buffer_len);
        std.debug.assert(options.stream_write_buffer.len >= tls.Client.min_buffer_len);

        const now = std.Io.Clock.real.now(io);

        switch (options.ca) {
            .system => try ensureBundle(io, bundle, bundle_lock, gpa, now),
            .insecure_skip_verify => {},
        }

        var entropy: [tls.Client.Options.entropy_len]u8 = undefined;
        io.random(&entropy);

        self.* = .{
            .stream_reader = stream.reader(io, options.stream_read_buffer),
            .stream_writer = stream.writer(io, options.stream_write_buffer),
            .client = undefined,
        };

        self.client = try tls.Client.init(
            &self.stream_reader.interface,
            &self.stream_writer.interface,
            .{
                .host = .{ .explicit = options.host },
                .ca = switch (options.ca) {
                    .system => .{ .bundle = .{
                        .gpa = gpa,
                        .io = io,
                        .lock = bundle_lock,
                        .bundle = bundle,
                    } },
                    .insecure_skip_verify => .no_verification,
                },
                .read_buffer = options.read_buffer,
                .write_buffer = options.write_buffer,
                .entropy = &entropy,
                .realtime_now = now,
            },
        );
    }

    fn ensureBundle(
        io: std.Io,
        bundle: *Certificate.Bundle,
        bundle_lock: *std.Io.RwLock,
        gpa: std.mem.Allocator,
        now: std.Io.Timestamp,
    ) InitError!void {
        {
            try bundle_lock.lockShared(io);
            defer bundle_lock.unlockShared(io);
            if (bundle.map.count() != 0) return;
        }

        try bundle_lock.lock(io);
        defer bundle_lock.unlock(io);
        if (bundle.map.count() != 0) return;

        // A partial rescan leaves entries in `map`, which the fast path above would
        // read as "already loaded". Reset so the next init retries a full rescan.
        bundle.rescan(gpa, io, now) catch |err| {
            bundle.deinit(gpa);
            bundle.* = .empty;
            return switch (err) {
                error.Canceled => error.Canceled,
                else => error.CertificateBundleLoadFailure,
            };
        };
    }

    /// Plaintext reader.
    pub fn reader(self: *TlsStream) *std.Io.Reader {
        return &self.client.reader;
    }

    /// Plaintext writer.
    pub fn writer(self: *TlsStream) *std.Io.Writer {
        return &self.client.writer;
    }

    /// Pushes buffered plaintext all the way to the socket.
    ///
    /// Both flushes are required. `tls.Client.flush` only encrypts what the
    /// plaintext writer holds into the socket writer's buffer and calls
    /// `advance` (crypto/tls/Client.zig:999); it never flushes that writer, so
    /// on its own it leaves the record sitting in this process. A caller that
    /// then waits for a reply waits until the peer gives up.
    pub fn flush(self: *TlsStream) std.Io.Writer.Error!void {
        try self.client.writer.flush();
        try self.stream_writer.interface.flush();
    }

    /// Sends close_notify and flushes the socket. Does not close the underlying
    /// stream; the caller owns it.
    pub fn close(self: *TlsStream) void {
        self.client.end() catch |err| {
            log.debug("close_notify failed: {s}", .{@errorName(err)});
        };
        self.stream_writer.interface.flush() catch |err| {
            log.debug("flush after close_notify failed: {s}", .{@errorName(err)});
        };
    }
};

test "classify maps handshake errors" {
    try std.testing.expectEqual(ErrorClass.handshake, classify(error.TlsAlert));
    try std.testing.expectEqual(ErrorClass.handshake, classify(error.TlsIllegalParameter));
    try std.testing.expectEqual(ErrorClass.handshake, classify(error.InsufficientEntropy));
}

test "classify maps certificate errors" {
    try std.testing.expectEqual(ErrorClass.certificate, classify(error.CertificateExpired));
    try std.testing.expectEqual(ErrorClass.certificate, classify(error.CertificateHostMismatch));
    try std.testing.expectEqual(ErrorClass.certificate, classify(error.TlsCertificateNotVerified));
    try std.testing.expectEqual(ErrorClass.certificate, classify(error.CertificateBundleLoadFailure));
}

test "classify maps io errors" {
    try std.testing.expectEqual(ErrorClass.io, classify(error.ReadFailed));
    try std.testing.expectEqual(ErrorClass.io, classify(error.WriteFailed));
    try std.testing.expectEqual(ErrorClass.io, classify(error.Canceled));
    try std.testing.expectEqual(ErrorClass.io, classify(error.ConnectionResetByPeer));
}

test "classify maps protocol errors" {
    try std.testing.expectEqual(ErrorClass.protocol, classify(error.TlsDecodeError));
    try std.testing.expectEqual(ErrorClass.protocol, classify(error.InvalidEncoding));
    try std.testing.expectEqual(ErrorClass.protocol, classify(error.TlsBadLength));
    try std.testing.expectEqual(ErrorClass.protocol, classify(error.TlsConnectionTruncated));
    try std.testing.expectEqual(ErrorClass.protocol, classify(error.SomethingNobodyMapped));
}

test "classify covers every InitError member" {
    inline for (@typeInfo(tls.Client.InitError).error_set.?) |member| {
        _ = classifyInit(@field(tls.Client.InitError, member.name));
    }
}
