//! 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 `server/handler.zig`
//! 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;

/// `tcp://[` + the longest IPv6 text form + `]:65535`, the widest spelling
/// `identityText` can produce.
pub const max_identity_len: usize = "tcp://[".len + 45 + "]:65535".len;

pub const ForwardClient = struct {
    resolver: validate.Resolver,
    /// The resolver as text, owned here so the `transport.Client` out-parameter
    /// has something stable to borrow: `validate.Resolver` is a parsed address,
    /// and a caller logging the exchange needs its spelling.
    identity_buf: [max_identity_len]u8 = undefined,
    identity_len: usize = 0,
    /// 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);
        var self: ForwardClient = .{
            .resolver = resolver,
            .frame_buf = frame_buf,
            .read_timeout = read_timeout,
        };
        self.identity_len = identityText(resolver, &self.identity_buf).len;
        return self;
    }

    /// `udp://192.168.1.1:53`, `tcp://[fd00::1]:53` — the same spelling
    /// `validate.parseResolver` accepts, so a log row names the configured
    /// value. Valid for as long as this client is.
    pub fn identity(self: *const ForwardClient) []const u8 {
        return self.identity_buf[0..self.identity_len];
    }

    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,
        selected: *?[]const u8,
    ) transport.ExchangeError![]u8 {
        const self: *ForwardClient = @ptrCast(@alignCast(ptr));
        // Set before the attempt: a failed forward-zone exchange still names
        // the resolver it was sent to.
        selected.* = self.identity();
        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.
    ///
    /// `read_timeout` bounds the WHOLE exchange, truncation fallback included:
    /// the instant is computed once here and every blocking step inside runs
    /// against it, so a truncated UDP answer followed by a stalled TCP retry
    /// costs one budget rather than two.
    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;
        const expiry_at: std.Io.Clock.Timestamp = .fromNow(io, self.read_timeout);
        // The outcome is deliberately discarded. A forward zone has exactly one
        // configured resolver, so an expiry here is still that resolver failing
        // to answer in time: `error.Timeout` is peer evidence, and the
        // budget/peer distinction is upstream-pool policy.
        var outcome: transport.RaceOutcome = .completed;
        return transport.raceUntilTagged(io, expiry_at, &outcome, route, .{
            self,
            io,
            query,
            response_buf,
            expiry_at,
        }) catch |err| {
            switch (transport.group(err)) {
                .peer_fault, .local_resource => self.stats.failures += 1,
                .cancellation, .budget_exhausted => {},
            }
            return err;
        };
    }

    fn route(
        self: *ForwardClient,
        io: std.Io,
        query: []const u8,
        response_buf: []u8,
        expiry_at: std.Io.Clock.Timestamp,
    ) transport.ExchangeError![]u8 {
        if (self.resolver.scheme == .udp) {
            if (try self.exchangeUdp(io, query, response_buf, expiry_at)) |reply| return reply;
        }
        return self.tcpOnce(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,
        expiry_at: std.Io.Clock.Timestamp,
    ) 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);
        };

        // The exchange-wide instant, not a fresh duration: a discarded foreign
        // datagram restarts the receive, and a duration would hand each retry
        // the full budget again.
        const deadline: std.Io.Timeout = .{ .deadline = expiry_at };

        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;
        }
    }

    /// Unbounded on its own: `exchange` runs it inside the exchange-wide race,
    /// which is what cancels a stalled connect or read. It takes no budget of
    /// its own, so a truncation fallback does not start a second one.
    /// `ConnectOptions.timeout` is never set: the Threaded backend panics on it
    /// (Threaded.zig:12076).
    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);
    }
};

fn identityText(resolver: validate.Resolver, buf: *[max_identity_len]u8) []const u8 {
    var w: std.Io.Writer = .fixed(buf);
    w.writeAll(switch (resolver.scheme) {
        .udp => "udp://",
        .tcp => "tcp://",
    }) catch unreachable;

    const bracketed = switch (resolver.addr) {
        .ip4 => false,
        .ip6 => true,
    };
    if (bracketed) w.writeByte('[') catch unreachable;
    resolver.addr.format(&w) catch unreachable;
    if (bracketed) w.writeByte(']') catch unreachable;
    w.print(":{d}", .{resolver.port}) catch unreachable;
    return w.buffered();
}

/// 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;
const build_options = @import("build_options");
const name_mod = @import("../dns/name.zig");
const packet = @import("../dns/packet.zig");
const question = @import("../dns/question.zig");
const types = @import("../dns/types.zig");

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 client owns its resolver identity in both address families" {
    var buf = testBuf();
    const v4: ForwardClient = .init(
        try validate.parseResolver("udp://192.168.1.1:5300"),
        &buf,
        .{ .raw = .fromSeconds(1), .clock = .awake },
    );
    try testing.expectEqualStrings("udp://192.168.1.1:5300", v4.identity());

    var buf6 = testBuf();
    const v6: ForwardClient = .init(
        try validate.parseResolver("tcp://[fd00::1]:5353"),
        &buf6,
        .{ .raw = .fromSeconds(1), .clock = .awake },
    );
    try testing.expectEqualStrings("tcp://[fd00::1]:5353", v6.identity());

    // The borrow points into the client, not into `init`'s frame.
    try testing.expect(@intFromPtr(v6.identity().ptr) >= @intFromPtr(&v6));
}

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),
    );
}

const one_budget_ms = 400;
/// Late enough in the budget that a second, fresh budget for the TCP leg would
/// be unmistakable in the elapsed time.
const truncate_after_ms = 300;

/// Answers the first datagram late in the budget with TC=1, which sends the
/// client to TCP — where a listener that never accepts leaves it stalled.
fn truncatingThenStallingResolver(io: std.Io, socket: *const net.Socket) void {
    var buf: [2048]u8 = undefined;
    const msg = socket.receive(io, &buf) catch return;
    const request = packet.parse(msg.data) catch return;

    (std.Io.Clock.Duration{
        .raw = .fromMilliseconds(truncate_after_ms),
        .clock = .awake,
    }).sleep(io) catch return;

    // The question has to be echoed: `transport.validateResponse` runs before
    // the client reads the TC bit, so a bare header would come back as
    // `BadResponse` and never reach the TCP fallback this case is about.
    const q = packet.firstQuestion(request) orelse return;
    var reply_buf: [512]u8 = undefined;
    var b = packet.ResponseBuilder.init(&reply_buf, request.header, q) catch return;
    const reply = b.finish();

    var parsed = dns_header.parse(reply) catch return;
    parsed.flags.tc = true;
    dns_header.encode(parsed, reply[0..types.header_len]);
    socket.send(io, &msg.from, reply) catch return;
}

fn testQuery(io: std.Io, buf: []u8) ![]const u8 {
    var id_bytes: [2]u8 = undefined;
    io.random(&id_bytes);
    dns_header.encode(.{
        .id = std.mem.readInt(u16, &id_bytes, .big),
        .flags = .{
            .rcode = .no_error,
            .z = 0,
            .ra = false,
            .rd = true,
            .tc = false,
            .aa = false,
            .opcode = .query,
            .qr = false,
        },
        .qdcount = 1,
        .ancount = 0,
        .nscount = 0,
        .arcount = 0,
    }, buf[0..types.header_len]);

    var w: std.Io.Writer = .fixed(buf[types.header_len..]);
    try question.encode(.{
        .name = try name_mod.fromText("nas.lan"),
        .qtype = .a,
        .qclass = .in,
    }, &w);
    return buf[0 .. types.header_len + w.buffered().len];
}

test "one budget covers the udp leg, the TC=1 fallback and the tcp leg" {
    if (!build_options.integration) return error.SkipZigTest;

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

    const bind_address: net.IpAddress = try .parse("127.0.0.1", 0);
    const socket = try bind_address.bind(io, .{ .mode = .dgram });
    defer socket.close(io);
    const port = socket.address.ip4.port;

    // Bound but never accepted: the connect completes out of the kernel's
    // backlog and the read then waits forever, which is the stall a second
    // budget would be spent on.
    const tcp_address: net.IpAddress = try .parse("127.0.0.1", port);
    var tcp_listener = try tcp_address.listen(io, .{ .reuse_address = true });
    defer tcp_listener.deinit(io);

    var group: std.Io.Group = .init;
    defer group.cancel(io);
    try group.concurrent(io, truncatingThenStallingResolver, .{ io, &socket });

    var url_buf: [64]u8 = undefined;
    const url = try std.fmt.bufPrint(&url_buf, "udp://127.0.0.1:{d}", .{port});

    var frame_buf: [min_frame_buf]u8 = undefined;
    var fc: ForwardClient = .init(
        try validate.parseResolver(url),
        &frame_buf,
        .{ .raw = .fromMilliseconds(one_budget_ms), .clock = .awake },
    );

    var query_buf: [types.header_len + types.max_name_len + 4]u8 = undefined;
    const query = try testQuery(io, &query_buf);
    var response_buf: [2048]u8 = undefined;

    const started = std.Io.Clock.awake.now(io);
    try testing.expectError(error.Timeout, fc.exchange(io, query, &response_buf));
    const elapsed = started.durationTo(std.Io.Clock.awake.now(io)).nanoseconds;

    // Two budgets would spend 300 ms on the UDP leg and then a fresh 400 ms on
    // the stalled TCP leg. The bound sits between one budget and that sum, so
    // the double-budget shape cannot pass.
    try testing.expect(elapsed < @as(i96, one_budget_ms + truncate_after_ms / 2) * std.time.ns_per_ms);
    try testing.expectEqual(@as(u64, 1), fc.stats.udp_truncated);
    try testing.expectEqual(@as(u64, 1), fc.stats.failures);
}