//! End-to-end test for the whole resolver path: a client socket on the
//! loopback, the real UDP and TCP listeners, the real handler, the real pool,
//! and two in-process fake upstreams.
//!
//! This lives in its own file because it needs `@import("build_options")`, which
//! only exists when the compilation is driven by build.zig. The body is compiled
//! by every `zig build test` run, so it cannot rot, and skips at run time unless
//! `-Dintegration` is passed.
//!
//! Hermetic: every socket is on 127.0.0.1 and the upstreams are structs, so
//! nothing leaves the machine. No stream read in 0.16.0 takes a timeout, so the
//! TCP client side runs as one task raced against a budget and nothing can hang.

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

const handler = @import("handler.zig");
const tcp_server = @import("tcp_server.zig");
const udp_server = @import("udp_server.zig");
const model = @import("../config/model.zig");
const response = @import("../filter/response.zig");
const packet = @import("../dns/packet.zig");
const record = @import("../dns/record.zig");
const types = @import("../dns/types.zig");
const health = @import("../upstream/health.zig");
const pool = @import("../upstream/pool.zig");
const transport = @import("../upstream/transport.zig");

const testing = std.testing;

const blocking_defaults: model.Blocking = .{};
const blocking: response.Options = .{
    .mode = blocking_defaults.response,
    .ttl = blocking_defaults.ttl,
};
const forward_timeout: std.Io.Clock.Duration = .{
    .raw = model.readTimeout(.{}),
    .clock = .awake,
};

/// An upstream and nothing else optional: no filtering, no cache, no log. The
/// listeners and the pool are what this test exercises, so the handler is the
/// same bare one its own tests use.
fn bareHandler(client: transport.Client) handler.Handler {
    return .{
        .upstream = client,
        .blocking = blocking,
        .forward_read_timeout = forward_timeout,
    };
}

/// Long enough that a loopback round trip cannot lose to scheduling, short
/// enough that a broken server fails the run instead of hanging it.
const budget: std.Io.Timeout = .{ .duration = .{ .raw = .fromSeconds(5), .clock = .awake } };
const budget_duration: std.Io.Clock.Duration = .{ .raw = .fromSeconds(5), .clock = .awake };

/// A query for example.com A: RD set, one question, no OPT. Each step of the
/// test rewrites the ID so a reply can only match the query it belongs to.
const query_bytes =
    "\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00" ++
    "\x07example\x03com\x00\x00\x01\x00\x01";

const udp_id: u16 = 0x1234;
const tcp_id: u16 = 0x5678;
const backoff_id: u16 = 0x9abc;

/// The address every fake upstream answers with, and its TTL.
const answer_rdata = [4]u8{ 93, 184, 216, 34 };
const answer_ttl: u32 = 300;

/// One failure is enough to open a backoff window, so the third query has a
/// single deterministic outcome. The window outlives the test many times over.
const test_cfg: health.Config = .{
    .failure_threshold = 1,
    .base_backoff_ms = 60_000,
    .max_backoff_ms = 60_000,
};

/// Nothing in this test is slow, so these budgets only exist to stop a wedged
/// attempt from hanging the run.
const pool_timeouts: pool.Timeouts = .{
    .attempt = .{ .raw = .fromSeconds(10), .clock = .awake },
    .total = .{ .raw = .fromSeconds(30), .clock = .awake },
};

fn queryWithId(buf: *[query_bytes.len]u8, id: u16) []const u8 {
    buf.* = query_bytes.*;
    packet.setId(buf, id);
    return buf;
}

/// Echoes the question and appends one A record. This is the smallest thing a
/// real upstream could return that the handler forwards unchanged, so the
/// assertions below check bytes that travelled the whole path.
fn answerQuery(query: []const u8, response_buf: []u8) transport.ExchangeError![]u8 {
    const request = packet.parse(query) catch return error.BadResponse;
    const q = packet.firstQuestion(request) orelse return error.BadResponse;

    var b = packet.ResponseBuilder.init(response_buf, request.header, q) catch
        return error.ResponseTooLarge;
    b.addAnswer(q.name, .a, .in, answer_ttl, &answer_rdata) catch
        return error.ResponseTooLarge;
    return b.finish();
}

/// The healthy upstream. `calls` is atomic because the listener tasks run on
/// other threads than the one asserting.
const GoodUpstream = struct {
    calls: std.atomic.Value(u64) = .init(0),

    fn exchangeFn(
        ptr: *anyopaque,
        io: std.Io,
        query: []const u8,
        response_buf: []u8,
    ) transport.ExchangeError![]u8 {
        _ = io;
        const self: *GoodUpstream = @ptrCast(@alignCast(ptr));
        _ = self.calls.fetchAdd(1, .monotonic);
        return answerQuery(query, response_buf);
    }

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

/// Returns `fault` on the first `fail_first` calls and answers after that.
/// `fail_first` is `maxInt` for an upstream that never recovers.
const FaultyUpstream = struct {
    fault: transport.PeerFault,
    fail_first: u64,
    calls: std.atomic.Value(u64) = .init(0),

    fn exchangeFn(
        ptr: *anyopaque,
        io: std.Io,
        query: []const u8,
        response_buf: []u8,
    ) transport.ExchangeError![]u8 {
        _ = io;
        const self: *FaultyUpstream = @ptrCast(@alignCast(ptr));
        const seen = self.calls.fetchAdd(1, .monotonic);
        if (seen < self.fail_first) return self.fault;
        return answerQuery(query, response_buf);
    }

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

fn testEntry(url: []const u8, upstream_client: transport.Client, priority: i32) pool.Entry {
    return .{
        .endpoint = transport.Endpoint.parse(url) catch unreachable,
        .client = upstream_client,
        .priority = priority,
        .enabled = true,
        .health = .init,
    };
}

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

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

/// Runs the client side under a budget so a server that never answers fails the
/// test instead of hanging the run.
fn bounded(io: std.Io, comptime f: anytype, args: std.meta.ArgsTuple(@TypeOf(f))) !void {
    var outcomes: [2]Outcome = undefined;
    var race: std.Io.Select(Outcome) = .init(io, &outcomes);
    defer race.cancelDiscard();

    try race.concurrent(.work, f, args);
    try race.concurrent(.expiry, expire, .{ io, budget_duration });

    switch (try race.await()) {
        .work => |result| return result,
        .expiry => |result| {
            try result;
            return error.TestTimedOut;
        },
    }
}

/// Everything the client can check about one answer: it belongs to the query it
/// was sent for, it succeeded, and it carries the fake upstream's A record.
fn expectAnswer(reply: []const u8, id: u16) !void {
    const p = try packet.parse(reply);
    try testing.expectEqual(id, p.header.id);
    try testing.expectEqual(true, p.header.flags.qr);
    try testing.expectEqual(types.Rcode.no_error, p.header.flags.rcode);
    try testing.expectEqual(@as(u16, 1), p.header.qdcount);
    try testing.expectEqual(@as(u16, 1), p.header.ancount);

    const echoed = packet.firstQuestion(p) orelse return error.TestMissingQuestion;
    const asked = packet.firstQuestion(try packet.parse(query_bytes)).?;
    try testing.expectEqualSlices(u8, asked.name.wire(), echoed.name.wire());
    try testing.expectEqual(types.Type.a, echoed.qtype);

    var it = packet.answers(p);
    const rec = (try it.next()) orelse return error.TestMissingAnswer;
    try testing.expectEqual(types.Type.a, rec.rtype);
    try testing.expectEqual(answer_ttl, rec.ttl);
    try testing.expectEqual(answer_rdata, try record.rdataA(reply, rec));
}

/// One length-prefixed query and its answer on a fresh connection
/// (RFC 1035 §4.2.2).
fn tcpQuery(io: std.Io, address: net.IpAddress, id: u16) anyerror!void {
    var stream = try address.connect(io, .{ .mode = .stream });
    defer stream.close(io);

    var read_buf: [1024]u8 = undefined;
    var write_buf: [1024]u8 = undefined;
    var reader = stream.reader(io, &read_buf);
    var writer = stream.writer(io, &write_buf);

    var query_buf: [query_bytes.len]u8 = undefined;
    const query = queryWithId(&query_buf, id);

    try writer.interface.writeAll(&transport.framePrefix(@intCast(query.len)));
    try writer.interface.writeAll(query);
    try writer.interface.flush();

    const len = transport.parsePrefix((try reader.interface.takeArray(transport.prefix_len)).*);
    try expectAnswer(try reader.interface.take(len), id);
}

test "the whole resolver answers over udp and tcp and fails over to a healthy upstream" {
    if (!build_options.integration) return error.SkipZigTest;

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

    var bad: FaultyUpstream = .{
        .fault = error.ConnectFailed,
        .fail_first = std.math.maxInt(u64),
    };
    var good: GoodUpstream = .{};
    var entries = [_]pool.Entry{
        testEntry("https://bad.example/dns-query", bad.client(), 10),
        testEntry("tls://good.example", good.client(), 20),
    };
    var upstreams: pool.Pool = .init(&entries, test_cfg, pool_timeouts, 1);

    var h = bareHandler(upstreams.client());

    const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
    var udp = try udp_server.UdpServer.bind(gpa, io, listen_address, &h, .{ .max_in_flight = 4 });
    var tcp = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{ .max_connections = 2 });
    const udp_address = udp.boundAddress();
    const tcp_address = tcp.boundAddress();

    var group: std.Io.Group = .init;
    try group.concurrent(io, udp_server.UdpServer.serve, .{ &udp, io });
    try group.concurrent(io, tcp_server.TcpServer.serve, .{ &tcp, io });

    const client_address: net.IpAddress = try .parse("127.0.0.1", 0);
    var client = try client_address.bind(io, .{ .mode = .dgram });
    defer client.close(io);

    {
        var query_buf: [query_bytes.len]u8 = undefined;
        const query = queryWithId(&query_buf, udp_id);
        try client.send(io, &udp_address, query);

        var buf: [udp_server.max_datagram]u8 = undefined;
        const msg = try client.receiveTimeout(io, &buf, budget);
        try expectAnswer(msg.data, udp_id);
    }

    try bounded(io, tcpQuery, .{ io, tcp_address, tcp_id });

    try testing.expectEqual(@as(u64, 2), h.stats.queries.load(.monotonic));
    try testing.expectEqual(@as(u64, 0), h.stats.servfail.load(.monotonic));

    var snapshots: [2]pool.Snapshot = undefined;
    try testing.expectEqual(@as(usize, 2), try upstreams.snapshot(io, &snapshots));

    try testing.expectEqualStrings("https://bad.example/dns-query", snapshots[0].url);
    try testing.expect(snapshots[0].consecutive_failures >= 1);
    try testing.expect(snapshots[0].backoff_until != null);
    try testing.expect(!snapshots[0].available);

    try testing.expectEqualStrings("tls://good.example", snapshots[1].url);
    try testing.expect(snapshots[1].total_successes >= 2);

    // The failing entry is in backoff, so the third query must reach the
    // healthy entry without touching it.
    const bad_calls = bad.calls.load(.monotonic);
    {
        var query_buf: [query_bytes.len]u8 = undefined;
        const query = queryWithId(&query_buf, backoff_id);
        try client.send(io, &udp_address, query);

        var buf: [udp_server.max_datagram]u8 = undefined;
        const msg = try client.receiveTimeout(io, &buf, budget);
        try expectAnswer(msg.data, backoff_id);
    }
    try testing.expectEqual(bad_calls, bad.calls.load(.monotonic));
    try testing.expectEqual(@as(u64, 3), good.calls.load(.monotonic));

    udp.deinit(gpa, io);
    tcp.deinit(io);
    group.cancel(io);
}