//! Server-side TLS termination backed by Mbed TLS 3.6, exposing plaintext as
//! `std.Io.Reader` / `std.Io.Writer` so `std.http.Server` and the DoT server can
//! sit on top of any accepted `std.Io.net.Stream`.
//!
//! The Mbed TLS surface is declared by hand below. No `@cImport` — the context
//! structs stay opaque and are heap-allocated using sizes reported by
//! `src/platform/mbedtls_shim.c`.

const std = @import("std");
const Io = std.Io;
const Reader = std.Io.Reader;
const Writer = std.Io.Writer;
const assert = std.debug.assert;

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

/// Every Mbed TLS context is allocated at this alignment; `ServerContext.init`
/// asserts the shim agrees it is enough.
const context_alignment = 16;
const ContextAlignment: std.mem.Alignment = .fromByteUnits(context_alignment);

pub const InitError = error{
    OutOfMemory,
    /// The certificate chain PEM could not be parsed.
    CertParse,
    /// The private key PEM could not be parsed.
    KeyParse,
    /// The private key does not belong to the leaf certificate.
    KeyMismatch,
    /// Seeding the CTR-DRBG from the platform entropy source failed.
    EntropyFailed,
    /// Mbed TLS rejected the server configuration.
    ConfigFailed,
};

pub const AcceptError = error{
    OutOfMemory,
    /// Binding the connection to the shared configuration failed.
    SetupFailed,
    /// The peer went away before the handshake completed.
    PeerClosed,
    HandshakeFailed,
};

/// Parsed certificate chain, private key, RNG and `mbedtls_ssl_config` for one
/// listener. Create it once and share it across every accepted connection.
///
/// Safe to move: every Mbed TLS context lives on the heap, so the pointers Mbed
/// TLS stores internally stay valid.
pub const ServerContext = struct {
    entropy: Blocks.Entropy,
    drbg: Blocks.Drbg,
    cert: Blocks.Cert,
    key: Blocks.Key,
    config: Blocks.Config,

    const Blocks = struct {
        const Entropy = Block(EntropyContext);
        const Drbg = Block(CtrDrbgContext);
        const Cert = Block(X509Crt);
        const Key = Block(PkContext);
        const Config = Block(SslConfig);
    };

    /// `alpn`, when non-null, is a NULL-terminated list of protocol names in
    /// decreasing preference order (e.g. `&.{"http/1.1"}` as a
    /// comptime-constant `[_:null]?[*:0]const u8` array). Mbed TLS records the
    /// pointer, not a copy, so the array must outlive this `ServerContext` —
    /// pass a comptime-constant array, never stack or heap memory that can go
    /// away first. Clients that send no ALPN extension still connect; Mbed TLS
    /// only enforces the list against clients that offer one.
    pub fn init(
        gpa: std.mem.Allocator,
        cert_pem: [:0]const u8,
        key_pem: [:0]const u8,
        alpn: ?[*:null]const ?[*:0]const u8,
    ) InitError!ServerContext {
        assert(nx_max_context_alignment() <= context_alignment);
        assert(transcriptionsAgree());

        // TLS 1.3 is on in the stock config; its key schedule runs through PSA.
        const psa_status = psa_crypto_init();
        if (psa_status != 0) {
            log.err("psa_crypto_init failed: {d}", .{psa_status});
            return error.ConfigFailed;
        }

        var entropy = try Blocks.Entropy.alloc(gpa, nx_sizeof_entropy_context());
        errdefer entropy.free(gpa);
        mbedtls_entropy_init(entropy.ptr);
        errdefer mbedtls_entropy_free(entropy.ptr);

        var drbg = try Blocks.Drbg.alloc(gpa, nx_sizeof_ctr_drbg_context());
        errdefer drbg.free(gpa);
        mbedtls_ctr_drbg_init(drbg.ptr);
        errdefer mbedtls_ctr_drbg_free(drbg.ptr);

        const personalization = "nxdns-tls-server";
        try check(mbedtls_ctr_drbg_seed(
            drbg.ptr,
            mbedtls_entropy_func,
            entropy.ptr,
            personalization,
            personalization.len,
        ), "ctr_drbg_seed", error.EntropyFailed);

        var cert = try Blocks.Cert.alloc(gpa, nx_sizeof_x509_crt());
        errdefer cert.free(gpa);
        mbedtls_x509_crt_init(cert.ptr);
        errdefer mbedtls_x509_crt_free(cert.ptr);

        // PEM input must include the terminating zero byte in the length.
        try check(
            mbedtls_x509_crt_parse(cert.ptr, cert_pem.ptr, cert_pem.len + 1),
            "x509_crt_parse",
            error.CertParse,
        );

        var key = try Blocks.Key.alloc(gpa, nx_sizeof_pk_context());
        errdefer key.free(gpa);
        mbedtls_pk_init(key.ptr);
        errdefer mbedtls_pk_free(key.ptr);

        try check(mbedtls_pk_parse_key(
            key.ptr,
            key_pem.ptr,
            key_pem.len + 1,
            null,
            0,
            mbedtls_ctr_drbg_random,
            drbg.ptr,
        ), "pk_parse_key", error.KeyParse);

        // mbedtls_ssl_conf_own_cert accepts a mismatched pair; only this call
        // rejects one, and it needs an RNG for the blinded EC comparison.
        try check(mbedtls_pk_check_pair(
            nx_x509_crt_pk(cert.ptr),
            key.ptr,
            mbedtls_ctr_drbg_random,
            drbg.ptr,
        ), "pk_check_pair", error.KeyMismatch);

        var config = try Blocks.Config.alloc(gpa, nx_sizeof_ssl_config());
        errdefer config.free(gpa);
        mbedtls_ssl_config_init(config.ptr);
        errdefer mbedtls_ssl_config_free(config.ptr);

        try check(mbedtls_ssl_config_defaults(
            config.ptr,
            ssl_is_server,
            ssl_transport_stream,
            ssl_preset_default,
        ), "ssl_config_defaults", error.ConfigFailed);

        mbedtls_ssl_conf_rng(config.ptr, mbedtls_ctr_drbg_random, drbg.ptr);
        // nxdns never asks clients for a certificate.
        mbedtls_ssl_conf_authmode(config.ptr, ssl_verify_none);
        try check(
            mbedtls_ssl_conf_own_cert(config.ptr, cert.ptr, key.ptr),
            "ssl_conf_own_cert",
            error.ConfigFailed,
        );

        if (alpn) |protos| try check(
            mbedtls_ssl_conf_alpn_protocols(config.ptr, protos),
            "ssl_conf_alpn_protocols",
            error.ConfigFailed,
        );

        return .{
            .entropy = entropy,
            .drbg = drbg,
            .cert = cert,
            .key = key,
            .config = config,
        };
    }

    pub fn deinit(self: *ServerContext, gpa: std.mem.Allocator) void {
        mbedtls_ssl_config_free(self.config.ptr);
        mbedtls_pk_free(self.key.ptr);
        mbedtls_x509_crt_free(self.cert.ptr);
        mbedtls_ctr_drbg_free(self.drbg.ptr);
        mbedtls_entropy_free(self.entropy.ptr);

        self.config.free(gpa);
        self.key.free(gpa);
        self.cert.free(gpa);
        self.drbg.free(gpa);
        self.entropy.free(gpa);

        self.* = undefined;
    }
};

/// One TLS connection. Pinned once `accept` succeeds: the `std.Io` interfaces
/// are recovered from their addresses inside this struct, and Mbed TLS holds a
/// pointer to it as the BIO context.
pub const ServerStream = struct {
    ssl: Block(SslContext),
    io: Io,
    net_reader: Io.net.Stream.Reader,
    net_writer: Io.net.Stream.Writer,
    reader_interface: Reader,
    writer_interface: Writer,
    /// Set once the peer's close_notify is seen. A transport EOF without one
    /// is a truncated stream, not a close.
    peer_closed: bool,
    /// The concrete failure behind `error.ReadFailed`, in the arrangement
    /// `std.crypto.tls.Client.read_err` uses.
    read_err: ?ReadError,
    /// Most recent Mbed TLS code behind `read_err`.
    read_code: c_int,
    /// Most recent negative Mbed TLS code behind `error.WriteFailed`.
    write_code: c_int,
    /// The transport failure behind the last failed `bioRecv`. Mbed TLS only
    /// forwards its own generic code, so without this stash a routine idle
    /// cancel and a peer reset are indistinguishable to the caller.
    recv_cause: ?anyerror,
    /// The transport failure behind the last failed `bioSend`.
    send_cause: ?anyerror,

    pub const ReadError = error{
        /// The peer closed the transport without sending close_notify. Any
        /// data it withheld is indistinguishable from data an attacker cut
        /// off, so the stream never ends cleanly this way.
        TlsConnectionTruncated,
        /// Mbed TLS rejected the record; `read_code` holds its code.
        TlsFailed,
        /// nxdns canceled the read — an idle budget expiring or a shutdown,
        /// never a TLS fault. Kept distinct so callers do not count it.
        Canceled,
    };

    /// Zero-length transport buffers: Mbed TLS keeps its own record buffers, so
    /// a second layer of buffering under the BIO callbacks would only add
    /// copies. Mutable because `Stream.reader`/`Stream.writer` take `[]u8`.
    var no_transport_buffer: [0]u8 = .{};

    const reader_vtable: Reader.VTable = .{
        .stream = readerStream,
        .readVec = readerReadVec,
    };

    const writer_vtable: Writer.VTable = .{
        .drain = writerDrain,
    };

    /// Runs the TLS handshake over an already accepted TCP stream. `stream`
    /// must outlive the `ServerStream`; closing it stays the caller's job.
    ///
    /// `read_buffer` and `write_buffer` hold plaintext for the `reader` and
    /// `writer` interfaces.
    pub fn accept(
        self: *ServerStream,
        gpa: std.mem.Allocator,
        ctx: *ServerContext,
        io: Io,
        stream: *Io.net.Stream,
        read_buffer: []u8,
        write_buffer: []u8,
    ) AcceptError!void {
        var ssl = try Block(SslContext).alloc(gpa, nx_sizeof_ssl_context());
        errdefer ssl.free(gpa);
        mbedtls_ssl_init(ssl.ptr);
        errdefer mbedtls_ssl_free(ssl.ptr);

        self.* = .{
            .ssl = ssl,
            .io = io,
            .net_reader = stream.reader(io, &no_transport_buffer),
            .net_writer = stream.writer(io, &no_transport_buffer),
            .reader_interface = .{
                .vtable = &reader_vtable,
                .buffer = read_buffer,
                .seek = 0,
                .end = 0,
            },
            .writer_interface = .{
                .vtable = &writer_vtable,
                .buffer = write_buffer,
            },
            .peer_closed = false,
            .read_err = null,
            .read_code = 0,
            .write_code = 0,
            .recv_cause = null,
            .send_cause = null,
        };

        try check(mbedtls_ssl_setup(self.ssl.ptr, ctx.config.ptr), "ssl_setup", error.SetupFailed);
        mbedtls_ssl_set_bio(self.ssl.ptr, self, bioSend, bioRecv, null);

        while (true) {
            const rc = mbedtls_ssl_handshake(self.ssl.ptr);
            if (rc == 0) return;
            if (isRetry(rc)) continue;
            if (rc == err_conn_eof or rc == err_peer_close_notify) return error.PeerClosed;
            // A probe that connects and then drops reaches this on every
            // attempt, so a peer-driven handshake failure logs at debug.
            report("ssl_handshake", rc, self.level(self.recv_cause orelse self.send_cause));
            return error.HandshakeFailed;
        }
    }

    /// Plaintext read side. The `ServerStream` must not move after `accept`.
    pub fn reader(self: *ServerStream) *Reader {
        return &self.reader_interface;
    }

    /// Plaintext write side. Buffered — call `flush` to force records out.
    pub fn writer(self: *ServerStream) *Writer {
        return &self.writer_interface;
    }

    /// Flushes pending plaintext, sends close_notify, and frees the connection
    /// context. Does not close the underlying TCP stream.
    pub fn close(self: *ServerStream, gpa: std.mem.Allocator) void {
        self.writer_interface.flush() catch {
            log.debug("flush before close_notify failed: mbedtls {d}", .{self.write_code});
        };
        while (true) {
            const rc = mbedtls_ssl_close_notify(self.ssl.ptr);
            if (isRetry(rc)) continue;
            // close_notify against a socket the peer already dropped fails
            // every time; that is the peer's doing, not a fault worth a warn.
            if (rc != 0) report("ssl_close_notify", rc, self.level(self.send_cause orelse self.recv_cause));
            break;
        }
        mbedtls_ssl_free(self.ssl.ptr);
        self.ssl.free(gpa);
        self.* = undefined;
    }

    fn readerStream(r: *Reader, w: *Writer, limit: Io.Limit) Reader.StreamError!usize {
        _ = w;
        _ = limit;
        const self: *ServerStream = @alignCast(@fieldParentPtr("reader_interface", r));
        return self.readIntoBuffer();
    }

    fn readerReadVec(r: *Reader, data: [][]u8) Reader.Error!usize {
        _ = data;
        const self: *ServerStream = @alignCast(@fieldParentPtr("reader_interface", r));
        return self.readIntoBuffer();
    }

    /// Decrypts into `Reader.buffer` and reports zero bytes streamed, the
    /// arrangement `Reader.VTable` documents for buffer-writing implementations.
    fn readIntoBuffer(self: *ServerStream) Reader.Error!usize {
        const r = &self.reader_interface;
        if (self.read_err) |err| return self.failRead(err, self.read_code);
        if (self.peer_closed) return error.EndOfStream;

        if (r.end == r.buffer.len and r.seek != 0) {
            const pending = r.buffer[r.seek..r.end];
            @memmove(r.buffer[0..pending.len], pending);
            r.seek = 0;
            r.end = pending.len;
        }
        const dest = r.buffer[r.end..];
        if (dest.len == 0) return 0;

        while (true) {
            const rc = mbedtls_ssl_read(self.ssl.ptr, dest.ptr, dest.len);
            if (isRetry(rc)) continue;
            if (rc == err_peer_close_notify) {
                self.peer_closed = true;
                return error.EndOfStream;
            }
            // mbedtls_ssl_read reports a transport EOF that arrived without a
            // close_notify as 0 (ssl_msg.c maps MBEDTLS_ERR_SSL_CONN_EOF to 0
            // for application data); the raw code reaches us from the
            // handshake it may run internally. Both are a truncation attack
            // until the peer proves otherwise, so neither ends the stream
            // cleanly.
            if (rc == 0 or rc == err_conn_eof) {
                // Debug, not warn: any client that drops its TCP connection
                // reaches this, so a louder level would be a log-spam vector.
                log.debug("peer closed the transport without close_notify", .{});
                return self.failRead(error.TlsConnectionTruncated, rc);
            }
            if (rc < 0) {
                const err = readErrorFor(self.recv_cause);
                report("ssl_read", rc, self.level(self.recv_cause));
                return self.failRead(err, rc);
            }
            r.end += @intCast(rc);
            return 0;
        }
    }

    fn failRead(self: *ServerStream, err: ReadError, rc: c_int) error{ReadFailed} {
        self.read_err = err;
        self.read_code = rc;
        return error.ReadFailed;
    }

    fn level(self: *const ServerStream, cause: ?anyerror) std.log.Level {
        return causeLevel(self.peer_closed, cause);
    }

    fn writerDrain(w: *Writer, data: []const []const u8, splat: usize) Writer.Error!usize {
        const self: *ServerStream = @alignCast(@fieldParentPtr("writer_interface", w));

        const buffered = w.buffered();
        if (buffered.len != 0) return w.consume(try self.sslWrite(buffered));

        for (data[0 .. data.len - 1]) |chunk| {
            if (chunk.len == 0) continue;
            return w.consume(try self.sslWrite(chunk));
        }

        const last = data[data.len - 1];
        if (last.len == 0) return w.consume(0);
        var written: usize = 0;
        for (0..splat) |_| {
            const n = try self.sslWrite(last);
            written += n;
            if (n < last.len) break;
        }
        return w.consume(written);
    }

    fn sslWrite(self: *ServerStream, bytes: []const u8) Writer.Error!usize {
        while (true) {
            // A retry after WANT_WRITE must repeat the same arguments.
            const rc = mbedtls_ssl_write(self.ssl.ptr, bytes.ptr, bytes.len);
            if (isRetry(rc)) continue;
            if (rc < 0) {
                self.write_code = rc;
                report("ssl_write", rc, self.level(self.send_cause));
                return error.WriteFailed;
            }
            return @intCast(rc);
        }
    }

    fn bioSend(bio: ?*anyopaque, buf: [*]const u8, len: usize) callconv(.c) c_int {
        const self: *ServerStream = @ptrCast(@alignCast(bio.?));
        const w = &self.net_writer.interface;
        w.writeAll(buf[0..len]) catch return self.sendFailed();
        w.flush() catch return self.sendFailed();
        return @intCast(len);
    }

    /// Mbed TLS only ever sees `err_net_send_failed`, so the concrete cause is
    /// kept here for the caller and for the log level.
    fn sendFailed(self: *ServerStream) c_int {
        if (self.net_writer.err) |cause| self.send_cause = cause;
        return err_net_send_failed;
    }

    fn recvFailed(self: *ServerStream) c_int {
        if (self.net_reader.err) |cause| self.recv_cause = cause;
        return err_net_recv_failed;
    }

    fn bioRecv(bio: ?*anyopaque, buf: [*]u8, len: usize) callconv(.c) c_int {
        const self: *ServerStream = @ptrCast(@alignCast(bio.?));
        if (len == 0) return 0;
        var data: [1][]u8 = .{buf[0..len]};
        const n = self.net_reader.interface.readVec(&data) catch |err| switch (err) {
            error.EndOfStream => return 0,
            error.ReadFailed => return self.recvFailed(),
        };
        // A zero-length transport buffer makes a short read impossible, but the
        // interface permits it; ask Mbed TLS to come back rather than reporting EOF.
        if (n == 0) return err_want_read;
        return @intCast(n);
    }
};

/// Heap storage for one Mbed TLS context whose layout stays opaque to Zig.
fn Block(comptime T: type) type {
    return struct {
        ptr: *T,
        bytes: []align(context_alignment) u8,

        const Self = @This();

        fn alloc(gpa: std.mem.Allocator, size: usize) error{OutOfMemory}!Self {
            const bytes = try gpa.alignedAlloc(u8, ContextAlignment, size);
            return .{ .ptr = @ptrCast(bytes.ptr), .bytes = bytes };
        }

        fn free(self: Self, gpa: std.mem.Allocator) void {
            gpa.free(self.bytes);
        }
    };
}

/// Codes that mean "call me again with the same arguments".
///
/// MBEDTLS_ERR_SSL_{CRYPTO,ASYNC}_IN_PROGRESS are deliberately absent: nxdns
/// configures no async or opaque crypto, so those codes can only mean the
/// configuration is wrong. Retrying them would spin.
fn isRetry(rc: c_int) bool {
    return rc == err_want_read or rc == err_want_write;
}

fn check(rc: c_int, comptime op: []const u8, comptime failure: anytype) @TypeOf(failure)!void {
    if (rc == 0) return;
    // Setup and configuration failures are nxdns's own; no peer can cause them.
    report(op, rc, .warn);
    return failure;
}

fn report(comptime op: []const u8, rc: c_int, level: std.log.Level) void {
    var text: [160]u8 = undefined;
    mbedtls_strerror(rc, &text, text.len);
    const args = .{ op, std.mem.sliceTo(&text, 0), rc };
    switch (level) {
        .err => log.err("mbedtls_{s} failed: {s} ({d})", args),
        .warn => log.warn("mbedtls_{s} failed: {s} ({d})", args),
        .info => log.info("mbedtls_{s} failed: {s} ({d})", args),
        .debug => log.debug("mbedtls_{s} failed: {s} ({d})", args),
    }
}

/// A transport failure the peer or an nxdns shutdown produced, as opposed to a
/// local resource or configuration failure. The names come from
/// `std.Io.net.Stream.Reader.Error` and `.Writer.Error`.
fn isPeerCause(cause: ?anyerror) bool {
    const concrete = cause orelse return false;
    return switch (concrete) {
        error.Canceled,
        error.ConnectionResetByPeer,
        error.ConnectionRefused,
        error.Timeout,
        error.SocketUnconnected,
        error.HostUnreachable,
        error.NetworkUnreachable,
        => true,
        else => false,
    };
}

/// Peer-driven and shutdown-driven failures log at debug, for the reason the
/// truncation path already records: any client can reach these sites at will,
/// so a louder level is a log-spam vector. Local and configuration failures —
/// which no peer can provoke — keep warn.
fn causeLevel(peer_closed: bool, cause: ?anyerror) std.log.Level {
    if (peer_closed) return .debug;
    return if (isPeerCause(cause)) .debug else .warn;
}

/// A canceled transport read is nxdns closing the connection, not a TLS fault,
/// so it keeps its own error instead of collapsing into `TlsFailed`.
fn readErrorFor(cause: ?anyerror) ServerStream.ReadError {
    const concrete = cause orelse return error.TlsFailed;
    return if (concrete == error.Canceled) error.Canceled else error.TlsFailed;
}

// -- Mbed TLS 3.6.7 surface ------------------------------------------------
// Signatures transcribed from include/mbedtls/{ssl,x509_crt,pk,entropy,ctr_drbg,error}.h
// and include/psa/crypto.h of the pinned release.

const SslContext = opaque {};
const SslConfig = opaque {};
const X509Crt = opaque {};
const PkContext = opaque {};
const EntropyContext = opaque {};
const CtrDrbgContext = opaque {};

/// `mbedtls_f_rng_t` (platform_util.h)
const RngFn = fn (p_rng: ?*anyopaque, output: [*]u8, output_len: usize) callconv(.c) c_int;
/// `mbedtls_ssl_send_t` (ssl.h)
const SendFn = fn (ctx: ?*anyopaque, buf: [*]const u8, len: usize) callconv(.c) c_int;
/// `mbedtls_ssl_recv_t` (ssl.h)
const RecvFn = fn (ctx: ?*anyopaque, buf: [*]u8, len: usize) callconv(.c) c_int;
/// `mbedtls_ssl_recv_timeout_t` (ssl.h)
const RecvTimeoutFn = fn (ctx: ?*anyopaque, buf: [*]u8, len: usize, timeout: u32) callconv(.c) c_int;

const ssl_transport_stream: c_int = 0;
const ssl_is_server: c_int = 1;
const ssl_verify_none: c_int = 0;
const ssl_preset_default: c_int = 0;

const err_conn_eof: c_int = -0x7280;
const err_peer_close_notify: c_int = -0x7880;
const err_want_read: c_int = -0x6900;
const err_want_write: c_int = -0x6880;
const err_net_recv_failed: c_int = -0x004C;
const err_net_send_failed: c_int = -0x004E;

extern fn mbedtls_ssl_init(ssl: *SslContext) void;
extern fn mbedtls_ssl_free(ssl: *SslContext) void;
extern fn mbedtls_ssl_setup(ssl: *SslContext, conf: *const SslConfig) c_int;
extern fn mbedtls_ssl_handshake(ssl: *SslContext) c_int;
extern fn mbedtls_ssl_read(ssl: *SslContext, buf: [*]u8, len: usize) c_int;
extern fn mbedtls_ssl_write(ssl: *SslContext, buf: [*]const u8, len: usize) c_int;
extern fn mbedtls_ssl_close_notify(ssl: *SslContext) c_int;
extern fn mbedtls_ssl_set_bio(
    ssl: *SslContext,
    p_bio: ?*anyopaque,
    f_send: ?*const SendFn,
    f_recv: ?*const RecvFn,
    f_recv_timeout: ?*const RecvTimeoutFn,
) void;

extern fn mbedtls_ssl_config_init(conf: *SslConfig) void;
extern fn mbedtls_ssl_config_free(conf: *SslConfig) void;
extern fn mbedtls_ssl_config_defaults(
    conf: *SslConfig,
    endpoint: c_int,
    transport: c_int,
    preset: c_int,
) c_int;
extern fn mbedtls_ssl_conf_rng(conf: *SslConfig, f_rng: *const RngFn, p_rng: ?*anyopaque) void;
extern fn mbedtls_ssl_conf_authmode(conf: *SslConfig, authmode: c_int) void;
extern fn mbedtls_ssl_conf_own_cert(conf: *SslConfig, own_cert: *X509Crt, pk_key: *PkContext) c_int;
/// Records the pointer to `protos` (NULL-terminated, decreasing preference);
/// the list must outlive the configuration.
extern fn mbedtls_ssl_conf_alpn_protocols(conf: *SslConfig, protos: [*:null]const ?[*:0]const u8) c_int;

extern fn mbedtls_x509_crt_init(crt: *X509Crt) void;
extern fn mbedtls_x509_crt_free(crt: *X509Crt) void;
/// For PEM input, `buflen` must count the terminating zero byte.
extern fn mbedtls_x509_crt_parse(chain: *X509Crt, buf: [*]const u8, buflen: usize) c_int;

extern fn mbedtls_pk_init(ctx: *PkContext) void;
extern fn mbedtls_pk_free(ctx: *PkContext) void;
/// For PEM input, `keylen` must count the terminating zero byte.
extern fn mbedtls_pk_parse_key(
    ctx: *PkContext,
    key: [*]const u8,
    keylen: usize,
    pwd: ?[*]const u8,
    pwdlen: usize,
    f_rng: ?*const RngFn,
    p_rng: ?*anyopaque,
) c_int;

/// `f_rng` must not be NULL. Returns 0 when the pair matches.
extern fn mbedtls_pk_check_pair(
    pub_key: *const PkContext,
    prv: *const PkContext,
    f_rng: *const RngFn,
    p_rng: ?*anyopaque,
) c_int;

extern fn mbedtls_entropy_init(ctx: *EntropyContext) void;
extern fn mbedtls_entropy_free(ctx: *EntropyContext) void;
extern fn mbedtls_entropy_func(data: ?*anyopaque, output: [*]u8, len: usize) callconv(.c) c_int;

extern fn mbedtls_ctr_drbg_init(ctx: *CtrDrbgContext) void;
extern fn mbedtls_ctr_drbg_free(ctx: *CtrDrbgContext) void;
extern fn mbedtls_ctr_drbg_seed(
    ctx: *CtrDrbgContext,
    f_entropy: *const fn (?*anyopaque, [*]u8, usize) callconv(.c) c_int,
    p_entropy: ?*anyopaque,
    custom: ?[*]const u8,
    len: usize,
) c_int;
extern fn mbedtls_ctr_drbg_random(p_rng: ?*anyopaque, output: [*]u8, output_len: usize) callconv(.c) c_int;

extern fn mbedtls_strerror(errnum: c_int, buffer: [*]u8, buflen: usize) void;

/// `psa_status_t psa_crypto_init(void)` — TLS 1.3 runs its key schedule through PSA.
extern fn psa_crypto_init() i32;

// -- src/platform/mbedtls_shim.c -------------------------------------------

extern fn nx_sizeof_ssl_context() usize;
extern fn nx_sizeof_ssl_config() usize;
extern fn nx_sizeof_x509_crt() usize;
extern fn nx_sizeof_pk_context() usize;
extern fn nx_sizeof_entropy_context() usize;
extern fn nx_sizeof_ctr_drbg_context() usize;
extern fn nx_max_context_alignment() usize;
/// The public key of the first certificate in `crt`.
extern fn nx_x509_crt_pk(crt: *X509Crt) *PkContext;

extern fn nx_const_ssl_transport_stream() c_int;
extern fn nx_const_ssl_is_server() c_int;
extern fn nx_const_ssl_verify_none() c_int;
extern fn nx_const_ssl_preset_default() c_int;
extern fn nx_const_err_ssl_conn_eof() c_int;
extern fn nx_const_err_ssl_peer_close_notify() c_int;
extern fn nx_const_err_ssl_want_read() c_int;
extern fn nx_const_err_ssl_want_write() c_int;
extern fn nx_const_err_net_recv_failed() c_int;
extern fn nx_const_err_net_send_failed() c_int;

/// Every transcribed value beside the header macro it came from. The sizes and
/// alignments have always been read from the shim rather than transcribed;
/// these ten were the exception, and a renumbering after a version bump would
/// change behaviour with nothing to say so — a moved `close_notify` inverts the
/// truncation detection in `read` and `close`.
fn transcriptionsAgree() bool {
    return ssl_transport_stream == nx_const_ssl_transport_stream() and
        ssl_is_server == nx_const_ssl_is_server() and
        ssl_verify_none == nx_const_ssl_verify_none() and
        ssl_preset_default == nx_const_ssl_preset_default() and
        err_conn_eof == nx_const_err_ssl_conn_eof() and
        err_peer_close_notify == nx_const_err_ssl_peer_close_notify() and
        err_want_read == nx_const_err_ssl_want_read() and
        err_want_write == nx_const_err_ssl_want_write() and
        err_net_recv_failed == nx_const_err_net_recv_failed() and
        err_net_send_failed == nx_const_err_net_send_failed();
}

// -- tests -----------------------------------------------------------------

test "ServerContext.init accepts the fixture cert and key" {
    const fixtures = @import("test_fixtures");
    const gpa = std.testing.allocator;

    var ctx = try ServerContext.init(gpa, fixtures.cert_pem, fixtures.key_pem, null);
    defer ctx.deinit(gpa);
}

test "ServerContext.init rejects a truncated certificate" {
    const fixtures = @import("test_fixtures");
    const gpa = std.testing.allocator;

    const truncated = try truncate(gpa, fixtures.cert_pem);
    defer gpa.free(truncated);

    try std.testing.expectError(
        error.CertParse,
        ServerContext.init(gpa, truncated, fixtures.key_pem, null),
    );
}

test "ServerContext.init rejects a truncated private key" {
    const fixtures = @import("test_fixtures");
    const gpa = std.testing.allocator;

    const truncated = try truncate(gpa, fixtures.key_pem);
    defer gpa.free(truncated);

    try std.testing.expectError(
        error.KeyParse,
        ServerContext.init(gpa, fixtures.cert_pem, truncated, null),
    );
}

test "ServerContext.init rejects a key that is not the certificate's" {
    const fixtures = @import("test_fixtures");
    const gpa = std.testing.allocator;

    try std.testing.expectError(
        error.KeyMismatch,
        ServerContext.init(gpa, fixtures.cert_pem, fixtures.mismatched_key_pem, null),
    );
}

test "ServerContext.init accepts a NULL-terminated ALPN protocol list" {
    const fixtures = @import("test_fixtures");
    const gpa = std.testing.allocator;

    const protocols: [*:null]const ?[*:0]const u8 = &.{ "http/1.1", "dot" };
    var ctx = try ServerContext.init(gpa, fixtures.cert_pem, fixtures.key_pem, protocols);
    defer ctx.deinit(gpa);
}

test "the shim agrees with the alignment contexts are allocated at" {
    try std.testing.expect(nx_max_context_alignment() <= context_alignment);
    try std.testing.expect(nx_sizeof_ssl_context() > 0);
    try std.testing.expect(nx_sizeof_ssl_config() > 0);
    try std.testing.expect(nx_sizeof_x509_crt() > 0);
    try std.testing.expect(nx_sizeof_pk_context() > 0);
    try std.testing.expect(nx_sizeof_entropy_context() > 0);
    try std.testing.expect(nx_sizeof_ctr_drbg_context() > 0);
}

test "the shim agrees with every transcribed mbedtls constant" {
    try std.testing.expect(transcriptionsAgree());

    // Value by value, so a failure names the one that drifted instead of
    // reporting that something did.
    try std.testing.expectEqual(ssl_transport_stream, nx_const_ssl_transport_stream());
    try std.testing.expectEqual(ssl_is_server, nx_const_ssl_is_server());
    try std.testing.expectEqual(ssl_verify_none, nx_const_ssl_verify_none());
    try std.testing.expectEqual(ssl_preset_default, nx_const_ssl_preset_default());
    try std.testing.expectEqual(err_conn_eof, nx_const_err_ssl_conn_eof());
    try std.testing.expectEqual(err_peer_close_notify, nx_const_err_ssl_peer_close_notify());
    try std.testing.expectEqual(err_want_read, nx_const_err_ssl_want_read());
    try std.testing.expectEqual(err_want_write, nx_const_err_ssl_want_write());
    try std.testing.expectEqual(err_net_recv_failed, nx_const_err_net_recv_failed());
    try std.testing.expectEqual(err_net_send_failed, nx_const_err_net_send_failed());
}

test "readErrorFor keeps a canceled read out of the TLS failure bucket" {
    try std.testing.expectEqual(ServerStream.ReadError.Canceled, readErrorFor(error.Canceled));
    try std.testing.expectEqual(ServerStream.ReadError.TlsFailed, readErrorFor(error.ConnectionResetByPeer));
    try std.testing.expectEqual(ServerStream.ReadError.TlsFailed, readErrorFor(error.SystemResources));
    try std.testing.expectEqual(ServerStream.ReadError.TlsFailed, readErrorFor(null));
}

test "causeLevel logs peer misbehavior at debug and local faults at warn" {
    // Ruling 16: close_notify against a socket the peer already dropped.
    try std.testing.expectEqual(std.log.Level.debug, causeLevel(true, null));
    try std.testing.expectEqual(std.log.Level.debug, causeLevel(false, error.ConnectionResetByPeer));
    try std.testing.expectEqual(std.log.Level.debug, causeLevel(false, error.Canceled));
    try std.testing.expectEqual(std.log.Level.warn, causeLevel(false, error.SystemResources));
    // No stashed cause means Mbed TLS rejected the record itself.
    try std.testing.expectEqual(std.log.Level.warn, causeLevel(false, null));
}

test "isPeerCause separates peer and shutdown failures from local ones" {
    try std.testing.expect(isPeerCause(error.Canceled));
    try std.testing.expect(isPeerCause(error.ConnectionResetByPeer));
    try std.testing.expect(isPeerCause(error.Timeout));
    try std.testing.expect(isPeerCause(error.SocketUnconnected));
    try std.testing.expect(!isPeerCause(error.SystemResources));
    try std.testing.expect(!isPeerCause(error.AccessDenied));
    try std.testing.expect(!isPeerCause(error.NetworkDown));
    try std.testing.expect(!isPeerCause(null));
}

/// Drops the second half of a PEM document, keeping the header intact so the
/// parser fails on the body rather than on a missing "-----BEGIN" line.
fn truncate(gpa: std.mem.Allocator, pem: [:0]const u8) ![:0]const u8 {
    return gpa.dupeZ(u8, pem[0 .. pem.len / 2]);
}

const loopback_message = "nxdns loopback probe";

test "loopback echo between the mbedtls server and std.crypto.tls.Client" {
    const build_options = @import("build_options");
    if (!build_options.integration) return error.SkipZigTest;

    const fixtures = @import("test_fixtures");
    const gpa = std.testing.allocator;

    var threaded: Io.Threaded = .init(gpa, .{});
    defer threaded.deinit();
    const io = threaded.io();

    var ctx = try ServerContext.init(gpa, fixtures.cert_pem, fixtures.key_pem, null);
    defer ctx.deinit(gpa);

    const listen_address: Io.net.IpAddress = .{ .ip4 = .loopback(0) };
    var server = try listen_address.listen(io, .{ .reuse_address = true });
    defer server.deinit(io);

    var server_task = try io.concurrent(echoOnce, .{ gpa, &ctx, io, &server });
    const client_result = runEchoClient(io, server.socket.address);
    // A client that never connects would leave the server blocked in `accept`.
    const server_result = if (client_result) |_|
        server_task.await(io)
    else |_|
        server_task.cancel(io);

    try client_result;
    try server_result;
}

// std.crypto.tls.Client in Zig 0.16.0 cannot send the ALPN extension
// (Client.Options has no such field), so negotiation itself is untestable
// here; this proves the ruling that a client offering no ALPN still
// completes the handshake against a context advertising a list.
test "loopback echo succeeds against a context advertising ALPN" {
    const build_options = @import("build_options");
    if (!build_options.integration) return error.SkipZigTest;

    const fixtures = @import("test_fixtures");
    const gpa = std.testing.allocator;

    var threaded: Io.Threaded = .init(gpa, .{});
    defer threaded.deinit();
    const io = threaded.io();

    const protocols: [*:null]const ?[*:0]const u8 = &.{"http/1.1"};
    var ctx = try ServerContext.init(gpa, fixtures.cert_pem, fixtures.key_pem, protocols);
    defer ctx.deinit(gpa);

    const listen_address: Io.net.IpAddress = .{ .ip4 = .loopback(0) };
    var server = try listen_address.listen(io, .{ .reuse_address = true });
    defer server.deinit(io);

    var server_task = try io.concurrent(echoOnce, .{ gpa, &ctx, io, &server });
    const client_result = runEchoClient(io, server.socket.address);
    const server_result = if (client_result) |_|
        server_task.await(io)
    else |_|
        server_task.cancel(io);

    try client_result;
    try server_result;
}

test "a transport EOF without close_notify reads as a truncated stream" {
    const build_options = @import("build_options");
    if (!build_options.integration) return error.SkipZigTest;

    const fixtures = @import("test_fixtures");
    const gpa = std.testing.allocator;

    var threaded: Io.Threaded = .init(gpa, .{});
    defer threaded.deinit();
    const io = threaded.io();

    var ctx = try ServerContext.init(gpa, fixtures.cert_pem, fixtures.key_pem, null);
    defer ctx.deinit(gpa);

    const listen_address: Io.net.IpAddress = .{ .ip4 = .loopback(0) };
    var server = try listen_address.listen(io, .{ .reuse_address = true });
    defer server.deinit(io);

    var server_task = try io.concurrent(expectTruncation, .{ gpa, &ctx, io, &server });
    const client_result = runTruncatingClient(io, server.socket.address);
    const server_result = if (client_result) |_|
        server_task.await(io)
    else |_|
        server_task.cancel(io);

    try client_result;
    try server_result;
}

test "a canceled read stashes Canceled, not TlsFailed" {
    const build_options = @import("build_options");
    if (!build_options.integration) return error.SkipZigTest;

    const fixtures = @import("test_fixtures");
    const gpa = std.testing.allocator;

    var threaded: Io.Threaded = .init(gpa, .{});
    defer threaded.deinit();
    const io = threaded.io();

    var ctx = try ServerContext.init(gpa, fixtures.cert_pem, fixtures.key_pem, null);
    defer ctx.deinit(gpa);

    const listen_address: Io.net.IpAddress = .{ .ip4 = .loopback(0) };
    var server = try listen_address.listen(io, .{ .reuse_address = true });
    defer server.deinit(io);

    // Set once the server task is about to block in `readIntoBuffer`, so the
    // cancellation below lands on the read and not on the handshake.
    var reading: Io.Event = .unset;
    var server_task = try io.concurrent(expectCanceledRead, .{ gpa, &ctx, io, &server, &reading });

    var client = runIdleClient(io, server.socket.address) catch |err| {
        server_task.cancel(io) catch {};
        return err;
    };
    defer client.close(io);

    reading.waitUncancelable(io);
    try server_task.cancel(io);
}

/// Handshakes, then blocks in a plaintext read that only the cancel can end.
fn expectCanceledRead(
    gpa: std.mem.Allocator,
    ctx: *ServerContext,
    io: Io,
    server: *Io.net.Server,
    reading: *Io.Event,
) anyerror!void {
    var stream = try server.accept(io);
    defer stream.close(io);

    var read_buffer: [4096]u8 = undefined;
    var write_buffer: [4096]u8 = undefined;
    var tls: ServerStream = undefined;
    try tls.accept(gpa, ctx, io, &stream, &read_buffer, &write_buffer);
    defer tls.close(gpa);

    reading.set(io);

    var byte: [1]u8 = undefined;
    try std.testing.expectError(error.ReadFailed, tls.reader().readSliceAll(&byte));
    try std.testing.expectEqual(ServerStream.ReadError.Canceled, tls.read_err.?);
    try std.testing.expectEqual(@as(?anyerror, error.Canceled), tls.recv_cause);
    // A cancel is nxdns's own doing, never the peer ending the stream.
    try std.testing.expect(!tls.peer_closed);
}

/// A client that completes the handshake and then sends nothing. The caller
/// keeps it alive so the server's read has no other way to end.
const IdleClient = struct {
    stream: Io.net.Stream,

    fn close(self: *IdleClient, io: Io) void {
        self.stream.close(io);
    }
};

fn runIdleClient(io: Io, address: Io.net.IpAddress) !IdleClient {
    const tls = std.crypto.tls;

    var stream = try address.connect(io, .{ .mode = .stream });
    errdefer stream.close(io);

    // The handshake buffers die with this function; the connection outlives it
    // because nothing more is ever read from or written to it.
    var transport_read_buffer: [tls.Client.min_buffer_len]u8 = undefined;
    var transport_write_buffer: [tls.Client.min_buffer_len]u8 = undefined;
    var net_reader = stream.reader(io, &transport_read_buffer);
    var net_writer = stream.writer(io, &transport_write_buffer);

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

    var plaintext_read_buffer: [4096]u8 = undefined;
    var plaintext_write_buffer: [4096]u8 = undefined;

    var client = try tls.Client.init(&net_reader.interface, &net_writer.interface, .{
        .host = .no_verification,
        .ca = .no_verification,
        .read_buffer = &plaintext_read_buffer,
        .write_buffer = &plaintext_write_buffer,
        .entropy = &entropy,
        .realtime_now = Io.Timestamp.now(io, .real),
    });
    _ = &client;
    try net_writer.interface.flush();

    return .{ .stream = stream };
}

fn expectTruncation(
    gpa: std.mem.Allocator,
    ctx: *ServerContext,
    io: Io,
    server: *Io.net.Server,
) anyerror!void {
    var stream = try server.accept(io);
    defer stream.close(io);

    var read_buffer: [4096]u8 = undefined;
    var write_buffer: [4096]u8 = undefined;
    var tls: ServerStream = undefined;
    try tls.accept(gpa, ctx, io, &stream, &read_buffer, &write_buffer);
    defer tls.close(gpa);

    var received: [loopback_message.len]u8 = undefined;
    try tls.reader().readSliceAll(&received);
    try std.testing.expectEqualStrings(loopback_message, &received);

    var tail: [1]u8 = undefined;
    try std.testing.expectError(error.ReadFailed, tls.reader().readSliceAll(&tail));
    try std.testing.expectEqual(ServerStream.ReadError.TlsConnectionTruncated, tls.read_err.?);
    try std.testing.expect(!tls.peer_closed);
}

/// Writes one message, then drops the TCP connection without close_notify.
fn runTruncatingClient(io: Io, address: Io.net.IpAddress) anyerror!void {
    const tls = std.crypto.tls;

    var stream = try address.connect(io, .{ .mode = .stream });
    defer stream.close(io);

    var transport_read_buffer: [tls.Client.min_buffer_len]u8 = undefined;
    var transport_write_buffer: [tls.Client.min_buffer_len]u8 = undefined;
    var net_reader = stream.reader(io, &transport_read_buffer);
    var net_writer = stream.writer(io, &transport_write_buffer);

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

    var plaintext_read_buffer: [4096]u8 = undefined;
    var plaintext_write_buffer: [4096]u8 = undefined;

    var client = try tls.Client.init(&net_reader.interface, &net_writer.interface, .{
        .host = .no_verification,
        .ca = .no_verification,
        .read_buffer = &plaintext_read_buffer,
        .write_buffer = &plaintext_write_buffer,
        .entropy = &entropy,
        .realtime_now = Io.Timestamp.now(io, .real),
    });
    try net_writer.interface.flush();

    try client.writer.writeAll(loopback_message);
    try client.writer.flush();
    try net_writer.interface.flush();
}

fn echoOnce(
    gpa: std.mem.Allocator,
    ctx: *ServerContext,
    io: Io,
    server: *Io.net.Server,
) anyerror!void {
    var stream = try server.accept(io);
    defer stream.close(io);

    var read_buffer: [4096]u8 = undefined;
    var write_buffer: [4096]u8 = undefined;
    var tls: ServerStream = undefined;
    try tls.accept(gpa, ctx, io, &stream, &read_buffer, &write_buffer);
    defer tls.close(gpa);

    var received: [loopback_message.len]u8 = undefined;
    try tls.reader().readSliceAll(&received);
    try tls.writer().writeAll(&received);
    try tls.writer().flush();
}

fn runEchoClient(io: Io, address: Io.net.IpAddress) anyerror!void {
    const tls = std.crypto.tls;

    var stream = try address.connect(io, .{ .mode = .stream });
    defer stream.close(io);

    var transport_read_buffer: [tls.Client.min_buffer_len]u8 = undefined;
    var transport_write_buffer: [tls.Client.min_buffer_len]u8 = undefined;
    var net_reader = stream.reader(io, &transport_read_buffer);
    var net_writer = stream.writer(io, &transport_write_buffer);

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

    var plaintext_read_buffer: [4096]u8 = undefined;
    var plaintext_write_buffer: [4096]u8 = undefined;

    var client = try tls.Client.init(&net_reader.interface, &net_writer.interface, .{
        .host = .no_verification,
        .ca = .no_verification,
        .read_buffer = &plaintext_read_buffer,
        .write_buffer = &plaintext_write_buffer,
        .entropy = &entropy,
        .realtime_now = Io.Timestamp.now(io, .real),
    });
    try net_writer.interface.flush();

    try client.writer.writeAll(loopback_message);
    try client.writer.flush();
    try net_writer.interface.flush();

    var echoed: [loopback_message.len]u8 = undefined;
    try client.reader.readSliceAll(&echoed);
    try std.testing.expectEqualStrings(loopback_message, &echoed);

    try client.end();
    try net_writer.interface.flush();

    // The server's close_notify must arrive as a clean end of stream.
    var tail: [1]u8 = undefined;
    try std.testing.expectError(error.EndOfStream, client.reader.readSliceAll(&tail));
}
