milestone 7: serving pipeline, client tracking, pause and lifecycle

src/dns/edns.zig

@@ -10,6 +10,8 @@
 const std = @import("std");
 const types = @import("types.zig");
 const record = @import("record.zig");
+const question = @import("question.zig");
+const packet_mod = @import("packet.zig");
 const Writer = std.Io.Writer;
 
 /// The EDNS Client Subnet option code (RFC 7871 §6).
@@ -159,6 +161,111 @@ fn ttlFrom(opt: OptRecord) u32 {
         (@as(u32, @intFromBool(opt.do_bit)) << 15);
 }
 
+pub const StripResult = union(enum) {
+    /// The query needs no rewrite and nothing was written to `out`.
+    unchanged,
+    /// The rewritten query, a prefix of `out`.
+    rewritten: []u8,
+};
+
+/// Byte offset of RDLENGTH inside an OPT record as `encodeOpt` writes it: a
+/// one-byte root owner name, TYPE, CLASS and TTL come first.
+const opt_rdlength_offset = 1 + 2 + 2 + 4;
+
+/// Rebuilds `query` into `out` without its ECS option, so that a client's
+/// subnet never reaches the upstream resolver (PLAN §6.1). Every byte outside
+/// the OPT record is copied verbatim — ID, flags, all four section counts, the
+/// question and any record beside the OPT — and the OPT keeps its payload size,
+/// its flags and every option other than code 8.
+///
+/// `pkt` and `opt` are the caller's already-parsed views of `query`, and `out`
+/// must not overlap `query`.
+///
+/// One byte pattern does not survive the rewrite: a compression pointer in a
+/// record that follows the OPT and targets a byte inside the OPT's option list.
+/// Removing an option shifts those bytes, so the pointer then decodes to
+/// something else. RFC 1035 §4.1.4 only ever points a name at an earlier *name*,
+/// so such a pointer is malformed to begin with, and the upstream resolver
+/// answers the rewritten query with FORMERR instead of the original's answer.
+pub fn stripEcs(
+    query: []const u8,
+    pkt: packet_mod.Packet,
+    opt: OptRecord,
+    out: []u8,
+) error{ BadOption, Overflow }!StripResult {
+    std.debug.assert(query.ptr == pkt.bytes.ptr);
+    std.debug.assert(query.len == pkt.bytes.len);
+
+    // The rebuild reads `query` while it writes `out`, so an overlap would let
+    // it consume bytes it has already overwritten. The two buffers must be
+    // disjoint, and the caller finds that out here rather than in the packet it
+    // sends upstream.
+    std.debug.assert(@intFromPtr(out.ptr) + out.len <= @intFromPtr(query.ptr) or
+        @intFromPtr(query.ptr) + query.len <= @intFromPtr(out.ptr));
+
+    const opt_start = (try optRecordStart(pkt, opt)) orelse return .unchanged;
+    if ((try findOption(query, opt, ecs_option_code)) == null) return .unchanged;
+
+    var w = Writer.fixed(out);
+    w.writeAll(query[0..opt_start]) catch return error.Overflow;
+    encodeOpt(opt, &.{}, &w) catch |err| switch (err) {
+        error.OptionsTooLong => unreachable, // the list written here is empty
+        error.WriteFailed => return error.Overflow,
+    };
+
+    // RDLENGTH is only known once the list has been filtered, and the surviving
+    // options are not contiguous in `query`, so `encodeOpt` writes a placeholder
+    // and the options stream in behind it.
+    var kept_len: usize = 0;
+    var it = options(query, opt);
+    while (try it.next()) |o| {
+        if (o.code == ecs_option_code) continue;
+        var option_header: [4]u8 = undefined;
+        std.mem.writeInt(u16, option_header[0..2], o.code, .big);
+        std.mem.writeInt(u16, option_header[2..4], @intCast(o.data.len), .big);
+        w.writeAll(&option_header) catch return error.Overflow;
+        w.writeAll(o.data) catch return error.Overflow;
+        kept_len += option_header.len + o.data.len;
+    }
+
+    w.writeAll(query[opt.options.offset + opt.options.len ..]) catch return error.Overflow;
+
+    const message = w.buffered();
+    std.mem.writeInt(u16, message[opt_start + opt_rdlength_offset ..][0..2], @intCast(kept_len), .big);
+    return .{ .rewritten = message };
+}
+
+/// Where the record holding `opt` begins, or null when `pkt` has no such
+/// record. A record's start offset is not part of a parsed view — only a walk
+/// of every section before it reaches that byte — so this repeats the walk
+/// `packet.parse` already did.
+///
+/// The match is on the whole RDATA span, which keeps a hand-built `OptRecord`
+/// from pointing this function at a different record than the one it describes.
+/// A section that will not walk lands on `error.BadOption` for the same reason:
+/// a `Packet` from `packet.parse` always walks, so only a `Packet` assembled by
+/// hand around unvalidated bytes gets here.
+fn optRecordStart(pkt: packet_mod.Packet, opt: OptRecord) error{BadOption}!?usize {
+    var pos: usize = types.header_len;
+    var q: u16 = 0;
+    while (q < pkt.header.qdcount) : (q += 1) {
+        const parsed = question.parse(pkt.bytes, pos) catch return error.BadOption;
+        pos = parsed.end;
+    }
+
+    const record_count = @as(u32, pkt.header.ancount) +
+        @as(u32, pkt.header.nscount) + @as(u32, pkt.header.arcount);
+    var r: u32 = 0;
+    while (r < record_count) : (r += 1) {
+        const parsed = record.parse(pkt.bytes, pos) catch return error.BadOption;
+        if (parsed.record.rtype == .opt and
+            parsed.record.rdata.offset == opt.options.offset and
+            parsed.record.rdata.len == opt.options.len) return pos;
+        pos = parsed.end;
+    }
+    return null;
+}
+
 /// The full 12-bit RCODE (RFC 6891 §6.1.3): the OPT record supplies the upper
 /// eight bits, the header the lower four. Without an OPT record the value is
 /// just the header's four bits.
@@ -406,6 +513,172 @@ test "encodeOpt reports a short buffer" {
     try testing.expectError(error.WriteFailed, encodeOpt(opt, "", &w));
 }
 
+/// A query for example.com A whose OPT record carries an ECS option for
+/// 192.0.2.0/24 followed by a two-byte padding option, with DO set and a
+/// 4096-byte payload size.
+const ecs_query =
+    "\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x01" ++ // header
+    "\x07example\x03com\x00\x00\x01\x00\x01" ++ // question
+    "\x00\x00\x29\x10\x00\x00\x00\x80\x00\x00\x11" ++ // OPT, 17 rdata bytes
+    "\x00\x08\x00\x07\x00\x01\x18\x00\xc0\x00\x02" ++ // ECS
+    "\x00\x0c\x00\x02\x00\x00"; // padding
+
+/// `ecs_query` as `stripEcs` must rebuild it.
+const ecs_query_stripped =
+    "\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x01" ++
+    "\x07example\x03com\x00\x00\x01\x00\x01" ++
+    "\x00\x00\x29\x10\x00\x00\x00\x80\x00\x00\x06" ++
+    "\x00\x0c\x00\x02\x00\x00";
+
+const ParsedQuery = struct { pkt: packet_mod.Packet, opt: OptRecord };
+
+fn parseQuery(bytes: []const u8) !ParsedQuery {
+    const p = try packet_mod.parse(bytes);
+    return .{ .pkt = p, .opt = try parseOpt(bytes, packet_mod.findOptRecord(p).?) };
+}
+
+fn expectUnchanged(result: StripResult) !void {
+    switch (result) {
+        .unchanged => {},
+        .rewritten => return error.TestExpectedUnchanged,
+    }
+}
+
+test "stripEcs rebuilds the query without the ECS option" {
+    const parsed = try parseQuery(ecs_query);
+    var out: [512]u8 = undefined;
+    const result = try stripEcs(ecs_query, parsed.pkt, parsed.opt, &out);
+    const rewritten = result.rewritten;
+    try testing.expectEqualSlices(u8, ecs_query_stripped, rewritten);
+
+    const p = try packet_mod.parse(rewritten);
+    try testing.expectEqual(@as(u16, 0x1234), p.header.id);
+    try testing.expectEqual(@as(u16, 1), p.header.qdcount);
+    try testing.expectEqual(@as(u16, 1), p.header.arcount);
+
+    const opt = try parseOpt(rewritten, packet_mod.findOptRecord(p).?);
+    try testing.expectEqual(@as(u16, 4096), opt.udp_payload_size);
+    try testing.expectEqual(true, opt.do_bit);
+    try testing.expectEqual(@as(?Option, null), try findOption(rewritten, opt, ecs_option_code));
+    try testing.expectEqualSlices(u8, "\x00\x00", (try findOption(rewritten, opt, 12)).?.data);
+}
+
+test "stripEcs reports a query with nothing to strip as unchanged" {
+    const parsed = try parseQuery(ecs_query_stripped);
+    var out: [512]u8 = undefined;
+    try expectUnchanged(try stripEcs(ecs_query_stripped, parsed.pkt, parsed.opt, &out));
+
+    // No OPT record at all: the caller still holds an `OptRecord`, so the
+    // absence has to be found by the walk rather than assumed.
+    const no_opt = "\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00" ++
+        "\x07example\x03com\x00\x00\x01\x00\x01";
+    const p = try packet_mod.parse(no_opt);
+    try testing.expectEqual(@as(?record.Record, null), packet_mod.findOptRecord(p));
+    const absent: OptRecord = .{
+        .udp_payload_size = 4096,
+        .extended_rcode = 0,
+        .version = 0,
+        .do_bit = false,
+        .options = .{ .offset = 0, .len = 0 },
+    };
+    try expectUnchanged(try stripEcs(no_opt, p, absent, &out));
+}
+
+test "stripEcs keeps a record that follows the OPT" {
+    const head = "\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x02" ++
+        "\x07example\x03com\x00\x00\x01\x00\x01";
+    const trailing = "\x00\x00\x01\x00\x01\x00\x00\x00\x0a\x00\x04\x01\x02\x03\x04";
+    const query = head ++
+        "\x00\x00\x29\x10\x00\x00\x00\x80\x00\x00\x0b" ++
+        "\x00\x08\x00\x07\x00\x01\x18\x00\xc0\x00\x02" ++
+        trailing;
+
+    const parsed = try parseQuery(query);
+    var out: [512]u8 = undefined;
+    const rewritten = (try stripEcs(query, parsed.pkt, parsed.opt, &out)).rewritten;
+    try testing.expectEqualSlices(
+        u8,
+        head ++ "\x00\x00\x29\x10\x00\x00\x00\x80\x00\x00\x00" ++ trailing,
+        rewritten,
+    );
+    try testing.expectEqual(@as(u16, 2), (try packet_mod.parse(rewritten)).header.arcount);
+}
+
+test "stripEcs removes a repeated ECS option" {
+    const ecs = "\x00\x08\x00\x07\x00\x01\x18\x00\xc0\x00\x02";
+    const query = "\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x01" ++
+        "\x07example\x03com\x00\x00\x01\x00\x01" ++
+        "\x00\x00\x29\x10\x00\x00\x00\x80\x00\x00\x1c" ++
+        ecs ++ "\x00\x0c\x00\x02\x00\x00" ++ ecs;
+
+    const parsed = try parseQuery(query);
+    var out: [512]u8 = undefined;
+    const rewritten = (try stripEcs(query, parsed.pkt, parsed.opt, &out)).rewritten;
+    try testing.expectEqualSlices(u8, ecs_query_stripped, rewritten);
+}
+
+test "stripEcs reports an output buffer one byte short" {
+    const parsed = try parseQuery(ecs_query);
+
+    var exact: [ecs_query_stripped.len]u8 = undefined;
+    _ = (try stripEcs(ecs_query, parsed.pkt, parsed.opt, &exact)).rewritten;
+
+    var short: [ecs_query_stripped.len - 1]u8 = undefined;
+    try testing.expectError(error.Overflow, stripEcs(ecs_query, parsed.pkt, parsed.opt, &short));
+
+    // Every shorter buffer fails the same way, including one that cannot even
+    // hold the copied header.
+    var i: usize = 0;
+    while (i < short.len) : (i += 1) {
+        try testing.expectError(
+            error.Overflow,
+            stripEcs(ecs_query, parsed.pkt, parsed.opt, short[0..i]),
+        );
+    }
+}
+
+test "stripEcs rejects a malformed option list" {
+    // The record is well-formed; only its option list overruns, so `parseOpt`
+    // is the one thing that rejects this packet.
+    const query = "\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x01" ++
+        "\x07example\x03com\x00\x00\x01\x00\x01" ++
+        "\x00\x00\x29\x10\x00\x00\x00\x00\x00\x00\x08" ++
+        "\x00\x08\x00\x08\x00\x01\x18\x00";
+
+    const p = try packet_mod.parse(query);
+    const rec = packet_mod.findOptRecord(p).?;
+    try testing.expectError(error.BadOption, parseOpt(query, rec));
+
+    const hand_built: OptRecord = .{
+        .udp_payload_size = rec.class,
+        .extended_rcode = 0,
+        .version = 0,
+        .do_bit = false,
+        .options = rec.rdata,
+    };
+
+    var out: [512]u8 = undefined;
+    try testing.expectError(error.BadOption, stripEcs(query, p, hand_built, &out));
+}
+
+test "encodeOpt writes RDLENGTH where stripEcs patches it" {
+    const opt: OptRecord = .{
+        .udp_payload_size = 4096,
+        .extended_rcode = 0,
+        .version = 0,
+        .do_bit = true,
+        .options = .{ .offset = 0, .len = 0 },
+    };
+    var buf: [32]u8 = undefined;
+    var w = Writer.fixed(&buf);
+    try encodeOpt(opt, "\x00\x0c\x00\x02\x00\x00", &w);
+    const bytes = w.buffered();
+    try testing.expectEqual(
+        @as(u16, 6),
+        std.mem.readInt(u16, bytes[opt_rdlength_offset..][0..2], .big),
+    );
+}
+
 test "extendedRcode composes the twelve bits" {
     try testing.expectEqual(@as(u12, 3), extendedRcode(.nx_domain, null));