resolver transport: udp/tcp servers, doh/dot clients, pool failover with health
This commit is contained in:
@@ -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));
|
||||
}
|
||||
Reference in New Issue
Block a user