//! The UDP/53 listener.
//!
//! One receive loop feeds a bounded pool of in-flight tasks. A datagram that
//! finds no free slot is dropped and counted, never queued: an unbounded queue
//! only moves the failure from "visible drop" to "invisible latency", which is
//! the failure mode PLAN §1 names.
//!
//! Every slot carries its own query and reply buffer, so answering a datagram
//! allocates nothing. The pool is sized once in `bind` and never grows.

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

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

/// The EDNS ceiling this server advertises. A larger datagram arrives
/// truncated and is dropped rather than guessed at.
pub const max_datagram = 4096;

comptime {
    // `handler.handle` asserts this range on the buffers it is given.
    std.debug.assert(max_datagram >= handler.udp_limit_min);
    std.debug.assert(max_datagram >= handler.udp_limit_max);
    std.debug.assert(transport.max_message_len >= handler.udp_limit_min);
}

/// How long a receive waits before the loop re-reads the stop flag. `deinit`
/// therefore returns within this much time, without depending on a blocked
/// `receive` being unblocked by a close — a guarantee POSIX does not make.
const poll_interval: std.Io.Clock.Duration = .{ .raw = .fromMilliseconds(100), .clock = .awake };

pub const Options = struct {
    max_in_flight: u16 = 64,
};

pub const Stats = struct {
    received: std.atomic.Value(u64) = .init(0),
    /// The datagram arrived with `flags.trunc`: the kernel discarded its tail.
    dropped_oversize: std.atomic.Value(u64) = .init(0),
    dropped_no_slot: std.atomic.Value(u64) = .init(0),
    /// The handler answered `.drop` — an upstream exchange canceled by
    /// shutdown. The handler counts its own malformed drops; this counter is
    /// what keeps the listener's "every dropped datagram is counted" rule true
    /// for the drops it cannot see.
    dropped_handler: std.atomic.Value(u64) = .init(0),
    receive_errors: std.atomic.Value(u64) = .init(0),
    send_errors: std.atomic.Value(u64) = .init(0),
};

/// A plain copy of `Stats`, the shape `metrics.counterGroup` walks for the
/// `nxdns_udp_server_*` families. Every counter is exported: the module doc
/// promises that a dropped datagram is counted, and a count nothing can read is
/// not a count.
pub const Snapshot = struct {
    received: u64,
    dropped_oversize: u64,
    dropped_no_slot: u64,
    dropped_handler: u64,
    receive_errors: u64,
    send_errors: u64,
};

/// Lifecycle of the receive loop. `serve` claims `.serving`, `deinit` publishes
/// `.closing`, and the two meet at `stopped` so no task touches the slots after
/// they are freed.
const State = enum(u32) { idle, serving, closing };

pub const UdpServer = struct {
    socket: std.Io.net.Socket,
    handler: *handler.Handler,
    slots: []Slot,
    mutex: std.Io.Mutex,
    stats: Stats,
    state: std.atomic.Value(State),
    stopped: std.Io.Event,

    pub const Slot = struct {
        query: [max_datagram]u8,
        /// The buffer the handler receives the upstream answer into, so it
        /// holds a whole DNS message, not a whole datagram. A reply between
        /// `max_datagram` and `transport.max_message_len` must reach the
        /// handler intact: only the handler can turn it into the TC=1 reply
        /// that makes the client retry over TCP (S5.3). A 4096-byte buffer
        /// would instead fail the exchange with `error.ResponseTooLarge` and
        /// answer SERVFAIL, which is wrong.
        ///
        /// The datagram actually sent stays bounded by `udpLimit` inside the
        /// handler, so nothing larger than 4096 bytes leaves this socket.
        ///
        /// Cost: 4096 + 65535 + the scratch below ≈ 74 KiB per slot, so the
        /// default 64 slots hold ≈ 4.6 MiB. The PLAN §18 budget is 100 MiB with
        /// ~1M blocked domains, so this pool takes about 5% of it.
        reply: [transport.max_message_len]u8,
        /// The handler's per-query working memory. It belongs to the slot so
        /// that answering a datagram still allocates nothing, and one slot
        /// serves one query at a time.
        scratch: handler.Scratch,
        from: std.Io.net.IpAddress,
        len: usize,
        /// Guarded by `UdpServer.mutex`.
        in_use: bool,
    };

    pub const BindError = std.Io.net.IpAddress.BindError || error{OutOfMemory};

    pub fn bind(
        gpa: std.mem.Allocator,
        io: std.Io,
        bind_address: std.Io.net.IpAddress,
        h: *handler.Handler,
        options: Options,
    ) BindError!UdpServer {
        std.debug.assert(options.max_in_flight > 0);

        const slots = try gpa.alloc(Slot, options.max_in_flight);
        errdefer gpa.free(slots);
        for (slots) |*slot| slot.in_use = false;

        const local = bind_address;
        const socket = try local.bind(io, .{ .mode = .dgram });

        return .{
            .socket = socket,
            .handler = h,
            .slots = slots,
            .mutex = .init,
            .stats = .{},
            .state = .init(.idle),
            .stopped = .unset,
        };
    }

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

    /// The counters, read one at a time. A scrape that lands mid-datagram can
    /// see a receive counted before its drop is; a lock would buy a consistency
    /// no consumer needs, and the receive loop takes that lock per datagram.
    pub fn snapshotStats(self: *const UdpServer) Snapshot {
        return .{
            .received = self.stats.received.load(.monotonic),
            .dropped_oversize = self.stats.dropped_oversize.load(.monotonic),
            .dropped_no_slot = self.stats.dropped_no_slot.load(.monotonic),
            .dropped_handler = self.stats.dropped_handler.load(.monotonic),
            .receive_errors = self.stats.receive_errors.load(.monotonic),
            .send_errors = self.stats.send_errors.load(.monotonic),
        };
    }

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

        var group: std.Io.Group = .init;
        self.receiveLoop(io, &group);

        // A reply that is half sent is worse than no reply, so the in-flight
        // tasks are awaited even when this task is being canceled.
        const prev = io.swapCancelProtection(.blocked);
        group.await(io) catch |err| switch (err) {
            error.Canceled => unreachable,
        };
        _ = io.swapCancelProtection(prev);

        self.stopped.set(io);
    }

    pub fn deinit(self: *UdpServer, gpa: std.mem.Allocator, io: std.Io) void {
        // `serve` reads this between receives, so it stops on its own and the
        // socket is closed only once nothing can be reading from it.
        if (self.state.swap(.closing, .acq_rel) == .serving) self.stopped.waitUncancelable(io);
        self.socket.close(io);
        gpa.free(self.slots);
        self.* = undefined;
    }

    fn receiveLoop(self: *UdpServer, io: std.Io, group: *std.Io.Group) void {
        // A datagram that finds no free slot still has to leave the socket
        // buffer, or the loop would spin on the same undeliverable packet.
        var overflow: [max_datagram]u8 = undefined;

        while (self.state.load(.acquire) == .serving) {
            const claimed = self.claim(io);
            const buffer = if (claimed) |index| &self.slots[index].query else &overflow;

            const msg = self.socket.receiveTimeout(io, buffer, .{ .duration = poll_interval }) catch |err| {
                if (claimed) |index| self.release(io, index);
                switch (err) {
                    error.Canceled => return,
                    // The stop flag is re-read at the top of the loop.
                    error.Timeout => continue,
                    else => {
                        bump(&self.stats.receive_errors);
                        log.debug("udp receive failed: {t}", .{err});
                        // One bad datagram must not kill the listener, and a
                        // persistent error must not turn it into a spin.
                        poll_interval.sleep(io) catch return;
                        continue;
                    },
                }
            };

            bump(&self.stats.received);

            const index = claimed orelse {
                bump(&self.stats.dropped_no_slot);
                continue;
            };

            // The kernel already threw the tail away, so the message cannot be
            // parsed and any answer would be a guess.
            if (msg.flags.trunc) {
                bump(&self.stats.dropped_oversize);
                self.release(io, index);
                continue;
            }

            const slot = &self.slots[index];
            slot.from = msg.from;
            slot.len = msg.data.len;

            group.concurrent(io, respondOne, .{ self, io, index }) catch |err| switch (err) {
                error.ConcurrencyUnavailable => {
                    bump(&self.stats.dropped_no_slot);
                    self.release(io, index);
                },
            };
        }
    }

    fn respondOne(self: *UdpServer, io: std.Io, index: usize) void {
        const slot = &self.slots[index];
        defer self.release(io, index);

        const outcome = self.handler.handle(
            io,
            .udp,
            address.NetAddress.fromIp(slot.from),
            slot.query[0..slot.len],
            &slot.reply,
            &slot.scratch,
        );
        switch (outcome) {
            .drop => bump(&self.stats.dropped_handler),
            .reply => |bytes| self.socket.send(io, &slot.from, bytes) catch |err| {
                bump(&self.stats.send_errors);
                log.debug("udp send to a client failed: {t}", .{err});
            },
        }
    }

    fn claim(self: *UdpServer, io: std.Io) ?usize {
        // 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 index = firstFree(self.slots) orelse return null;
        self.slots[index].in_use = true;
        return index;
    }

    fn release(self: *UdpServer, io: std.Io, index: usize) void {
        self.mutex.lockUncancelable(io);
        defer self.mutex.unlock(io);
        self.slots[index].in_use = false;
    }
};

/// The capacity rule, without the mutex, so it is testable without a backend.
fn firstFree(slots: []const UdpServer.Slot) ?usize {
    for (slots, 0..) |*slot, index| {
        if (!slot.in_use) return index;
    }
    return null;
}

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

const testing = std.testing;

fn testSlots(count: usize) ![]UdpServer.Slot {
    const slots = try testing.allocator.alloc(UdpServer.Slot, count);
    for (slots) |*slot| slot.in_use = false;
    return slots;
}

test "the slot pool hands out every index once" {
    const slots = try testSlots(3);
    defer testing.allocator.free(slots);

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

test "a full slot pool refuses instead of growing" {
    const slots = try testSlots(2);
    defer testing.allocator.free(slots);

    for (slots) |*slot| slot.in_use = true;
    try testing.expectEqual(@as(?usize, null), firstFree(slots));
}

test "a slot's reply buffer holds a whole DNS message" {
    const slots = try testSlots(1);
    defer testing.allocator.free(slots);

    // The handler asserts both ends of this range on the buffer it is given,
    // and an upstream reply of up to 65535 bytes has to fit so the handler can
    // truncate it instead of the exchange failing.
    try testing.expectEqual(@as(usize, transport.max_message_len), slots[0].reply.len);
    try testing.expect(slots[0].reply.len >= handler.udp_limit_min);
    try testing.expect(slots[0].reply.len > max_datagram);
}

test "the default slot pool stays inside the memory budget" {
    // 4096 + 65535 + scratch ≈ 74 KiB per slot; 64 slots ≈ 4.6 MiB, against the
    // 100 MiB of PLAN §18.
    const options: Options = .{};
    const pool_bytes = @sizeOf(UdpServer.Slot) * @as(usize, options.max_in_flight);
    try testing.expect(pool_bytes < 8 * 1024 * 1024);
}

test "a released slot is handed out again" {
    const slots = try testSlots(2);
    defer testing.allocator.free(slots);

    for (slots) |*slot| slot.in_use = true;
    slots[1].in_use = false;
    try testing.expectEqual(@as(?usize, 1), firstFree(slots));
}
