//! Plain UDP/TCP resolver client for conditional forward zones (PLAN §6.5).
//!
//! A forward zone points at a box on the LAN — a router, a NAS, an internal
//! resolver — which speaks port 53 and nothing else. `transport.Endpoint` knows
//! only `https://` and `tls://` by design, so the configuration for this client
//! comes from `validate.Resolver` instead. The interface it implements is the
//! same `transport.Client` every upstream implements, so the Phase 7 handler
//! treats a forward zone exactly like any other exchange.
//!
//! No health tracking and no backoff live here. `upstream/health.zig` and
//! `upstream/pool.zig` model the upstream *pool*, where failing over to a second
//! endpoint is the whole point. A forward zone has exactly one designated
//! resolver and no failover partner, so a backoff would only add latency to a
//! failure the caller already sees. Their absence is a decision, not an
//! oversight.
//!
//! One `ForwardClient` is used by one task at a time: `stats` is a plain struct
//! and `frame_buf` is not shared.

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

const transport = @import("../upstream/transport.zig");
const validate = @import("../config/validate.zig");
const dns_header = @import("../dns/header.zig");

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

/// RFC 1035 §4.2.2 length prefix for DNS over TCP.
/// The TCP path splits `frame_buf` between the socket writer and the socket
/// reader. Neither half has to hold a whole message — the reply is read
/// straight into the caller's `response_buf` — so this is a floor that keeps
/// each half large enough to frame a query in one write, not a capacity.
pub const min_frame_buf: usize = 1024;

pub const ForwardClient = struct {
    resolver: validate.Resolver,
    /// Caller-owned scratch for the TCP length-prefixed path.
    frame_buf: []u8,
    /// On the `.awake` clock at the caller's choosing, so a suspended host does
    /// not burn the budget while it sleeps.
    read_timeout: std.Io.Clock.Duration,
    stats: Stats = .{},

    pub const Stats = struct {
        queries: u64 = 0,
        /// TC=1 over UDP, so the exchange was retried over TCP.
        udp_truncated: u64 = 0,
        /// A datagram arrived from an address other than the resolver's. It was
        /// discarded and the receive retried within the remaining budget, which
        /// is invisible to the caller and would otherwise be an unrecorded
        /// failure mode.
        foreign_datagrams: u64 = 0,
        /// Exchanges that returned a peer fault or a local resource error.
        /// A cancellation is neither, so it is not counted.
        failures: u64 = 0,
    };

    /// An undersized `frame_buf` is a wiring bug in this process, not a runtime
    /// condition, so it is an assertion.
    pub fn init(
        resolver: validate.Resolver,
        frame_buf: []u8,
        read_timeout: std.Io.Clock.Duration,
    ) ForwardClient {
        std.debug.assert(frame_buf.len >= min_frame_buf);
        return .{
            .resolver = resolver,
            .frame_buf = frame_buf,
            .read_timeout = read_timeout,
        };
    }

    pub fn client(self: *ForwardClient) transport.Client {
        return .{ .ptr = self, .exchangeFn = exchangeFn };
    }

    fn exchangeFn(
        ptr: *anyopaque,
        io: std.Io,
        query: []const u8,
        response_buf: []u8,
    ) transport.ExchangeError![]u8 {
        const self: *ForwardClient = @ptrCast(@alignCast(ptr));
        return self.exchange(io, query, response_buf);
    }

    /// `.udp` resolvers send one datagram and fall back to TCP when the answer
    /// comes back with TC=1. `.tcp` resolvers skip straight to the TCP path.
    pub fn exchange(
        self: *ForwardClient,
        io: std.Io,
        query: []const u8,
        response_buf: []u8,
    ) transport.ExchangeError![]u8 {
        // The TCP length prefix is 16-bit, so a longer query cannot be framed.
        if (query.len > transport.max_message_len) return error.BufferTooSmall;
        if (response_buf.len == 0) return error.BufferTooSmall;

        self.stats.queries += 1;
        return self.route(io, query, response_buf) catch |err| {
            switch (transport.group(err)) {
                .peer_fault, .local_resource => self.stats.failures += 1,
                .cancellation => {},
            }
            return err;
        };
    }

    fn route(
        self: *ForwardClient,
        io: std.Io,
        query: []const u8,
        response_buf: []u8,
    ) transport.ExchangeError![]u8 {
        if (self.resolver.scheme == .udp) {
            if (try self.exchangeUdp(io, query, response_buf)) |reply| return reply;
        }
        return self.exchangeTcp(io, query, response_buf);
    }

    /// `null` means the resolver set TC=1 and the caller must retry over TCP.
    ///
    /// The socket is bound to the wildcard address of the resolver's family on
    /// an ephemeral port, so the kernel picks the source port for every
    /// exchange rather than this process reusing one.
    fn exchangeUdp(
        self: *ForwardClient,
        io: std.Io,
        query: []const u8,
        response_buf: []u8,
    ) transport.ExchangeError!?[]u8 {
        const dest = self.destination();
        const local = wildcardFor(dest);

        const socket = local.bind(io, .{ .mode = .dgram }) catch |err| {
            log.debug("forward resolver: udp bind failed: {s}", .{@errorName(err)});
            return transport.mapPhase(err, error.ConnectFailed);
        };
        defer transport.closeBlocked(io, &socket);

        socket.send(io, &dest, query) catch |err| {
            log.debug("forward resolver: udp send failed: {s}", .{@errorName(err)});
            return transport.mapPhase(err, error.SendFailed);
        };

        // A deadline, not a duration: a discarded foreign datagram restarts the
        // receive, and a duration would hand each retry the full budget again.
        const deadline = (std.Io.Timeout{ .duration = self.read_timeout }).toDeadline(io);

        while (true) {
            const msg = socket.receiveTimeout(io, response_buf, deadline) catch |err| switch (err) {
                error.Timeout => return error.Timeout,
                error.ConcurrencyUnavailable => return error.SystemResources,
                else => return transport.mapPhase(err, error.ReceiveFailed),
            };

            // Off-path spoofing is the reason the source address is checked at
            // all: the first datagram to arrive is not necessarily the
            // resolver's.
            if (!msg.from.eql(&dest)) {
                self.stats.foreign_datagrams += 1;
                continue;
            }

            // The kernel threw the tail away because `response_buf` was too
            // small, so the message cannot be parsed and TC=1 cannot be read
            // out of it.
            if (msg.flags.trunc) return error.ResponseTooLarge;

            const reply = response_buf[0..msg.data.len];
            try transport.validateResponse(query, reply);

            // Read after validation: acting on the TC bit of a message that has
            // not been matched to the query would let anything that reaches the
            // socket force a TCP connection.
            const parsed = dns_header.parse(reply) catch return error.BadResponse;
            if (parsed.flags.tc) {
                self.stats.udp_truncated += 1;
                return null;
            }
            return reply;
        }
    }

    /// The read budget bounds the whole TCP exchange through
    /// `transport.raceWithin`. `ConnectOptions.timeout` is never set: the
    /// Threaded backend panics on it (Threaded.zig:12076).
    fn exchangeTcp(
        self: *ForwardClient,
        io: std.Io,
        query: []const u8,
        response_buf: []u8,
    ) transport.ExchangeError![]u8 {
        return transport.raceWithin(io, self.read_timeout, tcpOnce, .{ self, io, query, response_buf });
    }

    fn tcpOnce(
        self: *ForwardClient,
        io: std.Io,
        query: []const u8,
        response_buf: []u8,
    ) transport.ExchangeError![]u8 {
        const dest = self.destination();

        const stream = dest.connect(io, .{ .mode = .stream }) catch |err| {
            log.debug("forward resolver: tcp connect failed: {s}", .{@errorName(err)});
            return transport.mapPhase(err, error.ConnectFailed);
        };
        defer transport.closeBlocked(io, &stream);

        const split = self.frame_buf.len / 2;
        var stream_writer = stream.writer(io, self.frame_buf[0..split]);
        var stream_reader = stream.reader(io, self.frame_buf[split..]);

        const w = &stream_writer.interface;
        const prefix = transport.framePrefix(@intCast(query.len));
        w.writeAll(&prefix) catch |err| return sendFailure(&stream_writer, err);
        w.writeAll(query) catch |err| return sendFailure(&stream_writer, err);
        w.flush() catch |err| return sendFailure(&stream_writer, err);

        const r = &stream_reader.interface;
        var prefix_bytes: [transport.prefix_len]u8 = undefined;
        r.readSliceAll(&prefix_bytes) catch |err| return receiveFailure(&stream_reader, err);

        // RFC 1035 §4.2.2 gives no meaning to a zero-length message.
        const len = transport.parsePrefix(prefix_bytes);
        if (len == 0) return error.BadResponse;
        if (len > response_buf.len) return error.ResponseTooLarge;
        r.readSliceAll(response_buf[0..len]) catch |err| return receiveFailure(&stream_reader, err);

        try transport.validateResponse(query, response_buf[0..len]);
        return response_buf[0..len];
    }

    fn destination(self: *const ForwardClient) net.IpAddress {
        return self.resolver.addr.toIp(self.resolver.port);
    }
};

/// The local address a datagram to `dest` is sent from: same family, port
/// chosen by the kernel.
fn wildcardFor(dest: net.IpAddress) net.IpAddress {
    return switch (dest) {
        .ip4 => .{ .ip4 = .unspecified(0) },
        .ip6 => .{ .ip6 = .unspecified(0) },
    };
}

/// `Io.Writer` collapses everything to `error.WriteFailed` and stashes the
/// cause. Unwrapping it is what keeps `error.Canceled` and the local resource
/// errors out of the peer fault group.
fn sendFailure(stream_writer: *const net.Stream.Writer, err: anyerror) transport.ExchangeError {
    const cause: anyerror = if (err == error.WriteFailed and stream_writer.err != null)
        stream_writer.err.?
    else
        err;
    return transport.mapPhase(cause, error.SendFailed);
}

fn receiveFailure(stream_reader: *const net.Stream.Reader, err: anyerror) transport.ExchangeError {
    const cause: anyerror = if (err == error.ReadFailed and stream_reader.err != null)
        stream_reader.err.?
    else
        err;
    return transport.mapPhase(cause, error.ReceiveFailed);
}

const testing = std.testing;

fn testBuf() [min_frame_buf]u8 {
    return undefined;
}

test "ForwardClient satisfies the Client interface" {
    var buf = testBuf();
    var fc: ForwardClient = .init(
        try validate.parseResolver("udp://192.168.1.1:53"),
        &buf,
        .{ .raw = .fromMilliseconds(500), .clock = .awake },
    );

    const iface: transport.Client = fc.client();
    try testing.expectEqual(@as(*anyopaque, @ptrCast(&fc)), iface.ptr);
    try testing.expectEqual(validate.ResolverScheme.udp, fc.resolver.scheme);
    try testing.expectEqual(@as(u16, 53), fc.resolver.port);
}

test "the stats struct starts at zero" {
    const stats: ForwardClient.Stats = .{};
    try testing.expectEqual(@as(u64, 0), stats.queries);
    try testing.expectEqual(@as(u64, 0), stats.udp_truncated);
    try testing.expectEqual(@as(u64, 0), stats.foreign_datagrams);
    try testing.expectEqual(@as(u64, 0), stats.failures);
}

test "init keeps a tcp resolver on the tcp path" {
    var buf = testBuf();
    const fc: ForwardClient = .init(
        try validate.parseResolver("tcp://[fd00::1]:5353"),
        &buf,
        .{ .raw = .fromSeconds(2), .clock = .awake },
    );
    try testing.expectEqual(validate.ResolverScheme.tcp, fc.resolver.scheme);
    try testing.expectEqual(@as(u16, 5353), fc.resolver.port);

    const dest = fc.destination();
    try testing.expectEqual(net.IpAddress.Family.ip6, std.meta.activeTag(dest));
    try testing.expectEqual(@as(u16, 5353), dest.getPort());
}

test "the destination carries the resolver's address and port" {
    var buf = testBuf();
    const fc: ForwardClient = .init(
        try validate.parseResolver("udp://192.168.1.1:5300"),
        &buf,
        .{ .raw = .fromSeconds(1), .clock = .awake },
    );
    const dest = fc.destination();
    try testing.expectEqualSlices(u8, &.{ 192, 168, 1, 1 }, &dest.ip4.bytes);
    try testing.expectEqual(@as(u16, 5300), dest.ip4.port);
}

test "only the resolver's own address and port count as its datagram" {
    const dest: net.IpAddress = .{ .ip4 = .{ .bytes = .{ 192, 168, 1, 1 }, .port = 53 } };

    const same: net.IpAddress = .{ .ip4 = .{ .bytes = .{ 192, 168, 1, 1 }, .port = 53 } };
    try testing.expect(same.eql(&dest));

    // A different host, the right host on a different port, and the right
    // address in the wrong family are each a datagram this client discards.
    const other_host: net.IpAddress = .{ .ip4 = .{ .bytes = .{ 192, 168, 1, 2 }, .port = 53 } };
    try testing.expect(!other_host.eql(&dest));

    const other_port: net.IpAddress = .{ .ip4 = .{ .bytes = .{ 192, 168, 1, 1 }, .port = 5353 } };
    try testing.expect(!other_port.eql(&dest));

    const mapped: net.IpAddress = .{ .ip6 = .fromIp4(.{ .bytes = .{ 192, 168, 1, 1 }, .port = 53 }) };
    try testing.expect(!mapped.eql(&dest));
}

test "the local socket matches the resolver's family and takes an ephemeral port" {
    const v4 = wildcardFor(.{ .ip4 = .{ .bytes = .{ 1, 1, 1, 1 }, .port = 53 } });
    try testing.expectEqual(net.IpAddress.Family.ip4, std.meta.activeTag(v4));
    try testing.expectEqual(@as(u16, 0), v4.getPort());
    try testing.expectEqualSlices(u8, &.{ 0, 0, 0, 0 }, &v4.ip4.bytes);

    const v6 = wildcardFor(.{ .ip6 = .unspecified(53) });
    try testing.expectEqual(net.IpAddress.Family.ip6, std.meta.activeTag(v6));
    try testing.expectEqual(@as(u16, 0), v6.getPort());
}

test "mapPhase keeps local resource and cancellation errors out of the peer fault group" {
    const local = [_]anyerror{
        error.OutOfMemory,
        error.SystemResources,
        error.ProcessFdQuotaExceeded,
        error.SystemFdQuotaExceeded,
        error.Unexpected,
    };
    for (local) |err| {
        try testing.expectEqual(
            transport.Group.local_resource,
            transport.group(transport.mapPhase(err, error.ReceiveFailed)),
        );
    }

    try testing.expectEqual(
        transport.ExchangeError.Canceled,
        transport.mapPhase(error.Canceled, error.ConnectFailed),
    );

    // A refused connection is the resolver's side, so it stays a peer fault.
    try testing.expectEqual(
        transport.ExchangeError.ConnectFailed,
        transport.mapPhase(error.ConnectionRefused, error.ConnectFailed),
    );
}

test "a stashed stream error is preferred over the collapsed one" {
    var stream_writer: net.Stream.Writer = undefined;
    stream_writer.err = error.Canceled;
    try testing.expectEqual(
        transport.ExchangeError.Canceled,
        sendFailure(&stream_writer, error.WriteFailed),
    );

    stream_writer.err = error.ConnectionResetByPeer;
    try testing.expectEqual(
        transport.ExchangeError.SendFailed,
        sendFailure(&stream_writer, error.WriteFailed),
    );

    var stream_reader: net.Stream.Reader = undefined;
    stream_reader.err = error.SystemResources;
    try testing.expectEqual(
        transport.ExchangeError.SystemResources,
        receiveFailure(&stream_reader, error.ReadFailed),
    );

    // A peer that closes mid-frame never reaches `err`, so the collapsed error
    // is what classifies it.
    stream_reader.err = null;
    try testing.expectEqual(
        transport.ExchangeError.ReceiveFailed,
        receiveFailure(&stream_reader, error.EndOfStream),
    );
}