//! The listener core the four stream listeners share.
//!
//! `tcp_server.zig`, `dot_server.zig`, `doh_server.zig` and `web/server.zig`
//! are the same machine wearing four transports: a fixed pre-allocated slot
//! pool, a claim rule where shutdown outranks capacity, an accept loop with one
//! error mapping, a mutex-ordered close dance, and — for everything that waits
//! on a peer — one select race against a budget. Milestone 18 ruling 1 puts
//! that machine here once, so a fix to it lands once.
//!
//! What stays outside: the per-connection serve function, the connection
//! payload (buffers, TLS context, arenas), and the TLS lifecycle. A TLS
//! listener's certificate pin, handshake, close_notify and release form one
//! ordered sequence that the Core has no business owning; what it does own is
//! `handshakeStage`, the exactly-once `handshook` flag whose absence was a real
//! leak in doh_server before the milestone-10 review hand-ported the fix.
//!
//! Shutdown is the one part worth reading twice. `deinit` publishes `.closing`,
//! shuts the listening socket down (which unblocks `accept` with
//! `error.SocketNotListening`) and shuts every `.active` connection down in the
//! same critical section that closes the door on new ones; `serve` then drains
//! its connection group so a half-written reply still goes out whole. A
//! *canceled* `serve` cannot drain, because a keep-alive peer has no deadline
//! of its own and one chatty client would stall the whole process's shutdown;
//! it cancels the group instead, at the cost of the one reply mid-write.
//! Either way `serve` returns only once no task can still touch a slot.

const std = @import("std");
const net = std.Io.net;
const Allocator = std.mem.Allocator;

const transport = @import("../upstream/transport.zig");

/// Lifecycle of the accept loop. `serve` claims `.serving`, `deinit` publishes
/// `.closing`, and the two meet at `stopped` so no task touches a connection
/// slot after it is freed.
pub const State = enum(u32) { idle, serving, closing };

/// `.closing` exists so `deinit` never shuts down a descriptor that its own
/// task is about to close: the transition to `.closing` happens under the mutex
/// before the close, and `deinit` only touches `.active` slots.
pub const ConnState = enum { free, active, closing };

/// Why the accept loop stopped, which decides what happens to the connections
/// still in flight.
pub const Stop = enum {
    /// `deinit` published `.closing`. It has already shut every live connection
    /// down, so each one is unblocked and finishing on its own.
    closing,
    /// This task is being canceled. Nothing has touched the connections.
    canceled,
};

/// What the accept loop does with a stream it has just accepted.
pub const Claim = union(enum) {
    /// The stream owns `conns[index]`.
    slot: usize,
    /// Every slot is taken. The stream is refused and the loop continues.
    at_capacity,
    /// `deinit` has started. The stream is closed and the loop returns.
    shutting_down,
};

/// How long the accept loop waits after an unexpected accept failure, so a
/// persistent one cannot turn the loop into a spin.
pub const retry_delay: std.Io.Clock.Duration = .{ .raw = .fromMilliseconds(100), .clock = .awake };

pub fn bump(counter: *std.atomic.Value(u64)) void {
    _ = counter.fetchAdd(1, .monotonic);
}

/// The counters every listener keeps. Listener-specific ones —
/// `tls_handshake_failures`, `bad_requests`, `requests` — live beside this in
/// the owning listener, and each listener's exported `Snapshot` stays a flat
/// hand-written struct so `/metrics` output does not depend on this layout.
///
/// `idle_timeouts` is bumped by the three DNS listeners; the web listener has
/// no idle race (its port is LAN-facing and the cancel path bounds shutdown),
/// so its copy stays zero and it exports no family at all.
pub const CoreStats = struct {
    connections: std.atomic.Value(u64) = .init(0),
    rejected_at_capacity: std.atomic.Value(u64) = .init(0),
    rejected_at_shutdown: std.atomic.Value(u64) = .init(0),
    accept_errors: std.atomic.Value(u64) = .init(0),
    connection_errors: std.atomic.Value(u64) = .init(0),
    idle_timeouts: std.atomic.Value(u64) = .init(0),
};

// ---------------------------------------------------------------------------
// the race harness
// ---------------------------------------------------------------------------

pub const Outcome = union(enum) {
    op: anyerror!void,
    expiry: std.Io.Cancelable!void,
};

pub const Result = enum { ok, timed_out, failed, canceled };

/// Runs one connection operation against a budget and cancels the loser. No
/// stream read or write in 0.16.0 accepts a timeout, so every wait on a peer
/// that owes nxdns bytes goes through here.
pub fn race(
    io: std.Io,
    budget: std.Io.Clock.Duration,
    comptime f: anytype,
    args: std.meta.ArgsTuple(@TypeOf(f)),
) Result {
    var outcomes: [2]Outcome = undefined;
    var select: std.Io.Select(Outcome) = .init(io, &outcomes);
    defer select.cancelDiscard();

    select.concurrent(.op, f, args) catch |err| switch (err) {
        error.ConcurrencyUnavailable => return .failed,
    };
    select.concurrent(.expiry, expire, .{ io, budget }) catch |err| switch (err) {
        error.ConcurrencyUnavailable => return .failed,
    };

    return switch (select.await() catch return .canceled) {
        .op => |result| if (result) |_| .ok else |err| switch (err) {
            error.Canceled => .canceled,
            else => .failed,
        },
        // A canceled sleep means this task is being torn down, not that the
        // peer went idle.
        .expiry => |result| if (result) |_| .timed_out else |_| .canceled,
    };
}

pub fn expire(io: std.Io, budget: std.Io.Clock.Duration) std.Io.Cancelable!void {
    return budget.sleep(io);
}

// ---------------------------------------------------------------------------
// the TLS handshake stage
// ---------------------------------------------------------------------------

/// Runs a TLS handshake under `budget` and owns the exactly-once cleanup of the
/// context it may have created.
///
/// `stage` is anything with `accept(self) anyerror!void` and `close(self) void`.
/// The select can report the expiry or the cancellation *after* `accept` has in
/// fact succeeded, so a context that exists must be closed on every losing path
/// and on no other: the flag below is written before the race joins its tasks,
/// which is what makes "exactly once" true. Getting this wrong leaks one
/// mbedTLS ssl context per stalled handshake, which is what doh_server did
/// until the milestone-10 review hand-ported dot_server's fix — the duplication
/// this helper exists to end.
///
/// The caller's required ordering, which stays in the caller because the
/// certificate pin and the plaintext close are the listener's own business:
///
/// 1. pin the certificate generation (`CertStore.acquire`, released on exit),
/// 2. call `handshakeStage`,
/// 3. on `.ok` only: serve the connection,
/// 4. close the TLS stream (close_notify + free the context),
/// 5. release the pin, then let the slot's `finish` close the TCP stream.
///
/// On any result other than `.ok` this function has already done step 4 for the
/// caller, and the caller must not repeat it.
pub fn handshakeStage(io: std.Io, budget: std.Io.Clock.Duration, stage: anytype) Result {
    const Stage = @TypeOf(stage);
    const run = struct {
        fn accept(s: Stage, handshook: *bool) anyerror!void {
            try s.accept();
            handshook.* = true;
        }
    }.accept;

    var handshook = false;
    const result = race(io, budget, run, .{ stage, &handshook });
    if (result != .ok and handshook) stage.close();
    return result;
}

// ---------------------------------------------------------------------------
// the framed-message helpers (RFC 1035 §4.2.2; tcp and dot)
// ---------------------------------------------------------------------------

/// `readSliceShort` rather than `readSliceAll`: a zero-length read is a client
/// that closed cleanly between messages, and only a partial prefix is an error.
pub fn readPrefix(reader: *std.Io.Reader, buf: *[transport.prefix_len]u8, out_len: *usize) anyerror!void {
    out_len.* = try reader.readSliceShort(buf);
}

pub fn readBody(reader: *std.Io.Reader, buf: []u8) anyerror!void {
    return reader.readSliceAll(buf);
}

pub fn writeReply(writer: *std.Io.Writer, prefix: *const [transport.prefix_len]u8, bytes: []const u8) anyerror!void {
    try writer.writeAll(prefix);
    try writer.writeAll(bytes);
    try writer.flush();
}

// ---------------------------------------------------------------------------
// the claim rule
// ---------------------------------------------------------------------------

/// The capacity rule, without the mutex, so it is testable without a backend.
/// `conns` is any slice whose element has a `state: ConnState`.
pub fn firstFree(conns: anytype) ?usize {
    for (conns, 0..) |*conn, index| {
        if (conn.state == .free) return index;
    }
    return null;
}

/// The whole claim rule, without the mutex. Shutdown outranks capacity: a free
/// slot is still refused once `deinit` has passed the connections.
pub fn decideClaim(conns: anytype, shutdown_begun: bool) Claim {
    if (shutdown_begun) return .shutting_down;
    const index = firstFree(conns) orelse return .at_capacity;
    return .{ .slot = index };
}

// ---------------------------------------------------------------------------
// the core
// ---------------------------------------------------------------------------

/// The slot pool, the accept loop and the shutdown protocol, parameterized over
/// the four things that genuinely differ between listeners.
///
/// `Cfg` declares:
///
/// - `Owner: type` — the listener struct that embeds this core in a field
///   named `core`. The accept loop recovers it with `@fieldParentPtr`, so an
///   owner must not move after `listen`.
/// - `ConnPayload: type` — the rest of one slot: message buffers, a TLS
///   context, a per-request arena. Never touched here.
/// - `serveConn: fn (*Owner, std.Io, usize) void` — one whole connection. The
///   core spawns it, and closes the slot when it returns.
/// - `read_buffer_len` / `write_buffer_len` — the stream staging buffers, which
///   every listener has and sizes differently.
/// - `log` — the owner's `std.log` scope, and `name` — the two or three letters
///   its messages already start with, so the log text does not change.
///
/// Optional, absent for most listeners:
///
/// - `refuse: fn (std.Io, net.Stream) void` — what an over-capacity accept does
///   with the stream. The default closes it, which is the only honest answer a
///   DNS listener can give; the web listener answers 503 first.
/// - `initPayload` / `deinitPayload` — for a payload that owns memory (the web
///   listener's per-connection arena). `initPayload` runs inside `listen`,
///   `deinitPayload` inside `deinit` after every connection task has joined.
pub fn Core(comptime Cfg: type) type {
    return struct {
        const Self = @This();

        gpa: Allocator,
        listener: net.Server,
        conns: []Conn,
        mutex: std.Io.Mutex,
        /// Guarded by `mutex`. `deinit` sets it in the same critical section
        /// that shuts the active connections down, so a stream that arrives
        /// after that scan can never claim a slot the scan will not visit again.
        shutdown_begun: bool,
        stats: CoreStats,
        run_state: std.atomic.Value(State),
        stopped: std.Io.Event,

        pub const Conn = struct {
            /// The stream staging buffers. For the plaintext listeners these
            /// feed the socket reader and writer; for the TLS ones they are the
            /// `ServerStream` plaintext buffers.
            read_buf: [Cfg.read_buffer_len]u8,
            write_buf: [Cfg.write_buffer_len]u8,
            payload: Cfg.ConnPayload,
            stream: net.Stream,
            /// The client, read off the accepted socket once at claim time:
            /// every message on this connection comes from the same peer.
            peer: net.IpAddress,
            /// Guarded by `Core.mutex`.
            state: ConnState,
        };

        pub const ListenError = net.IpAddress.ListenError || error{OutOfMemory};

        pub fn listen(
            gpa: Allocator,
            io: std.Io,
            listen_address: net.IpAddress,
            max_connections: u16,
        ) ListenError!Self {
            std.debug.assert(max_connections > 0);

            const conns = try gpa.alloc(Conn, max_connections);
            errdefer gpa.free(conns);
            for (conns) |*conn| {
                conn.state = .free;
                if (@hasDecl(Cfg, "initPayload")) Cfg.initPayload(&conn.payload, gpa);
            }

            const listener = try listen_address.listen(io, .{ .reuse_address = true });

            return .{
                .gpa = gpa,
                .listener = listener,
                .conns = conns,
                .mutex = .init,
                .shutdown_begun = false,
                .stats = .{},
                .run_state = .init(.idle),
                .stopped = .unset,
            };
        }

        pub fn deinit(self: *Self, io: std.Io) void {
            const was_serving = self.run_state.swap(.closing, .acq_rel) == .serving;

            // Shutting the listening socket down is the documented way to
            // unblock a pending `accept`: it fails with
            // `error.SocketNotListening`.
            const stream: net.Stream = .{ .socket = self.listener.socket };
            stream.shutdown(io, .both) catch |err| {
                Cfg.log.debug(Cfg.name ++ " listener shutdown failed: {t}", .{err});
            };

            // A live connection is blocked in a read that only its own budget
            // would end, which is seconds away. Shutting each one down bounds
            // this, and the same critical section closes the door on new ones.
            self.beginShutdown(io);

            if (was_serving) self.stopped.waitUncancelable(io);

            self.listener.deinit(io);
            if (@hasDecl(Cfg, "deinitPayload")) {
                for (self.conns) |*conn| Cfg.deinitPayload(&conn.payload);
            }
            self.gpa.free(self.conns);
            self.* = undefined;
        }

        /// The kernel-assigned address. A port of 0 in `listen` resolves here.
        pub fn boundAddress(self: *const Self) net.IpAddress {
            return self.listener.socket.address;
        }

        /// Accept loop. Returns when the task is canceled or `deinit` stops it.
        pub fn serve(self: *Self, io: std.Io) void {
            if (self.run_state.cmpxchgStrong(.idle, .serving, .acq_rel, .acquire) != null) return;

            var group: std.Io.Group = .init;
            switch (self.acceptLoop(io, &group)) {
                // `deinit` shut every live connection down before it published
                // `.closing`, so each one is already unblocked and ending on
                // its own. Awaiting them means a half-written reply still goes
                // out whole, and the wait is bounded by the shutdown, not the
                // peer.
                .closing => {
                    const prev = io.swapCancelProtection(.blocked);
                    group.await(io) catch |err| switch (err) {
                        error.Canceled => unreachable,
                    };
                    _ = io.swapCancelProtection(prev);
                },
                // Nothing has shut these connections down: `deinit` cannot run
                // until this task returns, and a peer may hold a connection
                // open indefinitely, so draining would let one client stall the
                // whole process's shutdown. `cancel` requests cancellation and
                // joins, so the slots are quiet — and the buffers still
                // unreferenced — by the time `serve` returns; the price is the
                // one reply that was mid-write.
                .canceled => group.cancel(io),
            }

            self.stopped.set(io);
        }

        /// The listener that embeds this core. Valid because the core is a
        /// field of it and neither may move after `listen`.
        pub fn owner(self: *Self) *Cfg.Owner {
            return @alignCast(@fieldParentPtr("core", self));
        }

        fn acceptLoop(self: *Self, io: std.Io, group: *std.Io.Group) Stop {
            while (self.run_state.load(.acquire) == .serving) {
                const stream = self.listener.accept(io) catch |err| switch (err) {
                    error.Canceled => return .canceled,
                    // `deinit` shuts the listening socket down to unblock
                    // exactly this call, so it is the shutdown path arriving
                    // early.
                    error.SocketNotListening => return .closing,
                    else => {
                        bump(&self.stats.accept_errors);
                        Cfg.log.debug(Cfg.name ++ " accept failed: {t}", .{err});
                        retry_delay.sleep(io) catch return .canceled;
                        continue;
                    },
                };

                const index = switch (self.claim(io, stream)) {
                    .slot => |index| index,
                    // Refusing now is honest; a queue would only hide the
                    // overload.
                    .at_capacity => {
                        bump(&self.stats.rejected_at_capacity);
                        if (@hasDecl(Cfg, "refuse")) Cfg.refuse(io, stream) else stream.close(io);
                        continue;
                    },
                    // `deinit` will not see this stream in any slot, so serving
                    // it would hold `deinit` for a whole idle budget.
                    .shutting_down => {
                        bump(&self.stats.rejected_at_shutdown);
                        stream.close(io);
                        return .closing;
                    },
                };

                group.concurrent(io, runConn, .{ self, io, index }) catch |err| switch (err) {
                    error.ConcurrencyUnavailable => {
                        bump(&self.stats.rejected_at_capacity);
                        self.finish(io, index);
                        continue;
                    },
                };

                bump(&self.stats.connections);
            }

            // The loop condition failed, which only `deinit` can cause.
            return .closing;
        }

        /// One connection task: the listener's own logic, then the slot close.
        /// Every early return inside `Cfg.serveConn` — and its own defers, such
        /// as a TLS close_notify — runs before the TCP stream is closed here.
        fn runConn(self: *Self, io: std.Io, index: usize) void {
            defer self.finish(io, index);
            Cfg.serveConn(self.owner(), io, index);
        }

        fn claim(self: *Self, io: std.Io, stream: net.Stream) Claim {
            // Uncancelable: this section takes no Io and never blocks on a
            // peer, so it cannot deadlock, and losing the lock mid-update would
            // leak a slot.
            self.mutex.lockUncancelable(io);
            defer self.mutex.unlock(io);

            const outcome = decideClaim(self.conns, self.shutdown_begun);
            switch (outcome) {
                .slot => |index| {
                    self.conns[index].stream = stream;
                    self.conns[index].peer = stream.socket.address;
                    self.conns[index].state = .active;
                },
                .at_capacity, .shutting_down => {},
            }
            return outcome;
        }

        fn finish(self: *Self, io: std.Io, index: usize) void {
            const conn = &self.conns[index];

            self.mutex.lockUncancelable(io);
            conn.state = .closing;
            self.mutex.unlock(io);

            // The socket is released even when this task is being torn down:
            // the next cancelable call would otherwise skip the close.
            const prev = io.swapCancelProtection(.blocked);
            conn.stream.close(io);
            _ = io.swapCancelProtection(prev);

            self.mutex.lockUncancelable(io);
            conn.state = .free;
            self.mutex.unlock(io);
        }

        /// Closes the door on new connections and unblocks the live ones. Both
        /// happen under one hold of the mutex: a `claim` that runs before this
        /// leaves an `.active` slot the loop below shuts down, and a `claim`
        /// that runs after it reads `shutdown_begun` and takes no slot at all.
        fn beginShutdown(self: *Self, io: std.Io) void {
            self.mutex.lockUncancelable(io);
            defer self.mutex.unlock(io);

            self.shutdown_begun = true;

            for (self.conns) |*conn| {
                if (conn.state != .active) continue;
                conn.stream.shutdown(io, .both) catch |err| {
                    Cfg.log.debug(Cfg.name ++ " connection shutdown failed: {t}", .{err});
                };
            }
        }
    };
}

// ---------------------------------------------------------------------------
// tests
// ---------------------------------------------------------------------------
//
// The claim rule is the whole of the shared state machine that can be tested
// without a backend, and it reads nothing but `state`, so these tests use a
// bare slot instead of instantiating a `Core`. They replace six copies that
// lived in tcp_server.zig and dot_server.zig and three more between
// doh_server.zig and web/server.zig.

const testing = std.testing;

const TestSlot = struct { state: ConnState };

fn testConns(count: usize) ![]TestSlot {
    const conns = try testing.allocator.alloc(TestSlot, count);
    for (conns) |*conn| conn.state = .free;
    return conns;
}

test "the connection pool hands out every slot once" {
    const conns = try testConns(3);
    defer testing.allocator.free(conns);

    for (0..conns.len) |expected| {
        const index = firstFree(conns) orelse return error.TestUnexpectedResult;
        try testing.expectEqual(expected, index);
        conns[index].state = .active;
    }
}

test "a full connection pool refuses instead of growing" {
    const conns = try testConns(2);
    defer testing.allocator.free(conns);

    for (conns) |*conn| conn.state = .active;
    try testing.expectEqual(@as(?usize, null), firstFree(conns));
    try testing.expectEqual(.at_capacity, std.meta.activeTag(decideClaim(conns, false)));
}

test "a closing slot is not reused until it is free" {
    const conns = try testConns(2);
    defer testing.allocator.free(conns);

    conns[0].state = .active;
    conns[1].state = .closing;
    try testing.expectEqual(@as(?usize, null), firstFree(conns));

    conns[1].state = .free;
    try testing.expectEqual(@as(?usize, 1), firstFree(conns));
    try testing.expectEqual(@as(usize, 1), decideClaim(conns, false).slot);
}

test "a claim takes the first free slot before shutdown" {
    const conns = try testConns(2);
    defer testing.allocator.free(conns);

    conns[0].state = .active;
    try testing.expectEqual(@as(usize, 1), decideClaim(conns, false).slot);
}

test "a claim after shutdown is refused even with a free slot" {
    const conns = try testConns(2);
    defer testing.allocator.free(conns);

    try testing.expectEqual(.shutting_down, std.meta.activeTag(decideClaim(conns, true)));

    // The refusal must not consume the slot: `deinit` frees it, nothing else.
    try testing.expectEqual(@as(?usize, 0), firstFree(conns));
}

test "shutdown outranks capacity" {
    const conns = try testConns(1);
    defer testing.allocator.free(conns);

    conns[0].state = .active;
    try testing.expectEqual(.at_capacity, std.meta.activeTag(decideClaim(conns, false)));
    try testing.expectEqual(.shutting_down, std.meta.activeTag(decideClaim(conns, true)));
}
