//! Whole DNS messages (RFC 1035 §4.1): structural validation, section
//! iteration, in-place mutation of a raw buffer, and a builder for synthesized
//! replies. Pure: no allocation, no `std.Io` beyond writing encoded bytes to a
//! caller's buffer.
//!
//! `parse` separates two failure classes because the server answers them
//! differently (PLAN §6.1):
//!
//!   - `error.Truncated` — fewer than the 12 header bytes arrived. There is no
//!     ID and no question to echo, so no reply is possible: the caller drops
//!     the datagram silently.
//!   - `WalkError` — the header is intact but a section is malformed. The
//!     caller can echo the ID and answer FORMERR.
//!
//! `name.ParseError.Truncated` therefore cannot pass through unchanged: a name
//! that runs off the end of a packet is a malformed section, not a short
//! header, and mapping it onto `error.Truncated` would turn a FORMERR into a
//! silent drop. `mapNameError` performs that translation once.

const std = @import("std");
const types = @import("types.zig");
const header = @import("header.zig");
const name = @import("name.zig");
const question = @import("question.zig");
const record = @import("record.zig");
const edns = @import("edns.zig");
const Writer = std.Io.Writer;

/// A malformed section in a packet whose header is intact.
pub const WalkError = error{
    /// A name is undecodable: a bad compression pointer, an over-long label or
    /// name, or too many pointer jumps.
    BadName,
    /// A question or record runs past the end of the packet. A count field
    /// larger than the records actually present lands here too, because the
    /// walk then reads past the last record.
    SectionOverrun,
    /// Bytes follow the last counted record. RFC 1035 §4.1 gives a message no
    /// padding outside its four sections, and EDNS padding (RFC 7830) lives
    /// inside the OPT record's option list, so nothing legitimate lands here.
    TrailingBytes,
    /// More than one OPT record. RFC 6891 §6.1.1 allows at most one per
    /// message and requires FORMERR for the rest.
    MultipleOptRecords,
};

pub const ParseError = error{Truncated} || WalkError;

fn mapNameError(err: name.ParseError) WalkError {
    return switch (err) {
        error.Truncated => error.SectionOverrun,
        error.LabelTooLong, error.NameTooLong, error.BadPointer, error.TooManyJumps => error.BadName,
    };
}

/// A validated view of a message: the bytes plus the decoded header. Sections
/// are not stored — iterators re-walk the bytes on demand, which keeps the
/// struct small enough to pass by value and free of pointers into itself.
pub const Packet = struct {
    bytes: []const u8,
    header: header.Header,
};

/// Decodes the header and walks every section once, bounds-checking each name
/// and record. The walk must land exactly on the end of the buffer and must
/// meet at most one OPT record, so a `Packet` from here has no bytes outside
/// its sections and at most one OPT.
pub fn parse(bytes: []const u8) ParseError!Packet {
    const h = header.parse(bytes) catch |err| switch (err) {
        error.Truncated => return error.Truncated,
    };

    var pos: usize = types.header_len;
    var q: u16 = 0;
    while (q < h.qdcount) : (q += 1) {
        const parsed = question.parse(bytes, pos) catch |err| return mapNameError(err);
        pos = parsed.end;
    }

    const record_count = recordCount(h);
    var opt_seen = false;
    var r: u32 = 0;
    while (r < record_count) : (r += 1) {
        const parsed = record.parse(bytes, pos) catch |err| return mapNameError(err);
        pos = parsed.end;
        if (parsed.record.rtype == .opt) {
            if (opt_seen) return error.MultipleOptRecords;
            opt_seen = true;
        }
    }

    if (pos != bytes.len) return error.TrailingBytes;
    return .{ .bytes = bytes, .header = h };
}

fn recordCount(h: header.Header) u32 {
    return @as(u32, h.ancount) + @as(u32, h.nscount) + @as(u32, h.arcount);
}

pub const QuestionIterator = struct {
    packet: []const u8,
    pos: usize,
    remaining: u16,

    pub fn next(self: *QuestionIterator) WalkError!?question.Question {
        if (self.remaining == 0) return null;
        const parsed = question.parse(self.packet, self.pos) catch |err| return mapNameError(err);
        self.pos = parsed.end;
        self.remaining -= 1;
        return parsed.question;
    }
};

/// Records of one section. The section's first byte is only reachable by
/// walking everything before it, so the first `next` call does that walk; the
/// iterator is otherwise inert.
pub const RecordIterator = struct {
    packet: []const u8,
    skip_questions: u16,
    skip_records: u32,
    remaining: u16,
    pos: usize,
    positioned: bool,

    pub fn next(self: *RecordIterator) WalkError!?record.Record {
        if (!self.positioned) try self.position();
        if (self.remaining == 0) return null;
        const parsed = record.parse(self.packet, self.pos) catch |err| return mapNameError(err);
        self.pos = parsed.end;
        self.remaining -= 1;
        return parsed.record;
    }

    fn position(self: *RecordIterator) WalkError!void {
        var pos: usize = types.header_len;
        var q: u16 = 0;
        while (q < self.skip_questions) : (q += 1) {
            const parsed = question.parse(self.packet, pos) catch |err| return mapNameError(err);
            pos = parsed.end;
        }
        var r: u32 = 0;
        while (r < self.skip_records) : (r += 1) {
            const parsed = record.parse(self.packet, pos) catch |err| return mapNameError(err);
            pos = parsed.end;
        }
        self.pos = pos;
        self.positioned = true;
    }
};

pub fn questions(p: Packet) QuestionIterator {
    return .{ .packet = p.bytes, .pos = types.header_len, .remaining = p.header.qdcount };
}

pub fn answers(p: Packet) RecordIterator {
    return recordIterator(p, 0, p.header.ancount);
}

pub fn authorities(p: Packet) RecordIterator {
    return recordIterator(p, p.header.ancount, p.header.nscount);
}

pub fn additionals(p: Packet) RecordIterator {
    return recordIterator(p, recordCount(p.header) - p.header.arcount, p.header.arcount);
}

fn recordIterator(p: Packet, skip_records: u32, remaining: u16) RecordIterator {
    return .{
        .packet = p.bytes,
        .skip_questions = p.header.qdcount,
        .skip_records = skip_records,
        .remaining = remaining,
        .pos = types.header_len,
        .positioned = false,
    };
}

/// The first question, or null when there is none. A `Packet` from `parse` has
/// a decodable question section, so the parse below only fails for a `Packet`
/// assembled by hand around unvalidated bytes.
pub fn firstQuestion(p: Packet) ?question.Question {
    var it = questions(p);
    return it.next() catch null;
}

/// The OPT record from the additional section, or null when there is none. A
/// `Packet` from `parse` carries at most one, so the search below finds either
/// nothing or that one record. The signature also accepts a `Packet` assembled
/// by hand around unvalidated bytes; for those the last OPT wins, which is what
/// a sequential walk of the section yields.
pub fn findOptRecord(p: Packet) ?record.Record {
    var it = additionals(p);
    var found: ?record.Record = null;
    while (it.next() catch return found) |rec| {
        if (rec.rtype == .opt) found = rec;
    }
    return found;
}

/// Overwrites the message ID in place. Used on the cache hit path, where a
/// stored response answers a new query.
pub fn setId(bytes: []u8, id: u16) void {
    std.debug.assert(bytes.len >= types.header_len);
    std.mem.writeInt(u16, bytes[0..2], id, .big);
}

/// TTL sits four bytes into the fixed fields and RDLENGTH's two bytes follow
/// it, so the TTL starts six bytes before the RDATA.
const ttl_bytes_before_rdata = 6;

/// Ages every record in place by `elapsed_seconds`, saturating at zero, and
/// returns the smallest resulting TTL — or null when the message carries no
/// record whose TTL means anything. The caller decides what a small or absent
/// TTL means; this function only does the arithmetic.
///
/// OPT records are skipped: RFC 6891 §6.1.3 reuses their TTL field for the
/// extended RCODE, the EDNS version and the DO bit, so subtracting from it
/// would corrupt the flags.
///
/// A TTL with its top bit set is a wire-level oddity that RFC 2181 §8 says to
/// treat as zero. That is a caching decision, so it stays with the caller and
/// this function ages such a value like any other.
///
/// A full structural validation runs before any byte changes, so malformed
/// input leaves the buffer exactly as it arrived. Aging and validating in one
/// pass would age the records before a malformed one and then report an error.
///
/// One case still ends mid-way: an owner name may be a compression pointer into
/// an earlier record's TTL field, and aging that TTL can make the name
/// undecodable. The walk below therefore keeps checking each record instead of
/// trusting the validation above. On any error the buffer holds a partly aged
/// message that no longer parses, so a caller that gets an error must discard
/// the buffer rather than send it.
pub fn decrementTtls(bytes: []u8, elapsed_seconds: u32) ParseError!?u32 {
    const p = try parse(bytes);

    var pos: usize = types.header_len;
    var q: u16 = 0;
    while (q < p.header.qdcount) : (q += 1) {
        const parsed = question.parse(bytes, pos) catch |err| return mapNameError(err);
        pos = parsed.end;
    }

    var minimum: ?u32 = null;
    const record_count = recordCount(p.header);
    var r: u32 = 0;
    while (r < record_count) : (r += 1) {
        const parsed = record.parse(bytes, pos) catch |err| return mapNameError(err);
        pos = parsed.end;
        if (parsed.record.rtype == .opt) continue;

        const aged = parsed.record.ttl -| elapsed_seconds;
        const ttl_offset = parsed.record.rdata.offset - ttl_bytes_before_rdata;
        std.mem.writeInt(u32, bytes[ttl_offset..][0..4], aged, .big);
        minimum = if (minimum) |m| @min(m, aged) else aged;
    }
    return minimum;
}

/// Encodes a reply into a caller buffer. Mechanism only: which RCODE to set
/// and which answers to add is the caller's policy.
///
/// Sections must be filled in wire order, so every `addAnswer` call has to
/// precede the OPT record — it belongs to the additional section, and an answer
/// written after it would land in the wrong section. Only one of `addOptEcho`
/// and `addOptWithRcode` runs, and only once, because RFC 6891 §6.1.1 allows
/// one OPT record per message. Both rules are programmer errors, so both are
/// assertions.
pub const ResponseBuilder = struct {
    writer: Writer,
    header: header.Header,
    opt_added: bool,

    pub const Error = Writer.Error || error{RdataTooLong};

    /// Copies the request's ID, opcode and RD bit, marks the message a
    /// response, and advertises recursion. The question is echoed when given,
    /// as RFC 1035 §4.1.2 expects of a reply.
    pub fn init(buf: []u8, request: header.Header, q: ?question.Question) Error!ResponseBuilder {
        // A DNS message carries a two-byte length prefix over TCP
        // (RFC 1035 §4.2.2), so 65535 bytes is the protocol maximum. Holding
        // the buffer to it keeps the section counters from overflowing: the
        // smallest possible record is 11 bytes.
        std.debug.assert(buf.len <= std.math.maxInt(u16));

        var self: ResponseBuilder = .{
            .writer = Writer.fixed(buf),
            .header = .{
                .id = request.id,
                .flags = .{
                    .rcode = .no_error,
                    .z = 0,
                    .ra = true,
                    .rd = request.flags.rd,
                    .tc = false,
                    .aa = false,
                    .opcode = request.flags.opcode,
                    .qr = true,
                },
                .qdcount = 0,
                .ancount = 0,
                .nscount = 0,
                .arcount = 0,
            },
            .opt_added = false,
        };

        var placeholder: [types.header_len]u8 = undefined;
        header.encode(self.header, &placeholder);
        try self.writer.writeAll(&placeholder);

        if (q) |echoed| {
            try question.encode(echoed, &self.writer);
            self.header.qdcount = 1;
        }
        return self;
    }

    pub fn setRcode(self: *ResponseBuilder, rcode: types.Rcode) void {
        self.header.flags.rcode = rcode;
    }

    pub fn setAuthoritative(self: *ResponseBuilder, aa: bool) void {
        self.header.flags.aa = aa;
    }

    /// `rdata` is written verbatim, so it must hold no compression pointers.
    pub fn addAnswer(
        self: *ResponseBuilder,
        owner: name.Name,
        rtype: types.Type,
        class: types.Class,
        ttl: u32,
        rdata: []const u8,
    ) Error!void {
        std.debug.assert(!self.opt_added);
        const rec: record.Record = .{
            .name = owner,
            .rtype = rtype,
            .class = @intFromEnum(class),
            .ttl = ttl,
            .rdata = .{ .offset = 0, .len = 0 },
        };
        try record.encode(rec, rdata, &self.writer);
        self.header.ancount += 1;
    }

    /// Answers an EDNS query with an EDNS reply: the requestor's payload size
    /// comes back unchanged and the DO bit passes through. No options are
    /// echoed — nxdns implements none of them.
    pub fn addOptEcho(self: *ResponseBuilder, request_opt: edns.OptRecord, do_bit: bool) Error!void {
        return self.addOptWithRcode(request_opt, do_bit, 0);
    }

    /// `addOptEcho` with the upper eight bits of a 12-bit RCODE, which only the
    /// OPT record can carry (RFC 6891 §6.1.3). Pair it with the header half from
    /// `edns.splitRcode`. The reply's EDNS version stays 0 either way: it states
    /// the highest version this server implements, not the version asked for.
    pub fn addOptWithRcode(
        self: *ResponseBuilder,
        request_opt: edns.OptRecord,
        do_bit: bool,
        extended_rcode: u8,
    ) Error!void {
        std.debug.assert(!self.opt_added);
        const opt: edns.OptRecord = .{
            .udp_payload_size = request_opt.udp_payload_size,
            .extended_rcode = extended_rcode,
            .version = 0,
            .do_bit = do_bit,
            .options = .{ .offset = 0, .len = 0 },
        };
        edns.encodeOpt(opt, &.{}, &self.writer) catch |err| switch (err) {
            error.OptionsTooLong => unreachable, // the option list is empty
            error.WriteFailed => return error.WriteFailed,
        };
        self.header.arcount += 1;
        self.opt_added = true;
    }

    /// Patches the counts into the reserved header bytes and returns the
    /// finished message, a prefix of the caller's buffer.
    pub fn finish(self: *ResponseBuilder) []u8 {
        const bytes = self.writer.buffered();
        header.encode(self.header, bytes[0..types.header_len]);
        return bytes;
    }
};

const testing = std.testing;

/// A query for example.com A with an EDNS(0) OPT record advertising 4096
/// bytes: id 0x1234, RD set, one question, one additional.
const query_bytes =
    "\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x01" ++ // header
    "\x07example\x03com\x00\x00\x01\x00\x01" ++ // question at 12, ends at 29
    "\x00\x00\x29\x10\x00\x00\x00\x00\x00\x00\x00"; // OPT at 29, ends at 40

/// The matching response: a CNAME to www.example.com and its A record, both
/// with compressed owner names, plus the echoed OPT record.
///   12 question, 29 CNAME (rdata at 41, ttl at 35),
///   47 A (rdata at 59, ttl at 53), 63 OPT, 74 end.
const response_bytes =
    "\x12\x34\x81\x80\x00\x01\x00\x02\x00\x00\x00\x01" ++ // header
    "\x07example\x03com\x00\x00\x01\x00\x01" ++ // question at 12
    "\xc0\x0c\x00\x05\x00\x01\x00\x00\x01\x2c\x00\x06\x03www\xc0\x0c" ++ // CNAME, ttl 300
    "\xc0\x29\x00\x01\x00\x01\x00\x00\x00\x3c\x00\x04\x5d\xb8\xd8\x22" ++ // A, ttl 60
    "\x00\x00\x29\x10\x00\x00\x00\x00\x00\x00\x00"; // OPT at 63

test "fixtures have the documented layout" {
    try testing.expectEqual(@as(usize, 40), query_bytes.len);
    try testing.expectEqual(@as(usize, 74), response_bytes.len);
}

test "parse a query" {
    const p = try parse(query_bytes);
    try testing.expectEqual(@as(u16, 0x1234), p.header.id);
    try testing.expectEqual(false, p.header.flags.qr);
    try testing.expectEqual(true, p.header.flags.rd);
    try testing.expectEqual(@as(u16, 1), p.header.qdcount);
    try testing.expectEqual(@as(u16, 1), p.header.arcount);

    const q = firstQuestion(p).?;
    try testing.expectEqualSlices(u8, (try name.fromText("example.com")).wire(), q.name.wire());
    try testing.expectEqual(types.Type.a, q.qtype);
    try testing.expectEqual(types.Class.in, q.qclass);
}

test "parse a response and walk every section" {
    const p = try parse(response_bytes);
    try testing.expectEqual(true, p.header.flags.qr);
    try testing.expectEqual(@as(u16, 2), p.header.ancount);

    var qit = questions(p);
    const q = (try qit.next()).?;
    try testing.expectEqual(types.Type.a, q.qtype);
    try testing.expectEqual(@as(?question.Question, null), try qit.next());

    var ait = answers(p);
    const cname = (try ait.next()).?;
    try testing.expectEqual(types.Type.cname, cname.rtype);
    try testing.expectEqual(@as(u32, 300), cname.ttl);
    try testing.expectEqualSlices(
        u8,
        (try name.fromText("example.com")).wire(),
        cname.name.wire(),
    );
    try testing.expectEqualSlices(
        u8,
        (try name.fromText("www.example.com")).wire(),
        (try record.rdataCname(response_bytes, cname)).wire(),
    );

    const a = (try ait.next()).?;
    try testing.expectEqual(types.Type.a, a.rtype);
    try testing.expectEqual(@as(u32, 60), a.ttl);
    try testing.expectEqualSlices(
        u8,
        (try name.fromText("www.example.com")).wire(),
        a.name.wire(),
    );
    try testing.expectEqual([4]u8{ 93, 184, 216, 34 }, try record.rdataA(response_bytes, a));
    try testing.expectEqual(@as(?record.Record, null), try ait.next());

    var nit = authorities(p);
    try testing.expectEqual(@as(?record.Record, null), try nit.next());

    var dit = additionals(p);
    const opt_rec = (try dit.next()).?;
    try testing.expectEqual(types.Type.opt, opt_rec.rtype);
    try testing.expectEqual(@as(?record.Record, null), try dit.next());
}

test "iterators are independent and re-walk on demand" {
    const p = try parse(response_bytes);
    var first = answers(p);
    var second = answers(p);
    const a1 = (try first.next()).?;
    const a2 = (try second.next()).?;
    try testing.expectEqual(a1.rdata.offset, a2.rdata.offset);
    _ = try first.next();
    const a2_second = (try second.next()).?;
    try testing.expectEqual(types.Type.a, a2_second.rtype);
}

test "authorities and additionals skip the sections before them" {
    // One question, one answer, one authority, one additional, all with root
    // owner names and 4-byte A rdata.
    const rec = "\x00\x00\x01\x00\x01\x00\x00\x00\x0a\x00\x04";
    const bytes = "\x00\x01\x81\x80\x00\x01\x00\x01\x00\x01\x00\x01" ++
        "\x00\x00\x01\x00\x01" ++ // question: root A IN
        rec ++ "\x01\x01\x01\x01" ++
        rec ++ "\x02\x02\x02\x02" ++
        rec ++ "\x03\x03\x03\x03";

    const p = try parse(bytes);
    var ait = answers(p);
    try testing.expectEqual([4]u8{ 1, 1, 1, 1 }, try record.rdataA(bytes, (try ait.next()).?));
    try testing.expectEqual(@as(?record.Record, null), try ait.next());

    var nit = authorities(p);
    try testing.expectEqual([4]u8{ 2, 2, 2, 2 }, try record.rdataA(bytes, (try nit.next()).?));
    try testing.expectEqual(@as(?record.Record, null), try nit.next());

    var dit = additionals(p);
    try testing.expectEqual([4]u8{ 3, 3, 3, 3 }, try record.rdataA(bytes, (try dit.next()).?));
    try testing.expectEqual(@as(?record.Record, null), try dit.next());
}

test "parse reports a short header as truncated" {
    var i: usize = 0;
    while (i < types.header_len) : (i += 1) {
        try testing.expectError(error.Truncated, parse(query_bytes[0..i]));
    }
    // Exactly a header with no sections is well-formed.
    const empty = try parse("\x00\x01\x81\x83" ++ "\x00" ** 8);
    try testing.expectEqual(@as(u16, 0), empty.header.qdcount);
    try testing.expectEqual(@as(?question.Question, null), firstQuestion(empty));
}

test "parse reports a truncated section as an overrun, not as truncation" {
    // The header promises a question that is not there.
    try testing.expectError(
        error.SectionOverrun,
        parse("\x00\x01\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00"),
    );
    // Every prefix that keeps the header but cuts a section.
    var i: usize = types.header_len;
    while (i < query_bytes.len) : (i += 1) {
        try testing.expectError(error.SectionOverrun, parse(query_bytes[0..i]));
    }
}

test "parse rejects a count larger than the records present" {
    // ancount claims two answers, one follows.
    const bytes = "\x00\x01\x81\x80\x00\x00\x00\x02\x00\x00\x00\x00" ++
        "\x00\x00\x01\x00\x01\x00\x00\x00\x0a\x00\x04\x01\x02\x03\x04";
    try testing.expectError(error.SectionOverrun, parse(bytes));

    // qdcount claims two questions, one follows.
    const two_questions = "\x00\x01\x01\x00\x00\x02\x00\x00\x00\x00\x00\x00" ++
        "\x00\x00\x01\x00\x01";
    try testing.expectError(error.SectionOverrun, parse(two_questions));
}

test "parse rejects an rdlength that overruns the packet" {
    const bytes = "\x00\x01\x81\x80\x00\x00\x00\x01\x00\x00\x00\x00" ++
        "\x00\x00\x01\x00\x01\x00\x00\x00\x0a\x00\x40\x01\x02\x03\x04";
    try testing.expectError(error.SectionOverrun, parse(bytes));
}

test "parse rejects a bad name in the question section" {
    // A pointer to itself: no chain can terminate.
    const loop = "\x00\x01\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00" ++
        "\xc0\x0c\x00\x01\x00\x01";
    try testing.expectError(error.BadName, parse(loop));

    // A forward pointer.
    const forward = "\x00\x01\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00" ++
        "\xc0\x12\x00\x01\x00\x01\x03com\x00";
    try testing.expectError(error.BadName, parse(forward));

    // A reserved label type, which is also what a 64-byte label looks like.
    const reserved = "\x00\x01\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00" ++
        "\x40abc\x00\x00\x01\x00\x01";
    try testing.expectError(error.BadName, parse(reserved));
}

test "parse rejects a bad name in a record" {
    const bytes = "\x00\x01\x81\x80\x00\x00\x00\x01\x00\x00\x00\x00" ++
        "\xc0\x0c\x00\x01\x00\x01\x00\x00\x00\x0a\x00\x04\x01\x02\x03\x04";
    try testing.expectError(error.BadName, parse(bytes));
}

test "parse rejects trailing bytes after the last record" {
    try testing.expectError(error.TrailingBytes, parse(query_bytes ++ "\xff\xff\xff"));
    // A single trailing byte counts too.
    try testing.expectError(error.TrailingBytes, parse(query_bytes ++ "\x00"));
    // A header-only message with a byte after it.
    try testing.expectError(error.TrailingBytes, parse("\x00\x01\x81\x83" ++ "\x00" ** 9));
}

test "parse rejects a second OPT record" {
    // Two OPT records in the additional section, payload sizes 512 and 1232.
    const two_opts = "\x00\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x02" ++
        "\x00\x00\x29\x02\x00\x00\x00\x00\x00\x00\x00" ++
        "\x00\x00\x29\x04\xd0\x00\x00\x00\x00\x00\x00";
    try testing.expectError(error.MultipleOptRecords, parse(two_opts));

    // One OPT in the answer section and one in the additional section: still
    // two OPT records in one message.
    const split = "\x00\x01\x01\x00\x00\x00\x00\x01\x00\x00\x00\x01" ++
        "\x00\x00\x29\x02\x00\x00\x00\x00\x00\x00\x00" ++
        "\x00\x00\x29\x04\xd0\x00\x00\x00\x00\x00\x00";
    try testing.expectError(error.MultipleOptRecords, parse(split));

    // One OPT beside a non-OPT additional record stays acceptable.
    const single = "\x00\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x02" ++
        "\x00\x00\x01\x00\x01\x00\x00\x00\x0a\x00\x04\x01\x02\x03\x04" ++
        "\x00\x00\x29\x10\x00\x00\x00\x00\x00\x00\x00";
    try testing.expectEqual(@as(u16, 2), (try parse(single)).header.arcount);
}

test "findOptRecord finds the OPT record and reads it" {
    const p = try parse(query_bytes);
    const rec = findOptRecord(p).?;
    const opt = try edns.parseOpt(query_bytes, rec);
    try testing.expectEqual(@as(u16, 4096), opt.udp_payload_size);
    try testing.expectEqual(false, opt.do_bit);

    const no_opt = try parse("\x00\x01\x01\x00" ++ "\x00" ** 8);
    try testing.expectEqual(@as(?record.Record, null), findOptRecord(no_opt));
}

test "findOptRecord skips a non-OPT additional record" {
    const bytes = "\x00\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x02" ++
        "\x00\x00\x01\x00\x01\x00\x00\x00\x0a\x00\x04\x01\x02\x03\x04" ++
        "\x00\x00\x29\x10\x00\x00\x00\x00\x00\x00\x00";
    const p = try parse(bytes);
    const opt = try edns.parseOpt(bytes, findOptRecord(p).?);
    try testing.expectEqual(@as(u16, 4096), opt.udp_payload_size);
}

test "setId overwrites the id and nothing else" {
    var bytes: [query_bytes.len]u8 = query_bytes.*;
    setId(&bytes, 0xbeef);
    try testing.expectEqual(@as(u16, 0xbeef), (try parse(&bytes)).header.id);
    try testing.expectEqualSlices(u8, query_bytes[2..], bytes[2..]);
}

test "decrementTtls ages every record and returns the minimum" {
    var bytes: [response_bytes.len]u8 = response_bytes.*;
    const minimum = (try decrementTtls(&bytes, 10)).?;
    try testing.expectEqual(@as(u32, 50), minimum);

    const p = try parse(&bytes);
    var it = answers(p);
    try testing.expectEqual(@as(u32, 290), (try it.next()).?.ttl);
    try testing.expectEqual(@as(u32, 50), (try it.next()).?.ttl);
}

test "decrementTtls saturates at zero" {
    var bytes: [response_bytes.len]u8 = response_bytes.*;
    try testing.expectEqual(@as(?u32, 0), try decrementTtls(&bytes, 1_000_000));

    const p = try parse(&bytes);
    var it = answers(p);
    try testing.expectEqual(@as(u32, 0), (try it.next()).?.ttl);
    try testing.expectEqual(@as(u32, 0), (try it.next()).?.ttl);
}

test "decrementTtls leaves the OPT flags alone" {
    // OPT with DO set and a 4096-byte payload size, so its TTL word is
    // 0x0000_8000 — a plain subtraction would clear the DO bit.
    const bytes_const = "\x12\x34\x81\x80\x00\x00\x00\x01\x00\x00\x00\x01" ++
        "\x00\x00\x01\x00\x01\x00\x00\x00\x64\x00\x04\x01\x02\x03\x04" ++
        "\x00\x00\x29\x10\x00\x00\x00\x80\x00\x00\x00";
    var bytes: [bytes_const.len]u8 = bytes_const.*;

    try testing.expectEqual(@as(?u32, 40), try decrementTtls(&bytes, 60));

    const p = try parse(&bytes);
    const opt = try edns.parseOpt(&bytes, findOptRecord(p).?);
    try testing.expectEqual(true, opt.do_bit);
    try testing.expectEqual(@as(u16, 4096), opt.udp_payload_size);
    try testing.expectEqual(@as(u8, 0), opt.extended_rcode);
}

test "decrementTtls ages authority and additional records too" {
    const rec = "\x00\x00\x01\x00\x01";
    const bytes_const = "\x00\x01\x81\x80\x00\x00\x00\x01\x00\x01\x00\x01" ++
        rec ++ "\x00\x00\x00\x64\x00\x04\x01\x01\x01\x01" ++ // ttl 100
        rec ++ "\x00\x00\x00\x1e\x00\x04\x02\x02\x02\x02" ++ // ttl 30
        rec ++ "\x00\x00\x00\x50\x00\x04\x03\x03\x03\x03"; // ttl 80
    var bytes: [bytes_const.len]u8 = bytes_const.*;

    try testing.expectEqual(@as(?u32, 10), try decrementTtls(&bytes, 20));

    const p = try parse(&bytes);
    var ait = answers(p);
    try testing.expectEqual(@as(u32, 80), (try ait.next()).?.ttl);
    var nit = authorities(p);
    try testing.expectEqual(@as(u32, 10), (try nit.next()).?.ttl);
    var dit = additionals(p);
    try testing.expectEqual(@as(u32, 60), (try dit.next()).?.ttl);
}

test "decrementTtls reports no minimum when nothing carries a ttl" {
    var query: [query_bytes.len]u8 = query_bytes.*;
    // The only record is the OPT, which is skipped.
    try testing.expectEqual(@as(?u32, null), try decrementTtls(&query, 5));
    try testing.expectEqualSlices(u8, query_bytes, &query);
}

test "decrementTtls propagates structural errors" {
    var short = [_]u8{0} ** 8;
    try testing.expectError(error.Truncated, decrementTtls(&short, 1));

    var overrun = "\x00\x01\x81\x80\x00\x01\x00\x00\x00\x00\x00\x00".*;
    try testing.expectError(error.SectionOverrun, decrementTtls(&overrun, 1));

    var bad_name = "\x00\x01\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\xc0\x0c\x00\x01\x00\x01".*;
    try testing.expectError(error.BadName, decrementTtls(&bad_name, 1));
}

test "decrementTtls leaves the buffer untouched when a later record is malformed" {
    // Two answers: a well-formed A record with TTL 100, then a record whose
    // owner name is a forward pointer.
    const bytes_const = "\x00\x01\x81\x80\x00\x00\x00\x02\x00\x00\x00\x00" ++
        "\x00\x00\x01\x00\x01\x00\x00\x00\x64\x00\x04\x01\x02\x03\x04" ++
        "\xc0\x40\x00\x01\x00\x01\x00\x00\x00\x64\x00\x00";
    var bytes: [bytes_const.len]u8 = bytes_const.*;

    try testing.expectError(error.BadName, decrementTtls(&bytes, 60));
    try testing.expectEqualSlices(u8, bytes_const, &bytes);
}

test "decrementTtls reports the packet that aging itself breaks" {
    // The first record's TTL word is 0x02_68_69_00, which also reads as the
    // name "hi." at offset 17. The second record's owner name points there, so
    // the packet parses — until aging rewrites those four bytes.
    const bytes_const = "\x00\x01\x81\x80\x00\x00\x00\x02\x00\x00\x00\x00" ++
        "\x00\x00\x01\x00\x01\x02\x68\x69\x00\x00\x04\x01\x02\x03\x04" ++
        "\xc0\x11\x00\x01\x00\x01\x00\x00\x00\x64\x00\x04\x05\x06\x07\x08";
    var bytes: [bytes_const.len]u8 = bytes_const.*;

    const p = try parse(&bytes);
    var it = answers(p);
    _ = (try it.next()).?;
    try testing.expectEqualSlices(u8, (try name.fromText("hi")).wire(), (try it.next()).?.name.wire());

    // Aging the first TTL turns the second owner name into a label that runs
    // off the end. The error is reported, not a panic, and the buffer is then
    // a partly aged message the caller must discard.
    try testing.expectError(error.SectionOverrun, decrementTtls(&bytes, 1000));
    try testing.expectError(error.SectionOverrun, parse(&bytes));
}

test "ResponseBuilder builds a reply that re-parses" {
    const request = try parse(query_bytes);
    const q = firstQuestion(request).?;
    const request_opt = try edns.parseOpt(query_bytes, findOptRecord(request).?);

    var buf: [512]u8 = undefined;
    var b = try ResponseBuilder.init(&buf, request.header, q);
    try b.addAnswer(q.name, .a, .in, 60, "\x0a\x00\x00\x01");
    try b.addOptEcho(request_opt, false);
    const bytes = b.finish();

    const p = try parse(bytes);
    try testing.expectEqual(request.header.id, p.header.id);
    try testing.expectEqual(true, p.header.flags.qr);
    try testing.expectEqual(true, p.header.flags.rd);
    try testing.expectEqual(true, p.header.flags.ra);
    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);
    try testing.expectEqual(@as(u16, 0), p.header.nscount);
    try testing.expectEqual(@as(u16, 1), p.header.arcount);

    const echoed = firstQuestion(p).?;
    try testing.expectEqualSlices(u8, q.name.wire(), echoed.name.wire());
    try testing.expectEqual(q.qtype, echoed.qtype);

    var it = answers(p);
    const a = (try it.next()).?;
    try testing.expectEqual(types.Type.a, a.rtype);
    try testing.expectEqual(@as(u32, 60), a.ttl);
    try testing.expectEqual([4]u8{ 10, 0, 0, 1 }, try record.rdataA(bytes, a));

    const opt = try edns.parseOpt(bytes, findOptRecord(p).?);
    try testing.expectEqual(@as(u16, 4096), opt.udp_payload_size);
    try testing.expectEqual(false, opt.do_bit);
}

test "a built reply decrements its ttls" {
    const request = try parse(query_bytes);
    const q = firstQuestion(request).?;

    var buf: [512]u8 = undefined;
    var b = try ResponseBuilder.init(&buf, request.header, q);
    try b.addAnswer(q.name, .a, .in, 90, "\x0a\x00\x00\x01");
    try b.addAnswer(q.name, .a, .in, 30, "\x0a\x00\x00\x02");
    const bytes = b.finish();

    try testing.expectEqual(@as(?u32, 5), try decrementTtls(bytes, 25));
    const p = try parse(bytes);
    var it = answers(p);
    try testing.expectEqual(@as(u32, 65), (try it.next()).?.ttl);
    try testing.expectEqual(@as(u32, 5), (try it.next()).?.ttl);
}

test "ResponseBuilder passes the DO bit through" {
    const request = try parse(query_bytes);
    const request_opt = try edns.parseOpt(query_bytes, findOptRecord(request).?);

    for ([_]bool{ false, true }) |do_bit| {
        var buf: [128]u8 = undefined;
        var b = try ResponseBuilder.init(&buf, request.header, firstQuestion(request).?);
        try b.addOptEcho(request_opt, do_bit);
        const bytes = b.finish();
        const p = try parse(bytes);
        const opt = try edns.parseOpt(bytes, findOptRecord(p).?);
        try testing.expectEqual(do_bit, opt.do_bit);
    }
}

test "addOptWithRcode carries the upper eight bits of a twelve-bit rcode" {
    const request = try parse(query_bytes);
    const request_opt = try edns.parseOpt(query_bytes, findOptRecord(request).?);
    const split = edns.splitRcode(edns.badvers);

    var buf: [128]u8 = undefined;
    var b = try ResponseBuilder.init(&buf, request.header, firstQuestion(request).?);
    b.setRcode(split.header);
    try b.addOptWithRcode(request_opt, false, split.extended);
    const bytes = b.finish();

    const p = try parse(bytes);
    const opt = try edns.parseOpt(bytes, findOptRecord(p).?);
    try testing.expectEqual(@as(u8, 0), opt.version);
    try testing.expectEqual(@as(u12, edns.badvers), edns.extendedRcode(p.header.flags.rcode, opt));
}

test "addOptEcho leaves the extended rcode clear" {
    const request = try parse(query_bytes);
    const request_opt = try edns.parseOpt(query_bytes, findOptRecord(request).?);

    var buf: [128]u8 = undefined;
    var b = try ResponseBuilder.init(&buf, request.header, firstQuestion(request).?);
    try b.addOptEcho(request_opt, false);
    const bytes = b.finish();

    const p = try parse(bytes);
    const opt = try edns.parseOpt(bytes, findOptRecord(p).?);
    try testing.expectEqual(@as(u8, 0), opt.extended_rcode);
}

test "ResponseBuilder sets an rcode and an empty answer section" {
    const request = try parse(query_bytes);

    var buf: [128]u8 = undefined;
    var b = try ResponseBuilder.init(&buf, request.header, firstQuestion(request).?);
    b.setRcode(.nx_domain);
    b.setAuthoritative(true);
    const bytes = b.finish();

    const p = try parse(bytes);
    try testing.expectEqual(types.Rcode.nx_domain, p.header.flags.rcode);
    try testing.expectEqual(true, p.header.flags.aa);
    try testing.expectEqual(@as(u16, 0), p.header.ancount);
    try testing.expectEqual(@as(usize, types.header_len + 17), bytes.len);
}

test "ResponseBuilder works without a question" {
    const request = try parse(query_bytes);

    var buf: [64]u8 = undefined;
    var b = try ResponseBuilder.init(&buf, request.header, null);
    b.setRcode(.form_err);
    const bytes = b.finish();

    try testing.expectEqual(@as(usize, types.header_len), bytes.len);
    const p = try parse(bytes);
    try testing.expectEqual(@as(u16, 0), p.header.qdcount);
    try testing.expectEqual(types.Rcode.form_err, p.header.flags.rcode);
}

test "ResponseBuilder keeps the opcode and a cleared RD bit" {
    const request: header.Header = .{
        .id = 0x4321,
        .flags = .{
            .rcode = .no_error,
            .z = 0,
            .ra = false,
            .rd = false,
            .tc = false,
            .aa = false,
            .opcode = .status,
            .qr = false,
        },
        .qdcount = 0,
        .ancount = 0,
        .nscount = 0,
        .arcount = 0,
    };

    var buf: [64]u8 = undefined;
    var b = try ResponseBuilder.init(&buf, request, null);
    const p = try parse(b.finish());
    try testing.expectEqual(@as(u16, 0x4321), p.header.id);
    try testing.expectEqual(types.Opcode.status, p.header.flags.opcode);
    try testing.expectEqual(false, p.header.flags.rd);
    try testing.expectEqual(@as(u3, 0), p.header.flags.z);
}

test "ResponseBuilder reports a buffer that is too small" {
    const request = try parse(query_bytes);
    const q = firstQuestion(request).?;

    var tiny: [8]u8 = undefined;
    try testing.expectError(error.WriteFailed, ResponseBuilder.init(&tiny, request.header, q));

    var no_room_for_question: [16]u8 = undefined;
    try testing.expectError(
        error.WriteFailed,
        ResponseBuilder.init(&no_room_for_question, request.header, q),
    );

    var no_room_for_answer: [32]u8 = undefined;
    var b = try ResponseBuilder.init(&no_room_for_answer, request.header, q);
    try testing.expectError(error.WriteFailed, b.addAnswer(q.name, .a, .in, 60, "\x01\x02\x03\x04"));
}
