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