//! The TCP/53 listener.
//!
//! RFC 1035 §4.2.2 frames every message with a 2-byte big-endian length, and
//! RFC 7766 §6.2.1.1 lets one connection carry several queries. Both are
//! implemented here: a connection is answered serially until the client closes
//! it or the idle budget runs out.
//!
//! The slot pool, the accept loop and the shutdown protocol are
//! `listener.Core`'s (milestone-18 ruling 1); this file is the per-connection
//! loop and nothing else. Connection slots are fixed and pre-allocated. Over
//! capacity the listener closes the new stream immediately and counts it; it
//! never queues, and it never allocates per connection.
//!
//! No stream read or write in 0.16.0 accepts a timeout, so every per-connection
//! operation is raced against `Options.idle_timeout` through `listener.race`
//! and the loser is canceled.

const std = @import("std");
const address = @import("../platform/address.zig");
const handler = @import("handler.zig");
const listener = @import("listener.zig");
const transport = @import("../upstream/transport.zig");

/// The stream buffers only stage the framing bytes. A message longer than this
/// is read straight into `Conn.query` and written straight from `Conn.reply`,
/// so making them larger would buy nothing.
const stream_buffer_len = 1024;

pub const Options = struct {
    max_connections: u16 = 64,
    /// RFC 7766 §6.2.3 recommends a few seconds of idle tolerance.
    idle_timeout: std.Io.Clock.Duration = .{ .raw = .fromSeconds(10), .clock = .awake },
};

/// TCP/53 keeps no counter of its own: the shared six are exactly what it
/// counts.
pub const Stats = listener.CoreStats;

/// A plain copy of `Stats`, the shape `metrics.counterGroup` walks for the
/// `nxdns_tcp_server_*` families. Every counter is exported, including the
/// two refusals: a listener that turns clients away at capacity is the thing an
/// operator most needs to see, and the module doc promises it is counted.
pub const Snapshot = struct {
    connections: u64,
    rejected_at_capacity: u64,
    rejected_at_shutdown: u64,
    accept_errors: u64,
    connection_errors: u64,
    idle_timeouts: u64,
};

pub const TcpServer = struct {
    core: listener.Core(Config),
    handler: *handler.Handler,
    options: Options,

    /// One slot is ~137 KiB, so the default 64 connections cost ~8.8 MiB, which
    /// is inside the PLAN §18 budget. The two message buffers cannot be shared
    /// or shrunk: the handler holds the query while the reply is built, and
    /// both ceilings are the 65535 bytes the length prefix can express.
    pub const Payload = struct {
        query: [transport.max_message_len]u8,
        reply: [transport.max_message_len]u8,
        /// The handler's per-query working memory. It belongs to the slot so
        /// that answering a message allocates nothing, and a connection is
        /// answered serially, so one query uses it at a time.
        scratch: handler.Scratch,
    };

    const Config = struct {
        pub const Owner = TcpServer;
        pub const ConnPayload = Payload;
        pub const serveConn = serveOne;
        pub const read_buffer_len = stream_buffer_len;
        pub const write_buffer_len = stream_buffer_len;
        pub const log = std.log.scoped(.tcp_server);
        pub const name = "tcp";
    };

    pub const Conn = listener.Core(Config).Conn;
    pub const ListenError = listener.Core(Config).ListenError;

    pub fn listen(
        gpa: std.mem.Allocator,
        io: std.Io,
        listen_address: std.Io.net.IpAddress,
        h: *handler.Handler,
        options: Options,
    ) ListenError!TcpServer {
        return .{
            .core = try listener.Core(Config).listen(gpa, io, listen_address, options.max_connections),
            .handler = h,
            .options = options,
        };
    }

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

    /// The counters, read one at a time. A scrape that lands mid-accept can see
    /// a connection counted before its outcome is; a lock would buy a
    /// consistency no consumer needs.
    pub fn snapshotStats(self: *const TcpServer) Snapshot {
        const stats = &self.core.stats;
        return .{
            .connections = stats.connections.load(.monotonic),
            .rejected_at_capacity = stats.rejected_at_capacity.load(.monotonic),
            .rejected_at_shutdown = stats.rejected_at_shutdown.load(.monotonic),
            .accept_errors = stats.accept_errors.load(.monotonic),
            .connection_errors = stats.connection_errors.load(.monotonic),
            .idle_timeouts = stats.idle_timeouts.load(.monotonic),
        };
    }

    /// Accept loop. Returns when the task is canceled or `deinit` stops it.
    pub fn serve(self: *TcpServer, io: std.Io) void {
        self.core.serve(io);
    }

    pub fn deinit(self: *TcpServer, io: std.Io) void {
        self.core.deinit(io);
        self.* = undefined;
    }

    /// One connection, answered serially until the client closes it, the idle
    /// budget runs out, or a framing error ends it. The core closes the slot
    /// when this returns.
    fn serveOne(self: *TcpServer, io: std.Io, index: usize) void {
        const conn = &self.core.conns[index];
        const stats = &self.core.stats;
        var reader = conn.stream.reader(io, &conn.read_buf);
        var writer = conn.stream.writer(io, &conn.write_buf);
        const budget = self.options.idle_timeout;

        while (true) {
            var prefix: [transport.prefix_len]u8 = undefined;
            var got: usize = 0;
            switch (listener.race(io, budget, listener.readPrefix, .{ &reader.interface, &prefix, &got })) {
                .ok => {},
                .timed_out => {
                    listener.bump(&stats.idle_timeouts);
                    return;
                },
                .canceled => return,
                .failed => {
                    listener.bump(&stats.connection_errors);
                    return;
                },
            }

            // A client that closes between messages has finished asking, which
            // is the normal end of a connection, not a failure.
            if (got == 0) return;
            if (got != transport.prefix_len) {
                listener.bump(&stats.connection_errors);
                return;
            }

            // RFC 1035 §4.2.2 gives no meaning to a zero-length message, and
            // the prefix is a u16 so it can never exceed `max_message_len`.
            const len = transport.parsePrefix(prefix);
            if (len == 0) {
                listener.bump(&stats.connection_errors);
                return;
            }

            switch (listener.race(io, budget, listener.readBody, .{ &reader.interface, conn.payload.query[0..len] })) {
                .ok => {},
                .canceled => return,
                // A half-sent message is a broken peer, not an idle one.
                .timed_out, .failed => {
                    listener.bump(&stats.connection_errors);
                    return;
                },
            }

            const outcome = self.handler.handle(
                io,
                .tcp,
                address.NetAddress.fromIp(conn.peer),
                conn.payload.query[0..len],
                &conn.payload.reply,
                &conn.payload.scratch,
            );
            const bytes = switch (outcome) {
                // There is no framing for "no answer", so the connection ends.
                .drop => return,
                .reply => |b| b,
            };

            const out = transport.framePrefix(@intCast(bytes.len));
            switch (listener.race(io, budget, listener.writeReply, .{ &writer.interface, &out, bytes })) {
                .ok => {},
                .canceled => return,
                .timed_out, .failed => {
                    listener.bump(&stats.connection_errors);
                    return;
                },
            }
        }
    }
};
