//! 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 header = @import("../dns/header.zig");
const packet = @import("../dns/packet.zig");
const types = @import("../dns/types.zig");
const transport = @import("../upstream/transport.zig");

const testing = std.testing;

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,
    ) transport.ExchangeError![]u8 {
        _ = io;
        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: handler.Handler = .{ .upstream = 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.stats.accepted.load(.monotonic));
    try testing.expectEqual(@as(u64, 0), server.stats.rejected_at_capacity.load(.monotonic));
    try testing.expectEqual(@as(u64, 2), h.stats.queries.load(.monotonic));

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

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: handler.Handler = .{ .upstream = 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.stats.accepted.load(.monotonic));
    try testing.expectEqual(@as(u64, 1), server.stats.idle_timeouts.load(.monotonic));
    try testing.expectEqual(@as(u64, 0), server.stats.connection_errors.load(.monotonic));

    server.deinit(gpa, 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: handler.Handler = .{ .upstream = 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.stats.connection_errors.load(.monotonic));
    try testing.expectEqual(@as(u64, 0), server.stats.idle_timeouts.load(.monotonic));

    server.deinit(gpa, 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: handler.Handler = .{ .upstream = 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(gpa, io);
    group.await(io) catch |err| switch (err) {
        error.Canceled => unreachable,
    };
}