//! Loopback tests for `tcp_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: one listener and one client on 127.0.0.1 and an in-process fake
//! upstream. No stream read in 0.16.0 takes a timeout, so the whole client side
//! of each test 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 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 },
    };
}

const budget: std.Io.Clock.Duration = .{ .raw = .fromSeconds(5), .clock = .awake };

/// Short enough to keep the idle-timeout test quick, long enough that a
/// loopback connect cannot lose to scheduling and time out on its own.
const short_idle: std.Io.Clock.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://tcp-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 };
    }
};

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

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

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

/// RFC 7766 §6.2.1.1: two queries on one connection, answered in order.
fn twoQueriesOnOneConnection(io: std.Io, address: net.IpAddress) anyerror!void {
    const remote = address;
    var stream = try remote.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);

    for (0..2) |_| {
        try writer.interface.writeAll(&transport.framePrefix(@intCast(query_bytes.len)));
        try writer.interface.writeAll(query_bytes);
        try writer.interface.flush();

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

/// The server must close an idle connection on its own, which the client sees
/// as end of stream.
fn waitForServerClose(io: std.Io, address: net.IpAddress) anyerror!void {
    const remote = address;
    var stream = try remote.connect(io, .{ .mode = .stream });
    defer stream.close(io);

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

    var sink: [64]u8 = undefined;
    const n = try reader.interface.readSliceShort(&sink);
    if (n != 0) return error.TestUnexpectedBytes;
}

test "two length-prefixed queries share one connection" {
    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 tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{ .max_connections = 2 });
    const server_address = server.boundAddress();

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

    try bounded(io, twoQueriesOnOneConnection, .{ io, server_address });

    try testing.expectEqual(@as(u64, 1), server.core.stats.connections.load(.monotonic));
    try testing.expectEqual(@as(u64, 0), server.core.stats.rejected_at_capacity.load(.monotonic));
    try testing.expectEqual(@as(u64, 2), h.stats.queries.load(.monotonic));

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

test "the claimed slot records the connecting client" {
    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 tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{ .max_connections = 2 });
    const server_address = server.boundAddress();

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

    try bounded(io, twoQueriesOnOneConnection, .{ io, server_address });

    // The only connection took slot 0, and the reply the client already read
    // was written after `claim` filled the slot in, so this read races nothing.
    // Without a real peer the handler would rate-limit, group and log every TCP
    // client under whatever the uninitialized slot happened to hold.
    const peer = server.core.conns[0].peer;
    try testing.expectEqual(net.IpAddress.ip4, std.meta.activeTag(peer));
    try testing.expectEqualSlices(u8, &[_]u8{ 127, 0, 0, 1 }, &peer.ip4.bytes);
    try testing.expect(peer.ip4.port != 0);

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

/// `std.Io.Group` takes only `Cancelable!void`, so the result is reported out
/// of band. `answered` is set either way — a client that fails early must not
/// leave the test blocked waiting for a reply that will never come, and `set`
/// is documented to have no effect the second time.
fn holdConnection(
    io: std.Io,
    address: net.IpAddress,
    answered: *std.Io.Event,
    result: *anyerror!void,
) void {
    result.* = holdConnectionOpen(io, address, answered);
    answered.set(io);
}

/// Holds a connection open the way RFC 7766 lets a real client hold one: one
/// query answered, then nothing, so the server sits blocked on the read for the
/// next message. Returns once the server ends the connection, however it ends
/// it — a canceled server closes the socket, which the client sees either as
/// end of stream or as a reset.
fn holdConnectionOpen(io: std.Io, address: net.IpAddress, answered: *std.Io.Event) anyerror!void {
    const remote = address;
    var stream = try remote.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);

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

    const len = transport.parsePrefix((try reader.interface.takeArray(transport.prefix_len)).*);
    try expectAnswersQuery(try reader.interface.take(len));

    answered.set(io);

    var sink: [64]u8 = undefined;
    _ = reader.interface.readSliceShort(&sink) catch {};
}

test "a canceled serve does not wait for a live connection" {
    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);
    // The idle budget is the whole time a drain would have to wait out, so it
    // is set far beyond any patience this test run has: if `serve` waits for
    // the connection instead of canceling it, the wait never ends.
    var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{
        .max_connections = 2,
        .idle_timeout = .{ .raw = .fromSeconds(600), .clock = .awake },
    });
    const server_address = server.boundAddress();

    var serving: std.Io.Group = .init;
    try serving.concurrent(io, tcp_server.TcpServer.serve, .{ &server, io });

    var answered: std.Io.Event = .unset;
    var client_result: anyerror!void = {};
    var client_group: std.Io.Group = .init;
    try client_group.concurrent(io, holdConnection, .{ io, server_address, &answered, &client_result });

    // The reply proves the connection is claimed and served, so the connection
    // task is now blocked reading the message that never comes. That is the
    // state a drain hangs in.
    answered.waitUncancelable(io);

    // This is the composition root's shutdown: cancel the task group before
    // anything it borrows is released, so no `deinit` has shut the connection
    // down. It must still return. A regression here does not fail the test, it
    // hangs the run — there is no way to bound a join that does not finish.
    serving.cancel(io);

    // Cancellation must not skip the per-connection cleanup: the slot is
    // released and the socket closed by `serveConn`'s defer, which runs on the
    // canceled path like any other.
    try testing.expectEqual(@as(?usize, 0), firstFreeSlot(&server));

    client_group.cancel(io);
    server.deinit(io);

    // Checked last: the connection had to be answered for the test to mean
    // anything, and the server is torn down before a failure is reported.
    try client_result;
}

/// The first slot the server would hand out, read after `serve` has returned so
/// nothing can be writing it.
fn firstFreeSlot(server: *const tcp_server.TcpServer) ?usize {
    for (server.core.conns, 0..) |*conn, index| {
        if (conn.state == .free) return index;
    }
    return null;
}

test "an idle connection is closed and counted" {
    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 tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{
        .max_connections = 2,
        .idle_timeout = short_idle,
    });
    const server_address = server.boundAddress();

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

    try bounded(io, waitForServerClose, .{ io, server_address });

    try testing.expectEqual(@as(u64, 1), server.core.stats.connections.load(.monotonic));
    try testing.expectEqual(@as(u64, 1), server.core.stats.idle_timeouts.load(.monotonic));
    try testing.expectEqual(@as(u64, 0), server.core.stats.connection_errors.load(.monotonic));

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

test "a zero-length message is a connection error" {
    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 tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{
        .max_connections = 2,
        .idle_timeout = short_idle,
    });
    const server_address = server.boundAddress();

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

    try bounded(io, sendZeroLength, .{ io, server_address });

    try testing.expectEqual(@as(u64, 1), server.core.stats.connection_errors.load(.monotonic));
    try testing.expectEqual(@as(u64, 0), server.core.stats.idle_timeouts.load(.monotonic));

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

/// A prefix of 0 announces a message RFC 1035 §4.2.2 gives no meaning to, so
/// the server closes rather than waiting for bytes that will never mean
/// anything.
fn sendZeroLength(io: std.Io, address: net.IpAddress) anyerror!void {
    const remote = address;
    var stream = try remote.connect(io, .{ .mode = .stream });
    defer stream.close(io);

    var write_buf: [64]u8 = undefined;
    var writer = stream.writer(io, &write_buf);
    try writer.interface.writeAll(&transport.framePrefix(0));
    try writer.interface.flush();

    var read_buf: [64]u8 = undefined;
    var reader = stream.reader(io, &read_buf);
    var sink: [64]u8 = undefined;
    const n = try reader.interface.readSliceShort(&sink);
    if (n != 0) return error.TestUnexpectedBytes;
}

test "deinit ends a serve loop that is blocked on accept" {
    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 tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{ .max_connections = 2 });

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

    // No client ever connects, so `serve` is inside an accept when this runs.
    server.deinit(io);
    group.await(io) catch |err| switch (err) {
        error.Canceled => unreachable,
    };
}