//! Loopback tests for `udp_server.zig`.
//!
//! 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: two sockets on 127.0.0.1 and an in-process fake upstream. Nothing
//! leaves the machine, and every wait carries a budget.

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

const handler = @import("handler.zig");
const udp_server = @import("udp_server.zig");
const model = @import("../config/model.zig");
const response = @import("../filter/response.zig");
const header = @import("../dns/header.zig");
const packet = @import("../dns/packet.zig");
const types = @import("../dns/types.zig");
const transport = @import("../upstream/transport.zig");
const upstream_owner = @import("../upstream/owner.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
/// listener is what these tests exercise, so the handler is the same bare one
/// its own tests use.
fn bareHandler(up: *upstream_owner.Owner) handler.Handler {
    return .{
        .upstream = up,
        .policy = .{ .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 } };

/// How long the test waits to be convinced that no reply is coming. A dropped
/// datagram produces nothing, so this budget is spent in full on every run.
const silence: std.Io.Timeout = .{ .duration = .{ .raw = .fromMilliseconds(300), .clock = .awake } };

/// A query for example.com A: id 0x1234, RD set, one question, no OPT.
const query_bytes =
    "\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00" ++
    "\x07example\x03com\x00\x00\x01\x00\x01";

/// The matching response: the question echoed plus one A record.
const response_bytes =
    "\x12\x34\x81\x80\x00\x01\x00\x01\x00\x00\x00\x00" ++
    "\x07example\x03com\x00\x00\x01\x00\x01" ++
    "\xc0\x0c\x00\x01\x00\x01\x00\x00\x01\x2c\x00\x04\x5d\xb8\xd8\x22";

/// Answers from a fixture and rewrites the ID, which is all the server needs
/// from an upstream. The real clients are exercised by their own tests.
const FakeUpstream = struct {
    reply: []const u8,

    fn exchangeFn(
        ptr: *anyopaque,
        io: std.Io,
        query: []const u8,
        response_buf: []u8,
        selected: *?[]const u8,
    ) transport.ExchangeError![]u8 {
        _ = io;
        selected.* = "fake://udp-server-upstream";
        const self: *FakeUpstream = @ptrCast(@alignCast(ptr));
        if (self.reply.len > response_buf.len) return error.ResponseTooLarge;
        @memcpy(response_buf[0..self.reply.len], self.reply);
        const bytes = response_buf[0..self.reply.len];
        packet.setId(bytes, (header.parse(query) catch unreachable).id);
        return bytes;
    }

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

fn expectAnswersQuery(reply: []const u8) !void {
    const p = try packet.parse(reply);
    try testing.expectEqual(@as(u16, 0x1234), 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);

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

test "a udp query is answered on the loopback" {
    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 fake: FakeUpstream = .{ .reply = response_bytes };
    var h_owner: upstream_owner.Borrowed = .{};
    var h = bareHandler(h_owner.client(fake.client()));

    const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
    var server = try udp_server.UdpServer.bind(gpa, io, listen_address, &h, .{ .max_in_flight = 4 });
    const server_address = server.boundAddress();

    var group: std.Io.Group = .init;
    try group.concurrent(io, udp_server.UdpServer.serve, .{ &server, 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);

    try client.send(io, &server_address, query_bytes);

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

    try testing.expectEqual(@as(u64, 1), server.stats.received.load(.monotonic));
    try testing.expectEqual(@as(u64, 0), server.stats.send_errors.load(.monotonic));
    try testing.expectEqual(@as(u64, 1), h.stats.queries.load(.monotonic));

    server.deinit(gpa, io);
    group.await(io) catch |err| switch (err) {
        error.Canceled => unreachable,
    };
}

test "a runt datagram is dropped and no reply is sent" {
    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 fake: FakeUpstream = .{ .reply = response_bytes };
    var h_owner: upstream_owner.Borrowed = .{};
    var h = bareHandler(h_owner.client(fake.client()));

    const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
    var server = try udp_server.UdpServer.bind(gpa, io, listen_address, &h, .{ .max_in_flight = 4 });
    const server_address = server.boundAddress();

    var group: std.Io.Group = .init;
    try group.concurrent(io, udp_server.UdpServer.serve, .{ &server, 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);

    try client.send(io, &server_address, query_bytes[0..5]);

    var buf: [udp_server.max_datagram]u8 = undefined;
    try testing.expectError(error.Timeout, client.receiveTimeout(io, &buf, silence));
    try testing.expectEqual(@as(u64, 1), server.stats.received.load(.monotonic));
    try testing.expectEqual(@as(u64, 1), h.stats.dropped_malformed.load(.monotonic));

    server.deinit(gpa, io);
    group.await(io) catch |err| switch (err) {
        error.Canceled => unreachable,
    };
}

test "an oversize datagram arrives truncated and is dropped" {
    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 fake: FakeUpstream = .{ .reply = response_bytes };
    var h_owner: upstream_owner.Borrowed = .{};
    var h = bareHandler(h_owner.client(fake.client()));

    const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
    var server = try udp_server.UdpServer.bind(gpa, io, listen_address, &h, .{ .max_in_flight = 4 });
    const server_address = server.boundAddress();

    var group: std.Io.Group = .init;
    try group.concurrent(io, udp_server.UdpServer.serve, .{ &server, 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);

    // The server receives into a `max_datagram` buffer, so the kernel discards
    // the tail of this one and reports it through `flags.trunc`.
    const oversize = try gpa.alloc(u8, udp_server.max_datagram * 2);
    defer gpa.free(oversize);
    @memcpy(oversize[0..query_bytes.len], query_bytes);
    @memset(oversize[query_bytes.len..], 0);
    try client.send(io, &server_address, oversize);

    var buf: [udp_server.max_datagram]u8 = undefined;
    try testing.expectError(error.Timeout, client.receiveTimeout(io, &buf, silence));
    try testing.expectEqual(@as(u64, 1), server.stats.dropped_oversize.load(.monotonic));
    try testing.expectEqual(@as(u64, 0), h.stats.queries.load(.monotonic));

    server.deinit(gpa, io);
    group.await(io) catch |err| switch (err) {
        error.Canceled => unreachable,
    };
}

test "deinit ends a serve loop that is blocked on receive" {
    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 fake: FakeUpstream = .{ .reply = response_bytes };
    var h_owner: upstream_owner.Borrowed = .{};
    var h = bareHandler(h_owner.client(fake.client()));

    const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
    var server = try udp_server.UdpServer.bind(gpa, io, listen_address, &h, .{ .max_in_flight = 4 });

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

    // No datagram ever arrives, so `serve` is inside a receive when this runs.
    server.deinit(gpa, io);
    group.await(io) catch |err| switch (err) {
        error.Canceled => unreachable,
    };
}
