resolver transport: udp/tcp servers, doh/dot clients, pool failover with health

This commit is contained in:
2026-08-01 12:36:50 +02:00
parent 346f2dc502
commit 17d0401f8a
15 changed files with 6062 additions and 0 deletions
+825
View File
@@ -0,0 +1,825 @@
//! The serving handler: raw bytes from a listener in, bytes to send back out.
//! It owns no socket, no clock and no allocator — the listener supplies the
//! buffers and the upstream supplies the answer.
//!
//! `handle` returns no error union. Every failure is either a DNS response the
//! client can act on or a counted drop, because a listener has nothing useful
//! to do with an error value: it cannot retry and it must not die. That is the
//! "every failure mode is visible" rule applied to the hot path — the counters
//! are the failure surface.
//!
//! Phase 7 extends this file with filtering, cache, local records and query
//! logging. None of that is here: this handler forwards every query.
const std = @import("std");
const edns = @import("../dns/edns.zig");
const header = @import("../dns/header.zig");
const packet = @import("../dns/packet.zig");
const question = @import("../dns/question.zig");
const types = @import("../dns/types.zig");
const transport = @import("../upstream/transport.zig");
/// Which listener a query arrived on. Only the UDP size limit depends on it.
pub const Transport = enum { udp, tcp };
/// RFC 1035 §4.2.1: 512 bytes is what a client accepts over UDP without EDNS,
/// and RFC 6891 §6.2.3 says a smaller advertised size must not be honoured
/// below that floor.
pub const udp_limit_min: u16 = types.max_udp_payload;
/// The ceiling nxdns advertises and accepts. Above this, fragmentation and
/// reflection amplification cost more than a TCP retry.
pub const udp_limit_max: u16 = 4096;
/// A synthesized reply is a header, at most one question and at most one OPT
/// record: 12 + (255 + 4) + 11 = 282 bytes worst case. Every buffer this file
/// builds into is at least `udp_limit_min`, so the builder cannot run out of
/// room and its `Writer.Error` is unreachable.
const max_synthetic_len = types.header_len + types.max_name_len + 4 + 11;
comptime {
std.debug.assert(max_synthetic_len <= udp_limit_min);
}
pub const Handler = struct {
/// In production this is `pool.client()`.
upstream: transport.Client,
stats: Stats = .{},
pub const Stats = struct {
queries: std.atomic.Value(u64) = .init(0),
dropped_malformed: std.atomic.Value(u64) = .init(0),
formerr: std.atomic.Value(u64) = .init(0),
notimp: std.atomic.Value(u64) = .init(0),
servfail: std.atomic.Value(u64) = .init(0),
truncated: std.atomic.Value(u64) = .init(0),
};
pub const Outcome = union(enum) {
/// A prefix of `response_buf`.
reply: []u8,
/// No response is possible or appropriate.
drop,
};
/// `response_buf.len` must be >= 512 and <= 65535.
pub fn handle(
self: *Handler,
io: std.Io,
which: Transport,
query: []const u8,
response_buf: []u8,
) Outcome {
std.debug.assert(response_buf.len >= udp_limit_min);
std.debug.assert(response_buf.len <= transport.max_message_len);
// Fewer than 12 bytes: there is no ID and no question, so no reply can
// be addressed to this query (PLAN §6.1).
const hdr = header.parse(query) catch {
bump(&self.stats.dropped_malformed);
return .drop;
};
// A response on a listener port is either a misdirected reply or a
// reflection attempt. Answering it would make this server the amplifier.
if (hdr.flags.qr) {
bump(&self.stats.dropped_malformed);
return .drop;
}
const p = packet.parse(query) catch |err| switch (err) {
error.Truncated => {
bump(&self.stats.dropped_malformed);
return .drop;
},
// The header is intact but a section is not. The question is what
// failed to parse, so nothing is echoed and no OPT is trusted.
else => return synthesize(hdr, null, null, .form_err, &self.stats.formerr, response_buf),
};
// RFC 6891 §6.1.1 puts the OPT record in the additional section, and §7
// answers a violation with FORMERR. `packet.parse` already rejects a
// second OPT wherever it sits; the case it accepts and this walk does
// not is a lone OPT in the answer or authority section, which
// `findOptRecord` cannot see.
if (containsOpt(packet.answers(p)) or containsOpt(packet.authorities(p))) {
const echo = if (hdr.qdcount == 1) packet.firstQuestion(p) else null;
return synthesize(hdr, echo, null, .form_err, &self.stats.formerr, response_buf);
}
// RFC 6891 §6.1.1 also gives the OPT record a root owner name. An OPT
// that is present but unusable must not be forwarded as if the query
// carried none: the reply's OPT echo would then disappear with nothing
// to say it ever existed.
const opt: ?edns.OptRecord = if (packet.findOptRecord(p)) |rec|
edns.parseOpt(query, rec) catch {
const echo = if (hdr.qdcount == 1) packet.firstQuestion(p) else null;
return synthesize(hdr, echo, null, .form_err, &self.stats.formerr, response_buf);
}
else
null;
if (hdr.flags.opcode != .query) {
return synthesize(hdr, null, opt, .not_imp, &self.stats.notimp, response_buf);
}
// RFC 9619: exactly one question, in both directions. Any other count
// has no valid interpretation, so there is no question to echo either.
if (hdr.qdcount != 1) {
return synthesize(hdr, null, opt, .form_err, &self.stats.formerr, response_buf);
}
// `parse` already walked the question section, so a packet with
// qdcount 1 always has a first question. The fallback keeps a
// hand-assembled `Packet` from turning into a crash.
const q = packet.firstQuestion(p) orelse
return synthesize(hdr, null, opt, .form_err, &self.stats.formerr, response_buf);
const reply = self.upstream.exchange(io, query, response_buf) catch |err| switch (transport.group(err)) {
// The process is shutting down. There is nothing to say, and the
// client is about to lose the socket anyway; the listener's own
// counters record the abandoned datagram.
.cancellation => return .drop,
.peer_fault, .local_resource => return synthesize(
hdr,
q,
opt,
.serv_fail,
&self.stats.servfail,
response_buf,
),
};
bump(&self.stats.queries);
if (which == .udp and reply.len > udpLimit(p)) {
// The upstream answer occupies `response_buf`, so the replacement
// is built beside it and copied over.
bump(&self.stats.truncated);
var scratch: [udp_limit_min]u8 = undefined;
const message = build(hdr, q, opt, .no_error, true, &scratch);
@memcpy(response_buf[0..message.len], message);
return .{ .reply = response_buf[0..message.len] };
}
return .{ .reply = reply };
}
};
/// Counts one synthesized reply and encodes it. The counter is passed in
/// because the counter *is* the record that this failure mode happened.
fn synthesize(
hdr: header.Header,
q: ?question.Question,
opt: ?edns.OptRecord,
rcode: types.Rcode,
counter: *std.atomic.Value(u64),
response_buf: []u8,
) Handler.Outcome {
bump(counter);
return .{ .reply = build(hdr, q, opt, rcode, false, response_buf) };
}
/// Whether the section holds a record of TYPE=OPT. A section that will not walk
/// counts as holding one: the caller's only answer to either condition is
/// FORMERR, and a `Packet` from `packet.parse` walks in full, so the fallback is
/// reachable only for a `Packet` assembled by hand around unvalidated bytes.
fn containsOpt(section: packet.RecordIterator) bool {
var it = section;
while (it.next() catch return true) |rec| {
if (rec.rtype == .opt) return true;
}
return false;
}
/// The query's advertised UDP payload size, clamped to `[512, 4096]`, or 512
/// when the query carries no usable OPT record.
pub fn udpLimit(query_packet: packet.Packet) u16 {
const rec = packet.findOptRecord(query_packet) orelse return udp_limit_min;
const opt = edns.parseOpt(query_packet.bytes, rec) catch return udp_limit_min;
return std.math.clamp(opt.udp_payload_size, udp_limit_min, udp_limit_max);
}
/// Encodes one synthesized reply. `ResponseBuilder.init` copies the request's
/// ID, opcode and RD bit and sets QR and RA, so every reply built here binds to
/// its request. `buf` is at least `udp_limit_min`, which `max_synthetic_len`
/// proves is enough, so the writer cannot fail.
fn build(
hdr: header.Header,
q: ?question.Question,
opt: ?edns.OptRecord,
rcode: types.Rcode,
tc: bool,
buf: []u8,
) []u8 {
var b = packet.ResponseBuilder.init(buf, hdr, q) catch unreachable;
b.setRcode(rcode);
b.header.flags.tc = tc;
// An EDNS query gets an EDNS reply, and the DO bit passes through
// untouched — nxdns validates no signatures, so it must not claim the
// client asked for none (PLAN §6.1).
if (opt) |o| b.addOptEcho(o, o.do_bit) catch unreachable;
return b.finish();
}
fn bump(counter: *std.atomic.Value(u64)) void {
_ = counter.fetchAdd(1, .monotonic);
}
const testing = std.testing;
/// 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";
const opt_len = 11;
const query_with_opt_len = query_bytes.len + opt_len;
/// `query_bytes` plus an OPT record advertising `payload_size`.
fn queryWithOpt(buf: *[query_with_opt_len]u8, payload_size: u16, do_bit: bool) []const u8 {
@memcpy(buf[0..query_bytes.len], query_bytes);
std.mem.writeInt(u16, buf[10..12], 1, .big); // arcount
const opt = buf[query_bytes.len..][0..opt_len];
@memset(opt, 0);
opt[2] = @intFromEnum(types.Type.opt);
std.mem.writeInt(u16, opt[3..5], payload_size, .big);
if (do_bit) opt[7] = 0x80; // the DO bit is bit 15 of the TTL word
return buf;
}
/// `query_bytes` plus an OPT record whose owner name is `com.` instead of the
/// root, which RFC 6891 §6.1.1 forbids.
const query_with_named_opt =
"\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x01" ++
"\x07example\x03com\x00\x00\x01\x00\x01" ++
"\x03com\x00\x00\x29\x10\x00\x00\x00\x00\x00\x00\x00";
/// `query_bytes` plus an OPT record whose RDATA ends in a three-byte option
/// header. The record itself fits the packet; only the option list is broken.
const query_with_bad_option =
"\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\x03\x00\x08\x00";
/// A bare OPT record: root owner, TYPE 41, 4096-byte payload, empty RDATA.
const opt_record = "\x00\x00\x29\x10\x00\x00\x00\x00\x00\x00\x00";
/// `query_bytes` with the OPT record counted into the answer section, which
/// RFC 6891 §6.1.1 forbids.
const query_with_opt_in_answer =
"\x12\x34\x01\x00\x00\x01\x00\x01\x00\x00\x00\x00" ++
"\x07example\x03com\x00\x00\x01\x00\x01" ++
opt_record;
/// The same OPT record, counted into the authority section instead.
const query_with_opt_in_authority =
"\x12\x34\x01\x00\x00\x01\x00\x00\x00\x01\x00\x00" ++
"\x07example\x03com\x00\x00\x01\x00\x01" ++
opt_record;
/// Two OPT records in the additional section. RFC 6891 §6.1.1 allows one.
const query_with_two_opts =
"\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x02" ++
"\x07example\x03com\x00\x00\x01\x00\x01" ++
opt_record ++ opt_record;
const FakeUpstream = struct {
reply: []const u8 = &.{},
err: ?transport.ExchangeError = null,
calls: usize = 0,
fn exchangeFn(
ptr: *anyopaque,
io: std.Io,
query: []const u8,
response_buf: []u8,
) transport.ExchangeError![]u8 {
_ = io;
const self: *FakeUpstream = @ptrCast(@alignCast(ptr));
self.calls += 1;
if (self.err) |e| return e;
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];
if (bytes.len >= types.header_len) {
packet.setId(bytes, (header.parse(query) catch unreachable).id);
}
return bytes;
}
fn client(self: *FakeUpstream) transport.Client {
return .{ .ptr = self, .exchangeFn = exchangeFn };
}
};
fn expectReply(outcome: Handler.Outcome) ![]u8 {
return switch (outcome) {
.reply => |bytes| bytes,
.drop => error.TestUnexpectedDrop,
};
}
test "a udp query is forwarded and the upstream reply is returned unchanged" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query_bytes, &buf));
try testing.expectEqualSlices(u8, response_bytes, reply);
try testing.expectEqual(@as(usize, 1), fake.calls);
try testing.expectEqual(@as(u64, 1), h.stats.queries.load(.monotonic));
try testing.expectEqual(@as(u64, 0), h.stats.truncated.load(.monotonic));
// A query that carries no OPT at all is not an EDNS violation.
try testing.expectEqual(@as(u64, 0), h.stats.formerr.load(.monotonic));
}
test "a response arriving on the listener port is dropped" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
var buf: [udp_limit_min]u8 = undefined;
try testing.expectEqual(Handler.Outcome.drop, h.handle(undefined, .udp, response_bytes, &buf));
try testing.expectEqual(@as(u64, 1), h.stats.dropped_malformed.load(.monotonic));
try testing.expectEqual(@as(usize, 0), fake.calls);
}
test "a query shorter than a header is dropped" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
var buf: [udp_limit_min]u8 = undefined;
const short = query_bytes[0..8];
try testing.expectEqual(Handler.Outcome.drop, h.handle(undefined, .udp, short, &buf));
try testing.expectEqual(@as(u64, 1), h.stats.dropped_malformed.load(.monotonic));
try testing.expectEqual(@as(usize, 0), fake.calls);
}
test "a query whose question runs off the end gets FORMERR" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
// The header promises a question; the packet ends after 12 bytes plus a
// partial name, which `packet.parse` reports as a section overrun.
var buf: [udp_limit_min]u8 = undefined;
const overrun = query_bytes[0 .. query_bytes.len - 3];
const outcome = h.handle(undefined, .udp, overrun, &buf);
const reply = try expectReply(outcome);
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.form_err, p.header.flags.rcode);
try testing.expectEqual(@as(u64, 1), h.stats.formerr.load(.monotonic));
}
test "a query with QDCOUNT 0 gets FORMERR with no question" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
const no_question = "\x12\x34\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00";
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, no_question, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.form_err, p.header.flags.rcode);
try testing.expectEqual(@as(u16, 0x1234), p.header.id);
try testing.expectEqual(@as(u16, 0), p.header.qdcount);
try testing.expectEqual(true, p.header.flags.qr);
try testing.expectEqual(true, p.header.flags.ra);
try testing.expectEqual(@as(u64, 1), h.stats.formerr.load(.monotonic));
try testing.expectEqual(@as(usize, 0), fake.calls);
}
test "a query with QDCOUNT 2 gets FORMERR" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
const two =
"\x12\x34\x01\x00\x00\x02\x00\x00\x00\x00\x00\x00" ++
"\x07example\x03com\x00\x00\x01\x00\x01" ++
"\x07example\x03com\x00\x00\x1c\x00\x01";
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, two, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.form_err, p.header.flags.rcode);
try testing.expectEqual(@as(u16, 0), p.header.qdcount);
try testing.expectEqual(@as(u64, 1), h.stats.formerr.load(.monotonic));
try testing.expectEqual(@as(usize, 0), fake.calls);
}
test "a structurally broken question gets FORMERR" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
// A compression pointer to itself: the name never terminates.
const loop = "\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00" ++
"\xc0\x0c\x00\x01\x00\x01";
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, loop, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.form_err, p.header.flags.rcode);
try testing.expectEqual(@as(u16, 0x1234), p.header.id);
try testing.expectEqual(@as(u16, 0), p.header.qdcount);
try testing.expectEqual(@as(u64, 1), h.stats.formerr.load(.monotonic));
try testing.expectEqual(@as(usize, 0), fake.calls);
}
test "an unimplemented opcode gets NOTIMP" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
var update: [query_bytes.len]u8 = query_bytes.*;
std.mem.writeInt(u16, update[2..4], 0x2900, .big); // opcode update, RD set
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, &update, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.not_imp, p.header.flags.rcode);
try testing.expectEqual(types.Opcode.update, p.header.flags.opcode);
try testing.expectEqual(@as(u16, 0x1234), p.header.id);
try testing.expectEqual(@as(u16, 0), p.header.qdcount);
try testing.expectEqual(@as(u64, 1), h.stats.notimp.load(.monotonic));
try testing.expectEqual(@as(usize, 0), fake.calls);
}
test "a peer fault becomes SERVFAIL with the question echoed" {
var fake: FakeUpstream = .{ .err = error.Timeout };
var h: Handler = .{ .upstream = fake.client() };
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query_bytes, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.serv_fail, p.header.flags.rcode);
try testing.expectEqual(@as(u16, 0x1234), p.header.id);
try testing.expectEqual(true, p.header.flags.rd);
try testing.expectEqual(true, p.header.flags.qr);
try testing.expectEqual(@as(u16, 1), p.header.qdcount);
try testing.expectEqual(@as(u16, 0), p.header.ancount);
const echoed = packet.firstQuestion(p).?;
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);
try testing.expectEqual(@as(u64, 1), h.stats.servfail.load(.monotonic));
try testing.expectEqual(@as(u64, 0), h.stats.queries.load(.monotonic));
}
test "a local resource failure becomes SERVFAIL, not a drop" {
var fake: FakeUpstream = .{ .err = error.OutOfMemory };
var h: Handler = .{ .upstream = fake.client() };
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query_bytes, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.serv_fail, p.header.flags.rcode);
try testing.expectEqual(@as(u16, 1), p.header.qdcount);
try testing.expectEqual(@as(u64, 1), h.stats.servfail.load(.monotonic));
}
test "a canceled exchange is dropped" {
var fake: FakeUpstream = .{ .err = error.Canceled };
var h: Handler = .{ .upstream = fake.client() };
var buf: [udp_limit_min]u8 = undefined;
try testing.expectEqual(Handler.Outcome.drop, h.handle(undefined, .udp, query_bytes, &buf));
try testing.expectEqual(@as(u64, 0), h.stats.servfail.load(.monotonic));
try testing.expectEqual(@as(u64, 0), h.stats.queries.load(.monotonic));
try testing.expectEqual(@as(u64, 0), h.stats.dropped_malformed.load(.monotonic));
}
test "udpLimit clamps the advertised payload size" {
try testing.expectEqual(@as(u16, 512), udpLimit(try packet.parse(query_bytes)));
const cases = [_]struct { advertised: u16, expected: u16 }{
.{ .advertised = 1232, .expected = 1232 },
.{ .advertised = 200, .expected = 512 },
.{ .advertised = 9000, .expected = 4096 },
.{ .advertised = 512, .expected = 512 },
.{ .advertised = 4096, .expected = 4096 },
};
for (cases) |c| {
var buf: [query_with_opt_len]u8 = undefined;
const q = queryWithOpt(&buf, c.advertised, false);
try testing.expectEqual(c.expected, udpLimit(try packet.parse(q)));
}
}
/// A valid response for `query_bytes` padded past 512 bytes with A records.
fn oversizeResponse(buf: []u8) []u8 {
const request = packet.parse(query_bytes) catch unreachable;
const q = packet.firstQuestion(request).?;
var b = packet.ResponseBuilder.init(buf, request.header, q) catch unreachable;
var i: usize = 0;
while (i < 32) : (i += 1) {
b.addAnswer(q.name, .a, .in, 300, "\x5d\xb8\xd8\x22") catch unreachable;
}
return b.finish();
}
test "an oversize udp reply is replaced by a truncated one" {
var reply_buf: [2048]u8 = undefined;
const oversize = oversizeResponse(&reply_buf);
try testing.expect(oversize.len > udp_limit_min);
var fake: FakeUpstream = .{ .reply = oversize };
var h: Handler = .{ .upstream = fake.client() };
var buf: [4096]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query_bytes, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(true, p.header.flags.tc);
try testing.expectEqual(types.Rcode.no_error, p.header.flags.rcode);
try testing.expectEqual(@as(u16, 0), p.header.ancount);
try testing.expectEqual(@as(u16, 1), p.header.qdcount);
try testing.expectEqual(@as(u16, 0x1234), p.header.id);
try testing.expectEqual(true, p.header.flags.rd);
const echoed = packet.firstQuestion(p).?;
const asked = packet.firstQuestion(try packet.parse(query_bytes)).?;
try testing.expectEqualSlices(u8, asked.name.wire(), echoed.name.wire());
try testing.expectEqual(@as(u64, 1), h.stats.truncated.load(.monotonic));
try testing.expectEqual(@as(u64, 1), h.stats.queries.load(.monotonic));
}
test "the same oversize reply passes through untouched over tcp" {
var reply_buf: [2048]u8 = undefined;
const oversize = oversizeResponse(&reply_buf);
var fake: FakeUpstream = .{ .reply = oversize };
var h: Handler = .{ .upstream = fake.client() };
var buf: [4096]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .tcp, query_bytes, &buf));
try testing.expectEqualSlices(u8, oversize, reply);
try testing.expectEqual(@as(u64, 0), h.stats.truncated.load(.monotonic));
try testing.expectEqual(@as(u64, 1), h.stats.queries.load(.monotonic));
}
test "a reply within the advertised EDNS limit is not truncated" {
var reply_buf: [2048]u8 = undefined;
const oversize = oversizeResponse(&reply_buf);
var fake: FakeUpstream = .{ .reply = oversize };
var h: Handler = .{ .upstream = fake.client() };
var query_buf: [query_with_opt_len]u8 = undefined;
const query = queryWithOpt(&query_buf, 4096, false);
var buf: [4096]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query, &buf));
try testing.expectEqual(oversize.len, reply.len);
try testing.expectEqual(@as(u64, 0), h.stats.truncated.load(.monotonic));
}
test "the DO bit passes through into a synthesized reply" {
for ([_]bool{ false, true }) |do_bit| {
var fake: FakeUpstream = .{ .err = error.Timeout };
var h: Handler = .{ .upstream = fake.client() };
var query_buf: [query_with_opt_len]u8 = undefined;
const query = queryWithOpt(&query_buf, 1232, do_bit);
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.serv_fail, p.header.flags.rcode);
const opt = try edns.parseOpt(reply, packet.findOptRecord(p).?);
try testing.expectEqual(do_bit, opt.do_bit);
try testing.expectEqual(@as(u16, 1232), opt.udp_payload_size);
}
}
test "the truncated reply echoes the OPT record" {
var reply_buf: [2048]u8 = undefined;
const oversize = oversizeResponse(&reply_buf);
var fake: FakeUpstream = .{ .reply = oversize };
var h: Handler = .{ .upstream = fake.client() };
// 512 is the advertised size, so the oversize answer still does not fit.
var query_buf: [query_with_opt_len]u8 = undefined;
const query = queryWithOpt(&query_buf, 512, true);
var buf: [4096]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(true, p.header.flags.tc);
const opt = try edns.parseOpt(reply, packet.findOptRecord(p).?);
try testing.expectEqual(true, opt.do_bit);
try testing.expectEqual(@as(u16, 512), opt.udp_payload_size);
try testing.expectEqual(@as(u64, 1), h.stats.truncated.load(.monotonic));
}
test "an OPT record with a non-root owner name gets FORMERR" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query_with_named_opt, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.form_err, p.header.flags.rcode);
try testing.expectEqual(@as(u16, 0x1234), p.header.id);
try testing.expectEqual(true, p.header.flags.qr);
try testing.expectEqual(true, p.header.flags.ra);
// The question parsed, so it is echoed; the OPT did not, so none is.
try testing.expectEqual(@as(u16, 1), p.header.qdcount);
const echoed = packet.firstQuestion(p).?;
const asked = packet.firstQuestion(try packet.parse(query_bytes)).?;
try testing.expectEqualSlices(u8, asked.name.wire(), echoed.name.wire());
try testing.expect(packet.findOptRecord(p) == null);
try testing.expectEqual(@as(u64, 1), h.stats.formerr.load(.monotonic));
try testing.expectEqual(@as(usize, 0), fake.calls);
}
test "an OPT record with a malformed option list gets FORMERR" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
// The packet as a whole is well-formed: only `parseOpt` rejects it.
const p_query = try packet.parse(query_with_bad_option);
try testing.expectError(
error.BadOption,
edns.parseOpt(query_with_bad_option, packet.findOptRecord(p_query).?),
);
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query_with_bad_option, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.form_err, p.header.flags.rcode);
try testing.expectEqual(@as(u16, 0x1234), p.header.id);
try testing.expect(packet.findOptRecord(p) == null);
try testing.expectEqual(@as(u64, 1), h.stats.formerr.load(.monotonic));
try testing.expectEqual(@as(usize, 0), fake.calls);
}
test "an OPT record whose rdata runs off the end gets FORMERR" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
// RDLENGTH claims three bytes and one follows, which `packet.parse` reports
// as a section overrun before the OPT is ever read.
const truncated = query_with_bad_option[0 .. query_with_bad_option.len - 2];
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, truncated, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.form_err, p.header.flags.rcode);
try testing.expectEqual(@as(u64, 1), h.stats.formerr.load(.monotonic));
try testing.expectEqual(@as(usize, 0), fake.calls);
}
test "an OPT record in the answer section gets FORMERR" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
// The message walks: only its placement of the OPT is illegal, so the
// handler's own section walk is what has to reject it.
const p_query = try packet.parse(query_with_opt_in_answer);
try testing.expect(packet.findOptRecord(p_query) == null);
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query_with_opt_in_answer, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.form_err, p.header.flags.rcode);
try testing.expectEqual(@as(u16, 0x1234), p.header.id);
// The question parsed, so it is echoed; the OPT is not echoed anywhere.
try testing.expectEqual(@as(u16, 1), p.header.qdcount);
const echoed = packet.firstQuestion(p).?;
const asked = packet.firstQuestion(try packet.parse(query_bytes)).?;
try testing.expectEqualSlices(u8, asked.name.wire(), echoed.name.wire());
try testing.expect(packet.findOptRecord(p) == null);
try testing.expectEqual(@as(u16, 0), p.header.arcount);
try testing.expectEqual(@as(u64, 1), h.stats.formerr.load(.monotonic));
try testing.expectEqual(@as(usize, 0), fake.calls);
}
test "an OPT record in the authority section gets FORMERR" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
const p_query = try packet.parse(query_with_opt_in_authority);
try testing.expect(packet.findOptRecord(p_query) == null);
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query_with_opt_in_authority, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.form_err, p.header.flags.rcode);
try testing.expectEqual(@as(u16, 0x1234), p.header.id);
try testing.expectEqual(@as(u16, 1), p.header.qdcount);
try testing.expect(packet.findOptRecord(p) == null);
try testing.expectEqual(@as(u64, 1), h.stats.formerr.load(.monotonic));
try testing.expectEqual(@as(usize, 0), fake.calls);
}
test "two OPT records in the additional section get FORMERR" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
// Here the rejection comes from `packet.parse`, one section walk earlier.
try testing.expectError(error.MultipleOptRecords, packet.parse(query_with_two_opts));
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query_with_two_opts, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(types.Rcode.form_err, p.header.flags.rcode);
try testing.expectEqual(@as(u16, 0x1234), p.header.id);
try testing.expect(packet.findOptRecord(p) == null);
try testing.expectEqual(@as(u64, 1), h.stats.formerr.load(.monotonic));
try testing.expectEqual(@as(usize, 0), fake.calls);
}
test "a query with a valid OPT record is still forwarded" {
var fake: FakeUpstream = .{ .reply = response_bytes };
var h: Handler = .{ .upstream = fake.client() };
var query_buf: [query_with_opt_len]u8 = undefined;
const query = queryWithOpt(&query_buf, 1232, true);
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, query, &buf));
try testing.expectEqualSlices(u8, response_bytes, reply);
try testing.expectEqual(@as(usize, 1), fake.calls);
try testing.expectEqual(@as(u64, 1), h.stats.queries.load(.monotonic));
try testing.expectEqual(@as(u64, 0), h.stats.formerr.load(.monotonic));
// The same query, answered from the SERVFAIL path, still echoes the OPT.
var fail: FakeUpstream = .{ .err = error.Timeout };
var h2: Handler = .{ .upstream = fail.client() };
const synthesized = try expectReply(h2.handle(undefined, .udp, query, &buf));
const p = try packet.parse(synthesized);
const opt = try edns.parseOpt(synthesized, packet.findOptRecord(p).?);
try testing.expectEqual(true, opt.do_bit);
try testing.expectEqual(@as(u16, 1232), opt.udp_payload_size);
try testing.expectEqual(@as(u64, 0), h2.stats.formerr.load(.monotonic));
}
test "every synthesized reply re-parses and binds to its request" {
const Case = struct {
query: []const u8,
err: ?transport.ExchangeError,
rcode: types.Rcode,
};
const broken_question = "\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00" ++
"\xc0\x0c\x00\x01\x00\x01";
const no_question = "\x12\x34\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00";
const update = "\x12\x34\x29\x00\x00\x01\x00\x00\x00\x00\x00\x00" ++
"\x07example\x03com\x00\x00\x01\x00\x01";
const cases = [_]Case{
.{ .query = broken_question, .err = null, .rcode = .form_err },
.{ .query = no_question, .err = null, .rcode = .form_err },
.{ .query = query_with_named_opt, .err = null, .rcode = .form_err },
.{ .query = query_with_bad_option, .err = null, .rcode = .form_err },
.{ .query = query_with_opt_in_answer, .err = null, .rcode = .form_err },
.{ .query = query_with_opt_in_authority, .err = null, .rcode = .form_err },
.{ .query = query_with_two_opts, .err = null, .rcode = .form_err },
.{ .query = update, .err = null, .rcode = .not_imp },
.{ .query = query_bytes, .err = error.BadResponse, .rcode = .serv_fail },
.{ .query = query_bytes, .err = error.Unexpected, .rcode = .serv_fail },
};
for (cases) |c| {
var fake: FakeUpstream = .{ .reply = response_bytes, .err = c.err };
var h: Handler = .{ .upstream = fake.client() };
var buf: [udp_limit_min]u8 = undefined;
const reply = try expectReply(h.handle(undefined, .udp, c.query, &buf));
const p = try packet.parse(reply);
try testing.expectEqual(c.rcode, p.header.flags.rcode);
try testing.expectEqual(@as(u16, 0x1234), p.header.id);
try testing.expectEqual(true, p.header.flags.qr);
try testing.expectEqual(true, p.header.flags.ra);
try testing.expectEqual(@as(u16, 0), p.header.ancount);
try testing.expectEqual(@as(u16, 0), p.header.nscount);
}
}
+289
View File
@@ -0,0 +1,289 @@
//! End-to-end test for the whole resolver path: a client socket on the
//! loopback, the real UDP and TCP listeners, the real handler, the real pool,
//! and two in-process fake upstreams.
//!
//! 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: every socket is on 127.0.0.1 and the upstreams are structs, so
//! nothing leaves the machine. No stream read in 0.16.0 takes a timeout, so the
//! TCP client side 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 udp_server = @import("udp_server.zig");
const packet = @import("../dns/packet.zig");
const record = @import("../dns/record.zig");
const types = @import("../dns/types.zig");
const health = @import("../upstream/health.zig");
const pool = @import("../upstream/pool.zig");
const transport = @import("../upstream/transport.zig");
const testing = std.testing;
/// 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 } };
const budget_duration: std.Io.Clock.Duration = .{ .raw = .fromSeconds(5), .clock = .awake };
/// A query for example.com A: RD set, one question, no OPT. Each step of the
/// test rewrites the ID so a reply can only match the query it belongs to.
const query_bytes =
"\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00" ++
"\x07example\x03com\x00\x00\x01\x00\x01";
const udp_id: u16 = 0x1234;
const tcp_id: u16 = 0x5678;
const backoff_id: u16 = 0x9abc;
/// The address every fake upstream answers with, and its TTL.
const answer_rdata = [4]u8{ 93, 184, 216, 34 };
const answer_ttl: u32 = 300;
/// One failure is enough to open a backoff window, so the third query has a
/// single deterministic outcome. The window outlives the test many times over.
const test_cfg: health.Config = .{
.failure_threshold = 1,
.base_backoff_ms = 60_000,
.max_backoff_ms = 60_000,
};
/// Nothing in this test is slow, so this budget only exists to stop a wedged
/// attempt from hanging the run.
const attempt_timeout: std.Io.Clock.Duration = .{ .raw = .fromSeconds(10), .clock = .awake };
fn queryWithId(buf: *[query_bytes.len]u8, id: u16) []const u8 {
buf.* = query_bytes.*;
packet.setId(buf, id);
return buf;
}
/// Echoes the question and appends one A record. This is the smallest thing a
/// real upstream could return that the handler forwards unchanged, so the
/// assertions below check bytes that travelled the whole path.
fn answerQuery(query: []const u8, response_buf: []u8) transport.ExchangeError![]u8 {
const request = packet.parse(query) catch return error.BadResponse;
const q = packet.firstQuestion(request) orelse return error.BadResponse;
var b = packet.ResponseBuilder.init(response_buf, request.header, q) catch
return error.ResponseTooLarge;
b.addAnswer(q.name, .a, .in, answer_ttl, &answer_rdata) catch
return error.ResponseTooLarge;
return b.finish();
}
/// The healthy upstream. `calls` is atomic because the listener tasks run on
/// other threads than the one asserting.
const GoodUpstream = struct {
calls: std.atomic.Value(u64) = .init(0),
fn exchangeFn(
ptr: *anyopaque,
io: std.Io,
query: []const u8,
response_buf: []u8,
) transport.ExchangeError![]u8 {
_ = io;
const self: *GoodUpstream = @ptrCast(@alignCast(ptr));
_ = self.calls.fetchAdd(1, .monotonic);
return answerQuery(query, response_buf);
}
fn client(self: *GoodUpstream) transport.Client {
return .{ .ptr = self, .exchangeFn = exchangeFn };
}
};
/// Returns `fault` on the first `fail_first` calls and answers after that.
/// `fail_first` is `maxInt` for an upstream that never recovers.
const FaultyUpstream = struct {
fault: transport.PeerFault,
fail_first: u64,
calls: std.atomic.Value(u64) = .init(0),
fn exchangeFn(
ptr: *anyopaque,
io: std.Io,
query: []const u8,
response_buf: []u8,
) transport.ExchangeError![]u8 {
_ = io;
const self: *FaultyUpstream = @ptrCast(@alignCast(ptr));
const seen = self.calls.fetchAdd(1, .monotonic);
if (seen < self.fail_first) return self.fault;
return answerQuery(query, response_buf);
}
fn client(self: *FaultyUpstream) transport.Client {
return .{ .ptr = self, .exchangeFn = exchangeFn };
}
};
fn testEntry(url: []const u8, upstream_client: transport.Client, priority: i32) pool.Entry {
return .{
.endpoint = transport.Endpoint.parse(url) catch unreachable,
.client = upstream_client,
.priority = priority,
.enabled = true,
.health = .init,
};
}
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_duration });
switch (try race.await()) {
.work => |result| return result,
.expiry => |result| {
try result;
return error.TestTimedOut;
},
}
}
/// Everything the client can check about one answer: it belongs to the query it
/// was sent for, it succeeded, and it carries the fake upstream's A record.
fn expectAnswer(reply: []const u8, id: u16) !void {
const p = try packet.parse(reply);
try testing.expectEqual(id, 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);
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);
var it = packet.answers(p);
const rec = (try it.next()) orelse return error.TestMissingAnswer;
try testing.expectEqual(types.Type.a, rec.rtype);
try testing.expectEqual(answer_ttl, rec.ttl);
try testing.expectEqual(answer_rdata, try record.rdataA(reply, rec));
}
/// One length-prefixed query and its answer on a fresh connection
/// (RFC 1035 §4.2.2).
fn tcpQuery(io: std.Io, address: net.IpAddress, id: u16) anyerror!void {
var stream = try address.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);
var query_buf: [query_bytes.len]u8 = undefined;
const query = queryWithId(&query_buf, id);
try writer.interface.writeAll(&tcp_server.framePrefix(@intCast(query.len)));
try writer.interface.writeAll(query);
try writer.interface.flush();
const len = tcp_server.parsePrefix((try reader.interface.takeArray(tcp_server.prefix_len)).*);
try expectAnswer(try reader.interface.take(len), id);
}
test "the whole resolver answers over udp and tcp and fails over to a healthy upstream" {
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 bad: FaultyUpstream = .{
.fault = error.ConnectFailed,
.fail_first = std.math.maxInt(u64),
};
var good: GoodUpstream = .{};
var entries = [_]pool.Entry{
testEntry("https://bad.example/dns-query", bad.client(), 10),
testEntry("tls://good.example", good.client(), 20),
};
var upstreams: pool.Pool = .init(&entries, test_cfg, attempt_timeout, 1);
var h: handler.Handler = .{ .upstream = upstreams.client() };
const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
var udp = try udp_server.UdpServer.bind(gpa, io, listen_address, &h, .{ .max_in_flight = 4 });
var tcp = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{ .max_connections = 2 });
const udp_address = udp.boundAddress();
const tcp_address = tcp.boundAddress();
var group: std.Io.Group = .init;
try group.concurrent(io, udp_server.UdpServer.serve, .{ &udp, io });
try group.concurrent(io, tcp_server.TcpServer.serve, .{ &tcp, 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);
{
var query_buf: [query_bytes.len]u8 = undefined;
const query = queryWithId(&query_buf, udp_id);
try client.send(io, &udp_address, query);
var buf: [udp_server.max_datagram]u8 = undefined;
const msg = try client.receiveTimeout(io, &buf, budget);
try expectAnswer(msg.data, udp_id);
}
try bounded(io, tcpQuery, .{ io, tcp_address, tcp_id });
try testing.expectEqual(@as(u64, 2), h.stats.queries.load(.monotonic));
try testing.expectEqual(@as(u64, 0), h.stats.servfail.load(.monotonic));
var snapshots: [2]pool.Snapshot = undefined;
try testing.expectEqual(@as(usize, 2), try upstreams.snapshot(io, &snapshots));
try testing.expectEqualStrings("https://bad.example/dns-query", snapshots[0].url);
try testing.expect(snapshots[0].consecutive_failures >= 1);
try testing.expect(snapshots[0].backoff_until != null);
try testing.expect(!snapshots[0].available);
try testing.expectEqualStrings("tls://good.example", snapshots[1].url);
try testing.expect(snapshots[1].total_successes >= 2);
// The failing entry is in backoff, so the third query must reach the
// healthy entry without touching it.
const bad_calls = bad.calls.load(.monotonic);
{
var query_buf: [query_bytes.len]u8 = undefined;
const query = queryWithId(&query_buf, backoff_id);
try client.send(io, &udp_address, query);
var buf: [udp_server.max_datagram]u8 = undefined;
const msg = try client.receiveTimeout(io, &buf, budget);
try expectAnswer(msg.data, backoff_id);
}
try testing.expectEqual(bad_calls, bad.calls.load(.monotonic));
try testing.expectEqual(@as(u64, 3), good.calls.load(.monotonic));
udp.deinit(gpa, io);
tcp.deinit(gpa, io);
group.cancel(io);
}
+501
View File
@@ -0,0 +1,501 @@
//! The TCP/53 listener.
//!
//! RFC 1035 §4.2.2 frames every message with a 2-byte big-endian length, and
//! RFC 7766 §6.2.1.1 lets one connection carry several queries. Both are
//! implemented here: a connection is answered serially until the client closes
//! it or the idle budget runs out.
//!
//! Connection slots are fixed and pre-allocated. Over capacity the listener
//! closes the new stream immediately and counts it; it never queues, and it
//! never allocates per connection.
//!
//! No stream read or write in 0.16.0 accepts a timeout, so every per-connection
//! operation is raced against `Options.idle_timeout` through `std.Io.Select` and
//! the loser is canceled.
const std = @import("std");
const handler = @import("handler.zig");
const transport = @import("../upstream/transport.zig");
const log = std.log.scoped(.tcp_server);
/// RFC 1035 §4.2.2: the message length prefix is two bytes, big-endian.
pub const prefix_len = 2;
/// The stream buffers only stage the framing bytes. A message longer than this
/// is read straight into `Conn.query` and written straight from `Conn.reply`,
/// so making them larger would buy nothing.
const stream_buffer_len = 1024;
/// How long the accept loop waits after an unexpected accept failure, so a
/// persistent one cannot turn the loop into a spin.
const retry_delay: std.Io.Clock.Duration = .{ .raw = .fromMilliseconds(100), .clock = .awake };
pub const Options = struct {
max_connections: u16 = 64,
/// RFC 7766 §6.2.3 recommends a few seconds of idle tolerance.
idle_timeout: std.Io.Clock.Duration = .{ .raw = .fromSeconds(10), .clock = .awake },
};
pub const Stats = struct {
accepted: std.atomic.Value(u64) = .init(0),
rejected_at_capacity: std.atomic.Value(u64) = .init(0),
rejected_at_shutdown: std.atomic.Value(u64) = .init(0),
accept_errors: std.atomic.Value(u64) = .init(0),
connection_errors: std.atomic.Value(u64) = .init(0),
idle_timeouts: std.atomic.Value(u64) = .init(0),
};
/// Lifecycle of the accept loop. `serve` claims `.serving`, `deinit` publishes
/// `.closing`, and the two meet at `stopped` so no task touches a connection
/// slot after it is freed.
const State = enum(u32) { idle, serving, closing };
/// `.closing` exists so `deinit` never shuts down a descriptor that its own
/// task is about to close: the transition to `.closing` happens under the mutex
/// before the close, and `deinit` only touches `.active` slots.
const ConnState = enum { free, active, closing };
/// What the accept loop does with a stream it has just accepted.
const Claim = union(enum) {
/// The stream owns `conns[index]`.
slot: usize,
/// Every slot is taken. The stream is closed and the loop continues.
at_capacity,
/// `deinit` has started. The stream is closed and the loop returns.
shutting_down,
};
pub const TcpServer = struct {
server: std.Io.net.Server,
handler: *handler.Handler,
conns: []Conn,
mutex: std.Io.Mutex,
/// Guarded by `mutex`. `deinit` sets it in the same critical section that
/// shuts the active connections down, so a stream that arrives after that
/// scan can never claim a slot the scan will not visit again.
shutdown_begun: bool,
options: Options,
stats: Stats,
state: std.atomic.Value(State),
stopped: std.Io.Event,
/// One slot is ~131 KiB, so the default 64 connections cost ~8.4 MiB, which
/// is inside the PLAN §18 budget. The two message buffers cannot be shared
/// or shrunk: the handler holds the query while the reply is built, and
/// both ceilings are the 65535 bytes the length prefix can express.
pub const Conn = struct {
query: [transport.max_message_len]u8,
reply: [transport.max_message_len]u8,
read_buf: [stream_buffer_len]u8,
write_buf: [stream_buffer_len]u8,
stream: std.Io.net.Stream,
/// Guarded by `TcpServer.mutex`.
state: ConnState,
};
pub const ListenError = std.Io.net.IpAddress.ListenError || error{OutOfMemory};
pub fn listen(
gpa: std.mem.Allocator,
io: std.Io,
address: std.Io.net.IpAddress,
h: *handler.Handler,
options: Options,
) ListenError!TcpServer {
std.debug.assert(options.max_connections > 0);
const conns = try gpa.alloc(Conn, options.max_connections);
errdefer gpa.free(conns);
for (conns) |*conn| conn.state = .free;
const local = address;
const server = try local.listen(io, .{ .reuse_address = true });
return .{
.server = server,
.handler = h,
.conns = conns,
.mutex = .init,
.shutdown_begun = false,
.options = options,
.stats = .{},
.state = .init(.idle),
.stopped = .unset,
};
}
/// The kernel-assigned address. A port of 0 in `listen` resolves here.
pub fn boundAddress(self: *const TcpServer) std.Io.net.IpAddress {
return self.server.socket.address;
}
/// Accept loop. Returns when the task is canceled or `deinit` stops it.
pub fn serve(self: *TcpServer, io: std.Io) void {
if (self.state.cmpxchgStrong(.idle, .serving, .acq_rel, .acquire) != null) return;
var group: std.Io.Group = .init;
self.acceptLoop(io, &group);
// A reply that is half written is worse than no reply, so the live
// connections are awaited even when this task is being canceled.
const prev = io.swapCancelProtection(.blocked);
group.await(io) catch |err| switch (err) {
error.Canceled => unreachable,
};
_ = io.swapCancelProtection(prev);
self.stopped.set(io);
}
pub fn deinit(self: *TcpServer, gpa: std.mem.Allocator, io: std.Io) void {
const was_serving = self.state.swap(.closing, .acq_rel) == .serving;
// Shutting the listening socket down is the documented way to unblock a
// pending `accept`: it fails with `error.SocketNotListening`.
const listener: std.Io.net.Stream = .{ .socket = self.server.socket };
listener.shutdown(io, .both) catch |err| {
log.debug("tcp listener shutdown failed: {t}", .{err});
};
// A live connection is blocked in a read that only the idle budget
// would end, which is seconds away. Shutting each one down bounds this,
// and the same critical section closes the door on new connections.
self.beginShutdown(io);
if (was_serving) self.stopped.waitUncancelable(io);
self.server.deinit(io);
gpa.free(self.conns);
self.* = undefined;
}
fn acceptLoop(self: *TcpServer, io: std.Io, group: *std.Io.Group) void {
while (self.state.load(.acquire) == .serving) {
const stream = self.server.accept(io) catch |err| switch (err) {
error.Canceled, error.SocketNotListening => return,
else => {
bump(&self.stats.accept_errors);
log.debug("tcp accept failed: {t}", .{err});
retry_delay.sleep(io) catch return;
continue;
},
};
const index = switch (self.claim(io, stream)) {
.slot => |index| index,
// Refusing now is honest; a queue would only hide the overload.
.at_capacity => {
bump(&self.stats.rejected_at_capacity);
stream.close(io);
continue;
},
// `deinit` will not see this stream in any slot, so serving it
// would hold `deinit` for the whole idle budget.
.shutting_down => {
bump(&self.stats.rejected_at_shutdown);
stream.close(io);
return;
},
};
group.concurrent(io, serveConn, .{ self, io, index }) catch |err| switch (err) {
error.ConcurrencyUnavailable => {
bump(&self.stats.rejected_at_capacity);
self.finish(io, index);
continue;
},
};
bump(&self.stats.accepted);
}
}
fn serveConn(self: *TcpServer, io: std.Io, index: usize) void {
defer self.finish(io, index);
const conn = &self.conns[index];
var reader = conn.stream.reader(io, &conn.read_buf);
var writer = conn.stream.writer(io, &conn.write_buf);
const budget = self.options.idle_timeout;
while (true) {
var prefix: [prefix_len]u8 = undefined;
var got: usize = 0;
switch (race(io, budget, readPrefix, .{ &reader.interface, &prefix, &got })) {
.ok => {},
.timed_out => {
bump(&self.stats.idle_timeouts);
return;
},
.canceled => return,
.failed => {
bump(&self.stats.connection_errors);
return;
},
}
// A client that closes between messages has finished asking, which
// is the normal end of a connection, not a failure.
if (got == 0) return;
if (got != prefix_len) {
bump(&self.stats.connection_errors);
return;
}
// RFC 1035 §4.2.2 gives no meaning to a zero-length message, and
// the prefix is a u16 so it can never exceed `max_message_len`.
const len = parsePrefix(prefix);
if (len == 0) {
bump(&self.stats.connection_errors);
return;
}
switch (race(io, budget, readBody, .{ &reader.interface, conn.query[0..len] })) {
.ok => {},
.canceled => return,
// A half-sent message is a broken peer, not an idle one.
.timed_out, .failed => {
bump(&self.stats.connection_errors);
return;
},
}
const bytes = switch (self.handler.handle(io, .tcp, conn.query[0..len], &conn.reply)) {
// There is no framing for "no answer", so the connection ends.
.drop => return,
.reply => |b| b,
};
const out = framePrefix(@intCast(bytes.len));
switch (race(io, budget, writeReply, .{ &writer.interface, &out, bytes })) {
.ok => {},
.canceled => return,
.timed_out, .failed => {
bump(&self.stats.connection_errors);
return;
},
}
}
}
fn claim(self: *TcpServer, io: std.Io, stream: std.Io.net.Stream) Claim {
// Uncancelable: this section takes no Io and never blocks on a peer, so
// it cannot deadlock, and losing the lock mid-update would leak a slot.
self.mutex.lockUncancelable(io);
defer self.mutex.unlock(io);
const outcome = decideClaim(self.conns, self.shutdown_begun);
switch (outcome) {
.slot => |index| {
self.conns[index].stream = stream;
self.conns[index].state = .active;
},
.at_capacity, .shutting_down => {},
}
return outcome;
}
fn finish(self: *TcpServer, io: std.Io, index: usize) void {
const conn = &self.conns[index];
self.mutex.lockUncancelable(io);
conn.state = .closing;
self.mutex.unlock(io);
// The socket is released even when this task is being torn down: the
// next cancelable call would otherwise skip the close.
const prev = io.swapCancelProtection(.blocked);
conn.stream.close(io);
_ = io.swapCancelProtection(prev);
self.mutex.lockUncancelable(io);
conn.state = .free;
self.mutex.unlock(io);
}
/// Closes the door on new connections and unblocks the live ones. Both
/// happen under one hold of the mutex: a `claim` that runs before this
/// leaves an `.active` slot the loop below shuts down, and a `claim` that
/// runs after it reads `shutdown_begun` and takes no slot at all.
fn beginShutdown(self: *TcpServer, io: std.Io) void {
self.mutex.lockUncancelable(io);
defer self.mutex.unlock(io);
self.shutdown_begun = true;
for (self.conns) |*conn| {
if (conn.state != .active) continue;
conn.stream.shutdown(io, .both) catch |err| {
log.debug("tcp connection shutdown failed: {t}", .{err});
};
}
}
};
/// RFC 1035 §4.2.2: the message length as a 2-byte big-endian integer.
pub fn framePrefix(len: u16) [prefix_len]u8 {
var out: [prefix_len]u8 = undefined;
std.mem.writeInt(u16, &out, len, .big);
return out;
}
pub fn parsePrefix(bytes: [prefix_len]u8) u16 {
return std.mem.readInt(u16, &bytes, .big);
}
/// The capacity rule, without the mutex, so it is testable without a backend.
fn firstFree(conns: []const TcpServer.Conn) ?usize {
for (conns, 0..) |*conn, index| {
if (conn.state == .free) return index;
}
return null;
}
/// The whole claim rule, without the mutex. Shutdown outranks capacity: a free
/// slot is still refused once `deinit` has passed the connections.
fn decideClaim(conns: []const TcpServer.Conn, shutdown_begun: bool) Claim {
if (shutdown_begun) return .shutting_down;
const index = firstFree(conns) orelse return .at_capacity;
return .{ .slot = index };
}
const Outcome = union(enum) {
op: anyerror!void,
expiry: std.Io.Cancelable!void,
};
const Result = enum { ok, timed_out, failed, canceled };
/// Runs one connection operation against the idle budget and cancels the loser.
fn race(
io: std.Io,
budget: std.Io.Clock.Duration,
comptime f: anytype,
args: std.meta.ArgsTuple(@TypeOf(f)),
) Result {
var outcomes: [2]Outcome = undefined;
var select: std.Io.Select(Outcome) = .init(io, &outcomes);
defer select.cancelDiscard();
select.concurrent(.op, f, args) catch |err| switch (err) {
error.ConcurrencyUnavailable => return .failed,
};
select.concurrent(.expiry, expire, .{ io, budget }) catch |err| switch (err) {
error.ConcurrencyUnavailable => return .failed,
};
return switch (select.await() catch return .canceled) {
.op => |result| if (result) |_| .ok else |err| switch (err) {
error.Canceled => .canceled,
else => .failed,
},
// A canceled sleep means this task is being torn down, not that the
// client went idle.
.expiry => |result| if (result) |_| .timed_out else |_| .canceled,
};
}
fn expire(io: std.Io, budget: std.Io.Clock.Duration) std.Io.Cancelable!void {
return budget.sleep(io);
}
/// `readSliceShort` rather than `readSliceAll`: a zero-length read is a client
/// that closed cleanly between messages, and only a partial prefix is an error.
fn readPrefix(reader: *std.Io.Reader, buf: *[prefix_len]u8, out_len: *usize) anyerror!void {
out_len.* = try reader.readSliceShort(buf);
}
fn readBody(reader: *std.Io.Reader, buf: []u8) anyerror!void {
return reader.readSliceAll(buf);
}
fn writeReply(writer: *std.Io.Writer, prefix: *const [prefix_len]u8, bytes: []const u8) anyerror!void {
try writer.writeAll(prefix);
try writer.writeAll(bytes);
try writer.flush();
}
fn bump(counter: *std.atomic.Value(u64)) void {
_ = counter.fetchAdd(1, .monotonic);
}
const testing = std.testing;
test "the length prefix is big-endian and round-trips" {
try testing.expectEqualSlices(u8, &.{ 0x00, 0x00 }, &framePrefix(0));
try testing.expectEqualSlices(u8, &.{ 0x01, 0x00 }, &framePrefix(256));
try testing.expectEqualSlices(u8, &.{ 0xff, 0xff }, &framePrefix(65535));
for ([_]u16{ 0, 1, 12, 512, 4096, 65534, 65535 }) |len| {
try testing.expectEqual(len, parsePrefix(framePrefix(len)));
}
}
test "the prefix ceiling is the message ceiling" {
try testing.expectEqual(@as(u16, transport.max_message_len), parsePrefix(.{ 0xff, 0xff }));
}
fn testConns(count: usize) ![]TcpServer.Conn {
const conns = try testing.allocator.alloc(TcpServer.Conn, count);
for (conns) |*conn| conn.state = .free;
return conns;
}
test "the connection pool hands out every slot once" {
const conns = try testConns(3);
defer testing.allocator.free(conns);
for (0..conns.len) |expected| {
const index = firstFree(conns) orelse return error.TestUnexpectedResult;
try testing.expectEqual(expected, index);
conns[index].state = .active;
}
}
test "a full connection pool refuses instead of growing" {
const conns = try testConns(2);
defer testing.allocator.free(conns);
for (conns) |*conn| conn.state = .active;
try testing.expectEqual(@as(?usize, null), firstFree(conns));
}
test "a closing slot is not reused until it is free" {
const conns = try testConns(2);
defer testing.allocator.free(conns);
conns[0].state = .active;
conns[1].state = .closing;
try testing.expectEqual(@as(?usize, null), firstFree(conns));
conns[1].state = .free;
try testing.expectEqual(@as(?usize, 1), firstFree(conns));
}
test "a claim takes the first free slot before shutdown" {
const conns = try testConns(2);
defer testing.allocator.free(conns);
conns[0].state = .active;
try testing.expectEqual(@as(usize, 1), decideClaim(conns, false).slot);
}
test "a claim after shutdown is refused even with a free slot" {
const conns = try testConns(2);
defer testing.allocator.free(conns);
try testing.expectEqual(.shutting_down, std.meta.activeTag(decideClaim(conns, true)));
// The refusal must not consume the slot: `deinit` frees it, nothing else.
try testing.expectEqual(@as(?usize, 0), firstFree(conns));
}
test "shutdown outranks capacity" {
const conns = try testConns(1);
defer testing.allocator.free(conns);
conns[0].state = .active;
try testing.expectEqual(.at_capacity, std.meta.activeTag(decideClaim(conns, false)));
try testing.expectEqual(.shutting_down, std.meta.activeTag(decideClaim(conns, true)));
}
+281
View File
@@ -0,0 +1,281 @@
//! 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(&tcp_server.framePrefix(@intCast(query_bytes.len)));
try writer.interface.writeAll(query_bytes);
try writer.interface.flush();
const len = tcp_server.parsePrefix((try reader.interface.takeArray(tcp_server.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(&tcp_server.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,
};
}
+300
View File
@@ -0,0 +1,300 @@
//! The UDP/53 listener.
//!
//! One receive loop feeds a bounded pool of in-flight tasks. A datagram that
//! finds no free slot is dropped and counted, never queued: an unbounded queue
//! only moves the failure from "visible drop" to "invisible latency", which is
//! the failure mode PLAN §1 names.
//!
//! Every slot carries its own query and reply buffer, so answering a datagram
//! allocates nothing. The pool is sized once in `bind` and never grows.
const std = @import("std");
const handler = @import("handler.zig");
const transport = @import("../upstream/transport.zig");
const log = std.log.scoped(.udp_server);
/// The EDNS ceiling this server advertises. A larger datagram arrives
/// truncated and is dropped rather than guessed at.
pub const max_datagram = 4096;
comptime {
// `handler.handle` asserts this range on the buffers it is given.
std.debug.assert(max_datagram >= handler.udp_limit_min);
std.debug.assert(max_datagram >= handler.udp_limit_max);
std.debug.assert(transport.max_message_len >= handler.udp_limit_min);
}
/// How long a receive waits before the loop re-reads the stop flag. `deinit`
/// therefore returns within this much time, without depending on a blocked
/// `receive` being unblocked by a close — a guarantee POSIX does not make.
const poll_interval: std.Io.Clock.Duration = .{ .raw = .fromMilliseconds(100), .clock = .awake };
pub const Options = struct {
max_in_flight: u16 = 64,
};
pub const Stats = struct {
received: std.atomic.Value(u64) = .init(0),
/// The datagram arrived with `flags.trunc`: the kernel discarded its tail.
dropped_oversize: std.atomic.Value(u64) = .init(0),
dropped_no_slot: std.atomic.Value(u64) = .init(0),
/// The handler answered `.drop` — an upstream exchange canceled by
/// shutdown. The handler counts its own malformed drops; this counter is
/// what keeps the listener's "every dropped datagram is counted" rule true
/// for the drops it cannot see.
dropped_handler: std.atomic.Value(u64) = .init(0),
receive_errors: std.atomic.Value(u64) = .init(0),
send_errors: std.atomic.Value(u64) = .init(0),
};
/// Lifecycle of the receive loop. `serve` claims `.serving`, `deinit` publishes
/// `.closing`, and the two meet at `stopped` so no task touches the slots after
/// they are freed.
const State = enum(u32) { idle, serving, closing };
pub const UdpServer = struct {
socket: std.Io.net.Socket,
handler: *handler.Handler,
slots: []Slot,
mutex: std.Io.Mutex,
stats: Stats,
state: std.atomic.Value(State),
stopped: std.Io.Event,
pub const Slot = struct {
query: [max_datagram]u8,
/// The buffer the handler receives the upstream answer into, so it
/// holds a whole DNS message, not a whole datagram. A reply between
/// `max_datagram` and `transport.max_message_len` must reach the
/// handler intact: only the handler can turn it into the TC=1 reply
/// that makes the client retry over TCP (S5.3). A 4096-byte buffer
/// would instead fail the exchange with `error.ResponseTooLarge` and
/// answer SERVFAIL, which is wrong.
///
/// The datagram actually sent stays bounded by `udpLimit` inside the
/// handler, so nothing larger than 4096 bytes leaves this socket.
///
/// Cost: 4096 + 65535 ≈ 68 KiB per slot, so the default 64 slots hold
/// ≈ 4.3 MiB. The PLAN §18 budget is 100 MB with ~1M blocked domains,
/// so this pool takes about 4% of it.
reply: [transport.max_message_len]u8,
from: std.Io.net.IpAddress,
len: usize,
/// Guarded by `UdpServer.mutex`.
in_use: bool,
};
pub const BindError = std.Io.net.IpAddress.BindError || error{OutOfMemory};
pub fn bind(
gpa: std.mem.Allocator,
io: std.Io,
address: std.Io.net.IpAddress,
h: *handler.Handler,
options: Options,
) BindError!UdpServer {
std.debug.assert(options.max_in_flight > 0);
const slots = try gpa.alloc(Slot, options.max_in_flight);
errdefer gpa.free(slots);
for (slots) |*slot| slot.in_use = false;
const local = address;
const socket = try local.bind(io, .{ .mode = .dgram });
return .{
.socket = socket,
.handler = h,
.slots = slots,
.mutex = .init,
.stats = .{},
.state = .init(.idle),
.stopped = .unset,
};
}
/// The kernel-assigned address. A port of 0 in `bind` resolves here.
pub fn boundAddress(self: *const UdpServer) std.Io.net.IpAddress {
return self.socket.address;
}
/// Receive loop. Returns when the task is canceled or `deinit` stops it.
pub fn serve(self: *UdpServer, io: std.Io) void {
if (self.state.cmpxchgStrong(.idle, .serving, .acq_rel, .acquire) != null) return;
var group: std.Io.Group = .init;
self.receiveLoop(io, &group);
// A reply that is half sent is worse than no reply, so the in-flight
// tasks are awaited even when this task is being canceled.
const prev = io.swapCancelProtection(.blocked);
group.await(io) catch |err| switch (err) {
error.Canceled => unreachable,
};
_ = io.swapCancelProtection(prev);
self.stopped.set(io);
}
pub fn deinit(self: *UdpServer, gpa: std.mem.Allocator, io: std.Io) void {
// `serve` reads this between receives, so it stops on its own and the
// socket is closed only once nothing can be reading from it.
if (self.state.swap(.closing, .acq_rel) == .serving) self.stopped.waitUncancelable(io);
self.socket.close(io);
gpa.free(self.slots);
self.* = undefined;
}
fn receiveLoop(self: *UdpServer, io: std.Io, group: *std.Io.Group) void {
// A datagram that finds no free slot still has to leave the socket
// buffer, or the loop would spin on the same undeliverable packet.
var overflow: [max_datagram]u8 = undefined;
while (self.state.load(.acquire) == .serving) {
const claimed = self.claim(io);
const buffer = if (claimed) |index| &self.slots[index].query else &overflow;
const msg = self.socket.receiveTimeout(io, buffer, .{ .duration = poll_interval }) catch |err| {
if (claimed) |index| self.release(io, index);
switch (err) {
error.Canceled => return,
// The stop flag is re-read at the top of the loop.
error.Timeout => continue,
else => {
bump(&self.stats.receive_errors);
log.debug("udp receive failed: {t}", .{err});
// One bad datagram must not kill the listener, and a
// persistent error must not turn it into a spin.
poll_interval.sleep(io) catch return;
continue;
},
}
};
bump(&self.stats.received);
const index = claimed orelse {
bump(&self.stats.dropped_no_slot);
continue;
};
// The kernel already threw the tail away, so the message cannot be
// parsed and any answer would be a guess.
if (msg.flags.trunc) {
bump(&self.stats.dropped_oversize);
self.release(io, index);
continue;
}
const slot = &self.slots[index];
slot.from = msg.from;
slot.len = msg.data.len;
group.concurrent(io, respondOne, .{ self, io, index }) catch |err| switch (err) {
error.ConcurrencyUnavailable => {
bump(&self.stats.dropped_no_slot);
self.release(io, index);
},
};
}
}
fn respondOne(self: *UdpServer, io: std.Io, index: usize) void {
const slot = &self.slots[index];
defer self.release(io, index);
switch (self.handler.handle(io, .udp, slot.query[0..slot.len], &slot.reply)) {
.drop => bump(&self.stats.dropped_handler),
.reply => |bytes| self.socket.send(io, &slot.from, bytes) catch |err| {
bump(&self.stats.send_errors);
log.debug("udp send to a client failed: {t}", .{err});
},
}
}
fn claim(self: *UdpServer, io: std.Io) ?usize {
// Uncancelable: this section takes no Io and never blocks on a peer, so
// it cannot deadlock, and losing the lock mid-update would leak a slot.
self.mutex.lockUncancelable(io);
defer self.mutex.unlock(io);
const index = firstFree(self.slots) orelse return null;
self.slots[index].in_use = true;
return index;
}
fn release(self: *UdpServer, io: std.Io, index: usize) void {
self.mutex.lockUncancelable(io);
defer self.mutex.unlock(io);
self.slots[index].in_use = false;
}
};
/// The capacity rule, without the mutex, so it is testable without a backend.
fn firstFree(slots: []const UdpServer.Slot) ?usize {
for (slots, 0..) |*slot, index| {
if (!slot.in_use) return index;
}
return null;
}
fn bump(counter: *std.atomic.Value(u64)) void {
_ = counter.fetchAdd(1, .monotonic);
}
const testing = std.testing;
fn testSlots(count: usize) ![]UdpServer.Slot {
const slots = try testing.allocator.alloc(UdpServer.Slot, count);
for (slots) |*slot| slot.in_use = false;
return slots;
}
test "the slot pool hands out every index once" {
const slots = try testSlots(3);
defer testing.allocator.free(slots);
for (0..slots.len) |expected| {
const index = firstFree(slots) orelse return error.TestUnexpectedResult;
try testing.expectEqual(expected, index);
slots[index].in_use = true;
}
}
test "a full slot pool refuses instead of growing" {
const slots = try testSlots(2);
defer testing.allocator.free(slots);
for (slots) |*slot| slot.in_use = true;
try testing.expectEqual(@as(?usize, null), firstFree(slots));
}
test "a slot's reply buffer holds a whole DNS message" {
const slots = try testSlots(1);
defer testing.allocator.free(slots);
// The handler asserts both ends of this range on the buffer it is given,
// and an upstream reply of up to 65535 bytes has to fit so the handler can
// truncate it instead of the exchange failing.
try testing.expectEqual(@as(usize, transport.max_message_len), slots[0].reply.len);
try testing.expect(slots[0].reply.len >= handler.udp_limit_min);
try testing.expect(slots[0].reply.len > max_datagram);
}
test "the default slot pool stays inside the memory budget" {
// 4096 + 65535 ≈ 68 KiB per slot; 64 slots ≈ 4.3 MiB, against the 100 MB
// of PLAN §18.
const options: Options = .{};
const pool_bytes = @sizeOf(UdpServer.Slot) * @as(usize, options.max_in_flight);
try testing.expect(pool_bytes < 8 * 1024 * 1024);
}
test "a released slot is handed out again" {
const slots = try testSlots(2);
defer testing.allocator.free(slots);
for (slots) |*slot| slot.in_use = true;
slots[1].in_use = false;
try testing.expectEqual(@as(?usize, 1), firstFree(slots));
}
+217
View File
@@ -0,0 +1,217 @@
//! 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 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;
/// 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,
) 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 };
}
};
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: handler.Handler = .{ .upstream = 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: handler.Handler = .{ .upstream = 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: handler.Handler = .{ .upstream = 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: handler.Handler = .{ .upstream = 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,
};
}