mokhtar/nxdns
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

initial commit

Browse Source
This commit is contained in:
mokhtar
2025-12-26 18:42:04 +01:00
commit d8d9ddfc53
52 changed files with 16863 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/server/handler.zig
+655
View File
@@ -0,0 +1,655 @@
const std = @import("std");
const Allocator = std.mem.Allocator;
const packet = @import("../dns/packet.zig");
const types = @import("../dns/types.zig");
const safesearch = @import("../filter/safe_search.zig");
const schema = @import("../storage/schema.zig");
const udp = @import("udp.zig");
const tcp = @import("tcp.zig");
const rate_limiter = @import("rate_limiter.zig");
/// Query log entry with full status tracking
pub const QueryLogEntry = struct {
timestamp: i64,
domain: []const u8,
client_ip: []const u8,
qtype: types.QType,
// Status tracking (Gap 1)
status: schema.QueryStatus,
// Attribution tracking (Gap 2)
list_id: ?i64,
rule_id: ?i64,
// Reply type tracking (Gap 3)
reply_type: schema.ReplyType,
// Protocol tracking (Gap 8)
protocol: schema.ClientProtocol,
// Performance
response_time_us: u64,
upstream: ?[]const u8,
reason: ?[]const u8,
/// DNSSEC validation status from upstream (AD bit)
dnssec_validated: bool,
// Legacy compatibility - computed from status
pub fn isDenied(self: QueryLogEntry) bool {
return self.status.isDenied();
}
};
/// Handler configuration
pub const HandlerConfig = struct {
blocking_response: BlockingResponse = .zero,
safe_search_enabled: bool = true,
rate_limit_enabled: bool = true,
rate_limit_qps: u32 = 20,
pub const BlockingResponse = enum {
zero, // Return 0.0.0.0
nxdomain, // Return NXDOMAIN
};
};
// Type-erased interfaces are used here to avoid circular imports:
// handler.zig defines interfaces, concrete types implement them.
// This allows handler to remain decoupled from specific implementations.
/// Result of denylist check with attribution
pub const DenyResult = struct {
denied: bool,
list_id: ?i64, // Which denylist source caused the block
rule_id: ?i64, // Which rule caused the block/allow
is_rule_allow: bool, // True if an allow rule overrode a denylist match
};
/// Interface for denylist checking
/// Gap 4: Uses u64 groups bitmask to support many-to-many client-group relationships
pub const Denylist = struct {
context: *anyopaque,
checkFn: *const fn (*anyopaque, []const u8, u64) DenyResult,
getGroupsFn: *const fn (*anyopaque, []const u8) u64,
/// Check if domain is denied for groups (bitmask), with full attribution
pub fn checkWithMask(self: Denylist, domain: []const u8, groups_mask: u64) DenyResult {
return self.checkFn(self.context, domain, groups_mask);
}
/// Legacy: Check if domain is denied for a single group
pub fn check(self: Denylist, domain: []const u8, group_id: u32) DenyResult {
const mask: u64 = if (group_id >= 64) 0 else (@as(u64, 1) << @intCast(group_id));
return self.checkWithMask(domain, mask);
}
/// Legacy compatibility - just returns boolean
pub fn isDenied(self: Denylist, domain: []const u8, group_id: u32) bool {
return self.check(domain, group_id).denied;
}
/// Get all groups for a client as a bitmask (Gap 4: many-to-many)
pub fn getGroupsForClient(self: Denylist, client_ip: []const u8) u64 {
return self.getGroupsFn(self.context, client_ip);
}
/// Legacy: Get primary group ID for a client
pub fn getGroupForClient(self: Denylist, client_ip: []const u8) u32 {
const mask = self.getGroupsForClient(client_ip);
if (mask == 1) return 0; // Only default group
var bit: u6 = 0;
while (bit < 64) : (bit += 1) {
if ((mask & (@as(u64, 1) << bit)) != 0) return bit;
}
return 0;
}
};
/// Interface for upstream DNS
pub const Upstream = struct {
context: *anyopaque,
queryFn: *const fn (*anyopaque, []const u8, Allocator) ?[]const u8,
pub fn query(self: Upstream, dns_packet: []const u8, allocator: Allocator) ?[]const u8 {
return self.queryFn(self.context, dns_packet, allocator);
}
};
/// Interface for DNS cache
/// Note: getCopy returns a mutable slice that the caller owns and must free
pub const Cache = struct {
context: *anyopaque,
getFn: *const fn (*anyopaque, []const u8, types.QType) ?[]u8,
putFn: *const fn (*anyopaque, []const u8, types.QType, []const u8, u32) void,
/// Get a copy of a cached response. Caller owns the returned memory.
pub fn getCopy(self: Cache, domain: []const u8, qtype: types.QType) ?[]u8 {
return self.getFn(self.context, domain, qtype);
}
pub fn put(self: Cache, domain: []const u8, qtype: types.QType, response: []const u8, ttl: u32) void {
self.putFn(self.context, domain, qtype, response, ttl);
}
};
/// Interface for query logging
pub const Logger = struct {
context: *anyopaque,
logFn: *const fn (*anyopaque, QueryLogEntry) void,
pub fn log(self: Logger, entry: QueryLogEntry) void {
self.logFn(self.context, entry);
}
};
/// Main DNS request handler
pub const Handler = struct {
denylist: ?Denylist,
cache: ?Cache,
upstream: ?Upstream,
logger: ?Logger,
limiter: ?*rate_limiter.RateLimiter,
config: HandlerConfig,
allocator: Allocator,
pub fn init(allocator: Allocator) Handler {
return Handler{
.denylist = null,
.cache = null,
.upstream = null,
.logger = null,
.limiter = null,
.config = HandlerConfig{},
.allocator = allocator,
};
}
pub fn setDenylist(self: *Handler, denylist: Denylist) void {
self.denylist = denylist;
}
pub fn setCache(self: *Handler, cache: Cache) void {
self.cache = cache;
}
pub fn setUpstream(self: *Handler, upstream: Upstream) void {
self.upstream = upstream;
}
pub fn setLogger(self: *Handler, logger: Logger) void {
self.logger = logger;
}
pub fn setRateLimiter(self: *Handler, limiter: *rate_limiter.RateLimiter) void {
self.limiter = limiter;
}
/// Handle a DNS query with protocol info
pub fn handleWithProtocol(
self: *Handler,
query_bytes: []const u8,
client_addr: std.net.Address,
protocol: schema.ClientProtocol,
allocator: Allocator,
) ?[]const u8 {
const start_time = std.time.microTimestamp();
// Parse the query
var query = packet.Packet.parse(query_bytes, allocator) catch |err| {
std.log.warn("Failed to parse DNS query: {}", .{err});
return self.createErrorResponse(query_bytes, types.RCode.FormErr, allocator);
};
defer query.deinit();
// Must have at least one question
if (query.questions.len == 0) {
return self.createErrorResponse(query_bytes, types.RCode.FormErr, allocator);
}
const question = query.questions[0];
// Get domain name as string (stack buffer, no allocation)
var domain_buf: [254]u8 = undefined;
const domain = question.name.toStringBuf(&domain_buf) orelse {
return self.createErrorResponse(query_bytes, types.RCode.ServFail, allocator);
};
// Get client IP as string
var client_ip_buf: [45]u8 = undefined;
const client_ip = formatAddress(client_addr, &client_ip_buf);
// Check rate limit
if (self.limiter) |limiter| {
if (!limiter.checkRequest(client_ip)) {
std.log.debug("Rate limited client: {s}", .{client_ip});
// Return REFUSED for rate-limited clients
return self.createErrorResponse(query_bytes, types.RCode.Refused, allocator);
}
}
// Get client's groups (Gap 4: many-to-many - returns bitmask)
var groups_mask: u64 = 1; // Default group (bit 0)
if (self.denylist) |dl| {
groups_mask = dl.getGroupsForClient(client_ip);
}
// Check denylist with full attribution (Gap 4: uses groups bitmask)
if (self.denylist) |dl| {
const result = dl.checkWithMask(domain, groups_mask);
if (result.denied) {
const response = self.createDeniedResponsePacket(&query, allocator) orelse return null;
// Determine status based on whether it was a denylist or rule
const status: schema.QueryStatus = if (result.rule_id != null)
.denied_rule
else
.denied_denylist;
self.logQuery(.{
.timestamp = std.time.timestamp(),
.domain = domain,
.client_ip = client_ip,
.qtype = question.qtype,
.status = status,
.list_id = result.list_id,
.rule_id = result.rule_id,
.reply_type = .ip, // Returns 0.0.0.0 or NXDOMAIN
.protocol = protocol,
.response_time_us = calcElapsedMicros(start_time),
.upstream = null,
.reason = "denylist",
.dnssec_validated = false,
});
return response;
} else if (result.is_rule_allow and result.rule_id != null) {
// Domain was on denylist but explicitly allowed by rule
// Continue to forward, but log the allow rule
}
}
// Check safe search enforcement
if (self.config.safe_search_enabled) {
if (safesearch.applySafeSearch(domain)) |safe_domain| {
const response = self.createSafeSearchResponsePacket(&query, safe_domain, allocator) orelse return null;
self.logQuery(.{
.timestamp = std.time.timestamp(),
.domain = domain,
.client_ip = client_ip,
.qtype = question.qtype,
.status = .denied_upstream, // Safe search is a kind of upstream block
.list_id = null,
.rule_id = null,
.reply_type = .cname,
.protocol = protocol,
.response_time_us = calcElapsedMicros(start_time),
.upstream = null,
.reason = "safesearch",
.dnssec_validated = false,
});
return response;
}
}
// Check cache
if (self.cache) |cache| {
if (cache.getCopy(domain, question.qtype)) |cached| {
// Update transaction ID in cached response
// cached is a mutable copy owned by caller
if (cached.len < 2) {
allocator.free(cached);
return null;
}
cached[0] = @intCast((query.header.id >> 8) & 0xFF);
cached[1] = @intCast(query.header.id & 0xFF);
self.logQuery(.{
.timestamp = std.time.timestamp(),
.domain = domain,
.client_ip = client_ip,
.qtype = question.qtype,
.status = .cached,
.list_id = null,
.rule_id = null,
.reply_type = .ip, // Assume IP for cached
.protocol = protocol,
.response_time_us = calcElapsedMicros(start_time),
.upstream = null,
.reason = "cache",
.dnssec_validated = false, // TODO: could store/retrieve from cache
});
return cached;
}
}
// Forward to upstream
if (self.upstream) |upstream| {
const response = upstream.query(query_bytes, allocator) orelse {
std.log.warn("Upstream query failed for {s}", .{domain});
self.logQuery(.{
.timestamp = std.time.timestamp(),
.domain = domain,
.client_ip = client_ip,
.qtype = question.qtype,
.status = .upstream_error,
.list_id = null,
.rule_id = null,
.reply_type = .servfail,
.protocol = protocol,
.response_time_us = calcElapsedMicros(start_time),
.upstream = null,
.reason = "upstream_failed",
.dnssec_validated = false,
});
return self.createErrorResponse(query_bytes, types.RCode.ServFail, allocator);
};
// Validate response is a valid DNS packet before processing
var response_pkt = packet.Packet.parse(response, allocator) catch |err| {
std.log.warn("Invalid upstream response for {s}: {}", .{ domain, err });
allocator.free(response);
return self.createErrorResponse(query_bytes, types.RCode.ServFail, allocator);
};
defer response_pkt.deinit();
// Determine reply type from response
const reply_type = self.detectReplyType(&response_pkt);
// Check for CNAME uncloaking - detect denied domains hiding behind CNAMEs
const dl = self.denylist orelse {
// No denylist - cache response and return it
if (self.cache) |cache| {
if (response_pkt.answers.len > 0) {
cache.put(domain, question.qtype, response, response_pkt.answers[0].ttl);
}
}
self.logQuery(.{
.timestamp = std.time.timestamp(),
.domain = domain,
.client_ip = client_ip,
.qtype = question.qtype,
.status = .forwarded,
.list_id = null,
.rule_id = null,
.reply_type = reply_type,
.protocol = protocol,
.response_time_us = calcElapsedMicros(start_time),
.upstream = "upstream",
.reason = null,
.dnssec_validated = response_pkt.header.ad,
});
return response;
};
// Check if any CNAME target is denied
if (self.findDeniedCname(&response_pkt, dl, groups_mask)) |_| {
allocator.free(response);
const denied_response = self.createDeniedResponsePacket(&query, allocator) orelse return null;
self.logQuery(.{
.timestamp = std.time.timestamp(),
.domain = domain,
.client_ip = client_ip,
.qtype = question.qtype,
.status = .denied_denylist,
.list_id = null, // TODO: Could track which list blocked the CNAME
.rule_id = null,
.reply_type = .ip,
.protocol = protocol,
.response_time_us = calcElapsedMicros(start_time),
.upstream = null,
.reason = "cname_uncloaking",
.dnssec_validated = false,
});
return denied_response;
}
// Cache the response
if (self.cache) |cache| {
if (response_pkt.answers.len > 0) {
cache.put(domain, question.qtype, response, response_pkt.answers[0].ttl);
}
}
self.logQuery(.{
.timestamp = std.time.timestamp(),
.domain = domain,
.client_ip = client_ip,
.qtype = question.qtype,
.status = .forwarded,
.list_id = null,
.rule_id = null,
.reply_type = reply_type,
.protocol = protocol,
.response_time_us = calcElapsedMicros(start_time),
.upstream = "upstream",
.reason = null,
.dnssec_validated = response_pkt.header.ad,
});
return response;
}
// No upstream configured
return self.createErrorResponse(query_bytes, types.RCode.ServFail, allocator);
}
/// Handle a DNS query (legacy - defaults to UDP protocol)
pub fn handle(
self: *Handler,
query_bytes: []const u8,
client_addr: std.net.Address,
allocator: Allocator,
) ?[]const u8 {
return self.handleWithProtocol(query_bytes, client_addr, .udp, allocator);
}
/// Detect reply type from response packet
fn detectReplyType(self: *Handler, response_pkt: *packet.Packet) schema.ReplyType {
_ = self;
// Check RCODE first
switch (response_pkt.header.rcode) {
.ServFail => return .servfail,
.NXDomain => return .nxdomain,
.Refused => return .refused,
else => {},
}
// No answers = NODATA
if (response_pkt.answers.len == 0) {
return .nodata;
}
// Check answer types
for (response_pkt.answers) |answer| {
if (answer.rtype == types.QType.A or answer.rtype == types.QType.AAAA) {
return .ip;
}
if (answer.rtype == types.QType.CNAME) {
return .cname;
}
}
return .unknown;
}
/// Check if any CNAME target in the response is denied
/// Checks ANSWER, AUTHORITY, and ADDITIONAL sections for completeness
/// Returns true if a denied CNAME was found
/// Gap 4: Takes groups bitmask for many-to-many support
fn findDeniedCname(self: *Handler, response_pkt: *packet.Packet, dl: Denylist, groups_mask: u64) ?bool {
_ = self;
// Check all sections where CNAMEs could appear
const sections = [_][]const packet.ResourceRecord{
response_pkt.answers,
response_pkt.authority,
response_pkt.additional,
};
for (sections) |section| {
for (section) |record| {
if (record.rtype != types.QType.CNAME) continue;
const cname = record.getCname() orelse continue;
var cname_buf: [254]u8 = undefined;
const cname_str = cname.toStringBuf(&cname_buf) orelse continue;
if (dl.checkWithMask(cname_str, groups_mask).denied) {
return true;
}
}
}
return null;
}
/// Calculate elapsed microseconds since start_time, handling potential negative values
fn calcElapsedMicros(start_time: i64) u64 {
const elapsed = std.time.microTimestamp() - start_time;
return if (elapsed > 0) @intCast(elapsed) else 0;
}
fn createDeniedResponsePacket(self: *Handler, query: *const packet.Packet, allocator: Allocator) ?[]const u8 {
var response = switch (self.config.blocking_response) {
.zero => packet.Packet.createDeniedResponse(query, allocator) catch return null,
.nxdomain => packet.Packet.createNxdomainResponse(query, allocator) catch return null,
};
defer response.deinit();
var buffer: [types.EDNS_DEFAULT_SIZE]u8 = undefined;
const len = response.encode(&buffer) catch return null;
return allocator.dupe(u8, buffer[0..len]) catch null;
}
fn createSafeSearchResponsePacket(self: *Handler, query: *const packet.Packet, safe_domain: []const u8, allocator: Allocator) ?[]const u8 {
_ = self;
var response = packet.Packet.createSafeSearchResponse(query, safe_domain, allocator) catch return null;
defer response.deinit();
var buffer: [types.EDNS_DEFAULT_SIZE]u8 = undefined;
const len = response.encode(&buffer) catch return null;
return allocator.dupe(u8, buffer[0..len]) catch null;
}
fn createErrorResponse(self: *Handler, query_bytes: []const u8, rcode: types.RCode, allocator: Allocator) ?[]const u8 {
_ = self;
if (query_bytes.len < types.DNS_HEADER_SIZE) return null;
// Create a minimal error response
var response_buf: [types.DNS_HEADER_SIZE]u8 = undefined;
@memcpy(&response_buf, query_bytes[0..types.DNS_HEADER_SIZE]);
// Set QR = 1 (response), RA = 1 (recursion available), and RCODE
response_buf[2] |= 0x80; // QR = 1
response_buf[3] = (response_buf[3] & 0xF0) | @intFromEnum(rcode);
response_buf[3] |= 0x80; // RA = 1
// Set counts to 0 for answers, authority, additional
response_buf[6] = 0;
response_buf[7] = 0;
response_buf[8] = 0;
response_buf[9] = 0;
response_buf[10] = 0;
response_buf[11] = 0;
return allocator.dupe(u8, &response_buf) catch null;
}
fn logQuery(self: *Handler, entry: QueryLogEntry) void {
if (self.logger) |logger| {
logger.log(entry);
}
}
/// Create a UDP handler wrapper
pub fn toUdpHandler(self: *Handler) udp.UdpServer.Handler {
return udp.UdpServer.Handler{
.context = self,
.handleFn = handleUdp,
};
}
/// Create a TCP handler wrapper
pub fn toTcpHandler(self: *Handler) tcp.TcpServer.Handler {
return tcp.TcpServer.Handler{
.context = self,
.handleFn = handleTcp,
};
}
fn handleUdp(ctx: *anyopaque, query: []const u8, client_addr: std.net.Address, allocator: Allocator) ?[]const u8 {
const self: *Handler = @ptrCast(@alignCast(ctx));
return self.handleWithProtocol(query, client_addr, .udp, allocator);
}
fn handleTcp(ctx: *anyopaque, query: []const u8, client_addr: std.net.Address, allocator: Allocator) ?[]const u8 {
const self: *Handler = @ptrCast(@alignCast(ctx));
return self.handleWithProtocol(query, client_addr, .tcp, allocator);
}
};
/// Format a network address to a string representation
/// Returns "0.0.0.0" as fallback for unknown address families (safer than "unknown" for rate limiting)
fn formatAddress(addr: std.net.Address, buf: []u8) []const u8 {
if (addr.any.family == std.posix.AF.INET) {
// IPv4 address
const bytes = @as(*const [4]u8, @ptrCast(&addr.in.sa.addr));
const result = std.fmt.bufPrint(buf, "{d}.{d}.{d}.{d}", .{
bytes[0], bytes[1], bytes[2], bytes[3],
}) catch return "0.0.0.0";
return result;
} else if (addr.any.family == std.posix.AF.INET6) {
// IPv6 address - use compressed format for common cases
const bytes = @as(*const [16]u8, @ptrCast(&addr.in6.sa.addr));
const result = std.fmt.bufPrint(buf, "{x:0>2}{x:0>2}:{x:0>2}{x:0>2}:{x:0>2}{x:0>2}:{x:0>2}{x:0>2}:{x:0>2}{x:0>2}:{x:0>2}{x:0>2}:{x:0>2}{x:0>2}:{x:0>2}{x:0>2}", .{
bytes[0], bytes[1], bytes[2], bytes[3],
bytes[4], bytes[5], bytes[6], bytes[7],
bytes[8], bytes[9], bytes[10], bytes[11],
bytes[12], bytes[13], bytes[14], bytes[15],
}) catch return "::";
return result;
}
// Unknown address family - use safe fallback that won't break rate limiting
std.log.warn("Unknown address family: {}", .{addr.any.family});
return "0.0.0.0";
}
test "Handler basic test" {
const testing = std.testing;
const allocator = testing.allocator;
var handler = Handler.init(allocator);
// Create a simple query
const query_bytes = [_]u8{
0x00, 0x01, // ID
0x01, 0x00, // Standard query, RD=1
0x00, 0x01, // QDCOUNT = 1
0x00, 0x00, // ANCOUNT = 0
0x00, 0x00, // NSCOUNT = 0
0x00, 0x00, // ARCOUNT = 0
0x07, 'e', 'x', 'a', 'm', 'p', 'l', 'e', // example
0x03, 'c', 'o', 'm', // com
0x00, // null
0x00, 0x01, // TYPE = A
0x00, 0x01, // CLASS = IN
};
const addr = std.net.Address.initIp4([4]u8{ 127, 0, 0, 1 }, 12345);
// Without upstream, should return SERVFAIL
const response = handler.handle(&query_bytes, addr, allocator);
if (response) |r| {
defer allocator.free(r);
try testing.expect(r.len >= types.DNS_HEADER_SIZE);
}
}
src/server/rate_limiter.zig
+241
View File
@@ -0,0 +1,241 @@
const std = @import("std");
const Allocator = std.mem.Allocator;
/// Rate limiter to prevent DNS amplification attacks
/// Uses a sliding window algorithm with per-IP tracking
pub const RateLimiter = struct {
/// Per-IP request tracking
clients: std.StringHashMapUnmanaged(ClientState),
allocator: Allocator,
mutex: std.Thread.Mutex,
/// Configuration
config: Config,
/// Statistics
stats: Stats,
pub const Config = struct {
/// Maximum requests per second per IP
max_qps: u32 = 20,
/// Time window for rate limiting (in milliseconds)
window_ms: u64 = 1000,
/// Maximum number of IPs to track (LRU eviction beyond this)
max_clients: usize = 10000,
/// Enable rate limiting
enabled: bool = true,
};
const ClientState = struct {
/// Request timestamps in the current window (circular buffer)
timestamps: [32]i64,
/// Current position in the buffer
pos: usize,
/// Number of requests in current window
count: u32,
/// Last activity timestamp (for LRU eviction)
last_seen: i64,
};
pub const Stats = struct {
/// Total requests processed
total_requests: u64 = 0,
/// Requests that were rate limited
rate_limited: u64 = 0,
/// Currently tracked clients
active_clients: usize = 0,
};
pub fn init(allocator: Allocator) RateLimiter {
return initWithConfig(allocator, Config{});
}
pub fn initWithConfig(allocator: Allocator, config: Config) RateLimiter {
return RateLimiter{
.clients = std.StringHashMapUnmanaged(ClientState){},
.allocator = allocator,
.mutex = std.Thread.Mutex{},
.config = config,
.stats = Stats{},
};
}
pub fn deinit(self: *RateLimiter) void {
var iter = self.clients.iterator();
while (iter.next()) |entry| {
self.allocator.free(entry.key_ptr.*);
}
self.clients.deinit(self.allocator);
}
/// Check if a request from this IP should be allowed
/// Returns true if allowed, false if rate limited
pub fn checkRequest(self: *RateLimiter, client_ip: []const u8) bool {
if (!self.config.enabled) return true;
self.mutex.lock();
defer self.mutex.unlock();
self.stats.total_requests += 1;
const now = std.time.milliTimestamp();
const window_start = now - @as(i64, @intCast(self.config.window_ms));
// Get or create client state
const gop = self.clients.getOrPut(self.allocator, client_ip) catch {
// On allocation failure, allow the request
return true;
};
if (!gop.found_existing) {
// New client - copy the key
gop.key_ptr.* = self.allocator.dupe(u8, client_ip) catch {
_ = self.clients.remove(client_ip);
return true;
};
// Initialize state
gop.value_ptr.* = ClientState{
.timestamps = [_]i64{0} ** 32,
.pos = 0,
.count = 0,
.last_seen = now,
};
// Evict oldest client if at capacity
if (self.clients.count() > self.config.max_clients) {
self.evictOldestLocked();
}
}
var state = gop.value_ptr;
state.last_seen = now;
// Count requests in current window
var count: u32 = 0;
for (state.timestamps) |ts| {
if (ts > window_start) {
count += 1;
}
}
// Check rate limit
if (count >= self.config.max_qps) {
self.stats.rate_limited += 1;
return false;
}
// Record this request
state.timestamps[state.pos] = now;
state.pos = (state.pos + 1) % state.timestamps.len;
state.count = count + 1;
self.stats.active_clients = self.clients.count();
return true;
}
/// Evict the oldest (least recently seen) client
fn evictOldestLocked(self: *RateLimiter) void {
var oldest_key: ?[]const u8 = null;
var oldest_time: i64 = std.math.maxInt(i64);
var iter = self.clients.iterator();
while (iter.next()) |entry| {
if (entry.value_ptr.last_seen < oldest_time) {
oldest_time = entry.value_ptr.last_seen;
oldest_key = entry.key_ptr.*;
}
}
if (oldest_key) |key| {
if (self.clients.fetchRemove(key)) |kv| {
self.allocator.free(kv.key);
}
}
}
/// Get current statistics
pub fn getStats(self: *RateLimiter) Stats {
self.mutex.lock();
defer self.mutex.unlock();
return self.stats;
}
/// Reset statistics
pub fn resetStats(self: *RateLimiter) void {
self.mutex.lock();
defer self.mutex.unlock();
self.stats = Stats{};
}
/// Clear all tracked clients
pub fn clear(self: *RateLimiter) void {
self.mutex.lock();
defer self.mutex.unlock();
var iter = self.clients.iterator();
while (iter.next()) |entry| {
self.allocator.free(entry.key_ptr.*);
}
self.clients.clearRetainingCapacity();
self.stats.active_clients = 0;
}
};
test "RateLimiter basic operations" {
const testing = std.testing;
const allocator = testing.allocator;
var limiter = RateLimiter.initWithConfig(allocator, .{
.max_qps = 5,
.window_ms = 1000,
});
defer limiter.deinit();
// First 5 requests should be allowed
for (0..5) |_| {
try testing.expect(limiter.checkRequest("192.168.1.1"));
}
// 6th request should be rate limited
try testing.expect(!limiter.checkRequest("192.168.1.1"));
// Different IP should still be allowed
try testing.expect(limiter.checkRequest("192.168.1.2"));
}
test "RateLimiter disabled" {
const testing = std.testing;
const allocator = testing.allocator;
var limiter = RateLimiter.initWithConfig(allocator, .{
.max_qps = 1,
.enabled = false,
});
defer limiter.deinit();
// All requests should be allowed when disabled
for (0..100) |_| {
try testing.expect(limiter.checkRequest("192.168.1.1"));
}
}
test "RateLimiter statistics" {
const testing = std.testing;
const allocator = testing.allocator;
var limiter = RateLimiter.initWithConfig(allocator, .{
.max_qps = 2,
.window_ms = 1000,
});
defer limiter.deinit();
_ = limiter.checkRequest("192.168.1.1"); // allowed
_ = limiter.checkRequest("192.168.1.1"); // allowed
_ = limiter.checkRequest("192.168.1.1"); // rate limited
const stats = limiter.getStats();
try testing.expectEqual(@as(u64, 3), stats.total_requests);
try testing.expectEqual(@as(u64, 1), stats.rate_limited);
}
src/server/shutdown.zig
+117
View File
@@ -0,0 +1,117 @@
const std = @import("std");
const posix = std.posix;
const linux = std.os.linux;
/// Unified shutdown coordinator using signalfd + eventfd
/// - signalfd: catches SIGINT/SIGTERM (main thread waits on this)
/// - eventfd: wakes all worker threads when shutdown requested
pub const ShutdownCoordinator = struct {
shutdown_requested: std.atomic.Value(bool),
signal_fd: posix.fd_t,
event_fd: posix.fd_t,
pub fn init() ShutdownCoordinator {
// Block SIGINT/SIGTERM so they go to signalfd instead of default handler
var mask = std.mem.zeroes(linux.sigset_t);
linux.sigaddset(&mask, linux.SIG.INT);
linux.sigaddset(&mask, linux.SIG.TERM);
_ = linux.sigprocmask(linux.SIG.BLOCK, &mask, null);
// Create signalfd
const sig_fd_raw = linux.signalfd(-1, &mask, linux.SFD.CLOEXEC);
const sig_fd: posix.fd_t = if (@as(isize, @bitCast(sig_fd_raw)) < 0) -1 else @intCast(sig_fd_raw);
// Create eventfd for waking threads
const evt_fd = posix.eventfd(0, linux.EFD.CLOEXEC) catch -1;
return ShutdownCoordinator{
.shutdown_requested = std.atomic.Value(bool).init(false),
.signal_fd = sig_fd,
.event_fd = evt_fd,
};
}
pub fn deinit(self: *ShutdownCoordinator) void {
if (self.signal_fd != -1) posix.close(self.signal_fd);
if (self.event_fd != -1) posix.close(self.event_fd);
}
/// Signal shutdown and wake all waiting threads
pub fn requestShutdown(self: *ShutdownCoordinator) void {
self.shutdown_requested.store(true, .release);
// Wake all threads waiting on eventfd
if (self.event_fd != -1) {
const val: u64 = 1;
_ = posix.write(self.event_fd, std.mem.asBytes(&val)) catch {};
}
}
pub fn isShutdownRequested(self: *ShutdownCoordinator) bool {
return self.shutdown_requested.load(.acquire);
}
/// Get eventfd for polling (worker threads use this)
pub fn getEventFd(self: *ShutdownCoordinator) posix.fd_t {
return self.event_fd;
}
/// Get signalfd for polling (main thread uses this)
pub fn getSignalFd(self: *ShutdownCoordinator) posix.fd_t {
return self.signal_fd;
}
/// Block until shutdown signal received (main thread calls this)
/// Returns true if signal was received, false on error
pub fn waitForSignal(self: *ShutdownCoordinator) bool {
if (self.signal_fd == -1) return false;
var fds = [1]posix.pollfd{
.{ .fd = self.signal_fd, .events = posix.POLL.IN, .revents = 0 },
};
_ = posix.poll(&fds, -1) catch return false;
if (fds[0].revents & posix.POLL.IN != 0) {
// Consume the signal
var siginfo: linux.signalfd_siginfo = undefined;
_ = posix.read(self.signal_fd, std.mem.asBytes(&siginfo)) catch {};
self.requestShutdown();
return true;
}
return false;
}
/// Block until shutdown is requested or timeout (worker threads use this)
/// timeout_ms: -1 for infinite wait
pub fn wait(self: *ShutdownCoordinator, timeout_ms: i32) bool {
if (self.isShutdownRequested()) return true;
if (self.event_fd == -1) return false;
var fds = [1]posix.pollfd{
.{ .fd = self.event_fd, .events = posix.POLL.IN, .revents = 0 },
};
_ = posix.poll(&fds, timeout_ms) catch return self.isShutdownRequested();
return self.isShutdownRequested();
}
};
test "ShutdownCoordinator basic" {
var coordinator = ShutdownCoordinator.init();
defer coordinator.deinit();
try std.testing.expect(!coordinator.isShutdownRequested());
coordinator.requestShutdown();
try std.testing.expect(coordinator.isShutdownRequested());
}
test "ShutdownCoordinator wait returns immediately after shutdown" {
var coordinator = ShutdownCoordinator.init();
defer coordinator.deinit();
coordinator.requestShutdown();
const result = coordinator.wait(1000);
try std.testing.expect(result);
}
src/server/tcp.zig
+373
View File
@@ -0,0 +1,373 @@
const std = @import("std");
const posix = std.posix;
const net = std.net;
const Allocator = std.mem.Allocator;
const packet = @import("../dns/packet.zig");
const types = @import("../dns/types.zig");
/// Connection task for the worker pool
const ConnectionTask = struct {
handle: posix.socket_t,
address: net.Address,
};
/// Simple bounded work queue for connection tasks
const WorkQueue = struct {
items: [QUEUE_SIZE]?ConnectionTask,
head: usize,
tail: usize,
count: usize,
mutex: std.Thread.Mutex,
not_empty: std.Thread.Condition,
not_full: std.Thread.Condition,
const QUEUE_SIZE = 64;
fn init() WorkQueue {
return .{
.items = [_]?ConnectionTask{null} ** QUEUE_SIZE,
.head = 0,
.tail = 0,
.count = 0,
.mutex = .{},
.not_empty = .{},
.not_full = .{},
};
}
/// Push a task, returns false if queue is full
fn tryPush(self: *WorkQueue, task: ConnectionTask) bool {
self.mutex.lock();
defer self.mutex.unlock();
if (self.count >= QUEUE_SIZE) {
return false;
}
self.items[self.tail] = task;
self.tail = (self.tail + 1) % QUEUE_SIZE;
self.count += 1;
self.not_empty.signal();
return true;
}
/// Pop a task, blocks until available or timeout
fn pop(self: *WorkQueue, running: *std.atomic.Value(bool)) ?ConnectionTask {
self.mutex.lock();
defer self.mutex.unlock();
while (self.count == 0) {
if (!running.load(.acquire)) {
return null;
}
// Wait with timeout to periodically check running flag
self.not_empty.timedWait(&self.mutex, 100 * std.time.ns_per_ms) catch {};
}
if (self.count == 0) return null;
const task = self.items[self.head];
self.items[self.head] = null;
self.head = (self.head + 1) % QUEUE_SIZE;
self.count -= 1;
self.not_full.signal();
return task;
}
fn wakeAll(self: *WorkQueue) void {
self.mutex.lock();
defer self.mutex.unlock();
self.not_empty.broadcast();
}
};
pub const TcpServer = struct {
listener: net.Server,
allocator: Allocator,
handler: *Handler,
running: std.atomic.Value(bool),
active_connections: std.atomic.Value(u32),
max_connections: u32,
connection_timeout_ms: u32,
work_queue: WorkQueue,
workers: []std.Thread,
num_workers: u32,
/// Default maximum concurrent TCP connections
pub const DEFAULT_MAX_CONNECTIONS: u32 = 100;
/// Default connection read/write timeout (30 seconds)
pub const DEFAULT_CONNECTION_TIMEOUT_MS: u32 = 30000;
/// Default number of worker threads
pub const DEFAULT_NUM_WORKERS: u32 = 8;
pub const Handler = struct {
context: *anyopaque,
handleFn: *const fn (*anyopaque, []const u8, std.net.Address, Allocator) ?[]const u8,
pub fn handle(self: Handler, query: []const u8, client_addr: std.net.Address, allocator: Allocator) ?[]const u8 {
return self.handleFn(self.context, query, client_addr, allocator);
}
};
pub const InitError = error{
ListenFailed,
} || net.Address.ListenError;
/// Initialize the TCP server
pub fn init(bind_addr: net.Address, handler: *Handler, allocator: Allocator) InitError!TcpServer {
return initWithConfig(bind_addr, handler, allocator, .{});
}
pub const Config = struct {
max_connections: u32 = DEFAULT_MAX_CONNECTIONS,
connection_timeout_ms: u32 = DEFAULT_CONNECTION_TIMEOUT_MS,
num_workers: u32 = DEFAULT_NUM_WORKERS,
};
/// Initialize the TCP server with custom configuration
pub fn initWithConfig(bind_addr: net.Address, handler: *Handler, allocator: Allocator, config: Config) InitError!TcpServer {
const listener = bind_addr.listen(.{
.reuse_address = true,
}) catch {
return error.ListenFailed;
};
return TcpServer{
.listener = listener,
.allocator = allocator,
.handler = handler,
.running = std.atomic.Value(bool).init(false),
.active_connections = std.atomic.Value(u32).init(0),
.max_connections = config.max_connections,
.connection_timeout_ms = config.connection_timeout_ms,
.work_queue = WorkQueue.init(),
.workers = &[_]std.Thread{},
.num_workers = config.num_workers,
};
}
/// Start the server loop
pub fn run(self: *TcpServer) !void {
self.running.store(true, .release);
// Start worker threads
self.workers = self.allocator.alloc(std.Thread, self.num_workers) catch |err| {
std.log.err("TCP: failed to allocate worker threads: {}", .{err});
return error.OutOfMemory;
};
errdefer self.allocator.free(self.workers);
var started: u32 = 0;
errdefer {
self.running.store(false, .release);
self.work_queue.wakeAll();
for (self.workers[0..started]) |w| w.join();
}
for (self.workers) |*worker| {
worker.* = std.Thread.spawn(.{}, workerLoop, .{self}) catch |err| {
std.log.err("TCP: failed to start worker thread: {}", .{err});
return error.ThreadSpawnFailed;
};
started += 1;
}
std.log.info("TCP: started {} worker threads", .{self.num_workers});
while (self.running.load(.acquire)) {
// Use poll with timeout to allow checking running flag
var fds = [1]posix.pollfd{
.{
.fd = self.listener.stream.handle,
.events = posix.POLL.IN,
.revents = 0,
},
};
const poll_result = posix.poll(&fds, 100) catch |err| {
std.log.warn("TCP poll error: {}", .{err});
continue;
};
// Timeout - check running flag and continue
if (poll_result == 0) continue;
// No connection pending
if (fds[0].revents & posix.POLL.IN == 0) continue;
const conn = self.listener.accept() catch |err| {
std.log.warn("TCP accept error: {}", .{err});
continue;
};
// Reserve a connection slot atomically (prevents TOCTOU race)
// fetchAdd returns the OLD value, so new count is old + 1
const old_count = self.active_connections.fetchAdd(1, .acq_rel);
if (old_count >= self.max_connections) {
// Over limit - rollback reservation and reject
_ = self.active_connections.fetchSub(1, .release);
std.log.warn("TCP: max connections ({}) reached, rejecting new connection", .{self.max_connections});
posix.close(conn.stream.handle);
continue;
}
// Submit to worker pool
const task = ConnectionTask{
.handle = conn.stream.handle,
.address = conn.address,
};
if (!self.work_queue.tryPush(task)) {
// Queue full - rollback reservation and reject
_ = self.active_connections.fetchSub(1, .release);
std.log.warn("TCP: work queue full, rejecting connection", .{});
posix.close(conn.stream.handle);
}
}
// Shutdown: wait for workers
self.work_queue.wakeAll();
for (self.workers) |w| w.join();
self.allocator.free(self.workers);
self.workers = &[_]std.Thread{};
}
/// Worker thread loop
fn workerLoop(self: *TcpServer) void {
while (self.running.load(.acquire)) {
if (self.work_queue.pop(&self.running)) |task| {
self.handleConnectionTask(task);
}
}
}
/// Maximum DNS message size for TCP (RFC 1035: 2-byte length prefix allows up to 65535)
const MAX_DNS_MESSAGE_SIZE: usize = 65535;
fn handleConnectionTask(self: *TcpServer, task: ConnectionTask) void {
// Count was already incremented in accept loop when slot was reserved
defer _ = self.active_connections.fetchSub(1, .release);
defer posix.close(task.handle);
// Set read/write timeout on the connection to prevent slow clients from blocking
const timeout = posix.timeval{
.sec = @intCast(self.connection_timeout_ms / 1000),
.usec = @intCast((self.connection_timeout_ms % 1000) * 1000),
};
posix.setsockopt(task.handle, posix.SOL.SOCKET, posix.SO.RCVTIMEO, std.mem.asBytes(&timeout)) catch {};
posix.setsockopt(task.handle, posix.SOL.SOCKET, posix.SO.SNDTIMEO, std.mem.asBytes(&timeout)) catch {};
var buffer: [MAX_DNS_MESSAGE_SIZE]u8 = undefined;
while (self.running.load(.acquire)) {
// Read 2-byte length prefix using posix
var len_buf: [2]u8 = undefined;
if (!readExact(task.handle, &len_buf)) break;
const length = std.mem.readInt(u16, &len_buf, .big);
// Validate: must be at least a DNS header, and within buffer capacity
if (length < types.DNS_HEADER_SIZE) {
std.log.debug("TCP: message too small ({} bytes), dropping connection", .{length});
break;
}
if (length > MAX_DNS_MESSAGE_SIZE) {
std.log.warn("TCP: message too large ({} bytes), dropping connection", .{length});
break;
}
// Read DNS message using posix
if (!readExact(task.handle, buffer[0..length])) break;
// Handle the query
const response = self.handler.handle(
buffer[0..length],
task.address,
self.allocator,
) orelse continue;
defer self.allocator.free(response);
// Validate response fits in TCP DNS message (u16 length prefix)
if (response.len > MAX_DNS_MESSAGE_SIZE) {
std.log.warn("TCP: response too large ({} bytes), dropping", .{response.len});
continue;
}
// Write length-prefixed response
var resp_len: [2]u8 = undefined;
std.mem.writeInt(u16, &resp_len, @intCast(response.len), .big);
if (!writeAll(task.handle, &resp_len)) break;
if (!writeAll(task.handle, response)) break;
}
}
/// Stop the server
pub fn stop(self: *TcpServer) void {
self.running.store(false, .release);
self.work_queue.wakeAll();
}
/// Wait for active connections to finish (with timeout)
pub fn waitForConnections(self: *TcpServer, timeout_ms: u64) void {
const start = std.time.milliTimestamp();
while (self.active_connections.load(.acquire) > 0) {
const elapsed: u64 = @intCast(std.time.milliTimestamp() - start);
if (elapsed >= timeout_ms) {
std.log.warn("TCP: {} connections still active after timeout", .{self.active_connections.load(.acquire)});
break;
}
std.posix.nanosleep(0, 10 * std.time.ns_per_ms);
}
}
/// Clean up resources
pub fn deinit(self: *TcpServer) void {
self.listener.deinit();
}
};
/// Read exactly the requested number of bytes using posix
fn readExact(handle: posix.socket_t, buf: []u8) bool {
var total_read: usize = 0;
while (total_read < buf.len) {
const n = posix.read(handle, buf[total_read..]) catch return false;
if (n == 0) return false; // EOF
total_read += n;
}
return true;
}
/// Write all bytes using posix
fn writeAll(handle: posix.socket_t, buf: []const u8) bool {
var total_written: usize = 0;
while (total_written < buf.len) {
const n = posix.write(handle, buf[total_written..]) catch return false;
if (n == 0) return false;
total_written += n;
}
return true;
}
test "TcpServer init" {
const testing = std.testing;
const allocator = testing.allocator;
const Handler = TcpServer.Handler;
var handler = Handler{
.context = undefined,
.handleFn = struct {
fn handle(_: *anyopaque, _: []const u8, _: std.net.Address, _: Allocator) ?[]const u8 {
return null;
}
}.handle,
};
const addr = net.Address.initIp4(.{ 127, 0, 0, 1 }, 0);
var server = TcpServer.init(addr, &handler, allocator) catch |err| {
std.debug.print("TcpServer init failed: {}\n", .{err});
return error.TestFailed;
};
defer server.deinit();
}
src/server/udp.zig
+512
View File
@@ -0,0 +1,512 @@
const std = @import("std");
const posix = std.posix;
const Allocator = std.mem.Allocator;
const packet = @import("../dns/packet.zig");
const types = @import("../dns/types.zig");
/// UDP query task for worker pool
const QueryTask = struct {
data: [types.EDNS_DEFAULT_SIZE]u8,
len: usize,
src_addr: posix.sockaddr,
addr_len: posix.socklen_t,
};
/// Simple bounded work queue for UDP query tasks
const WorkQueue = struct {
items: [QUEUE_SIZE]?QueryTask,
head: usize,
tail: usize,
count: usize,
mutex: std.Thread.Mutex,
not_empty: std.Thread.Condition,
const QUEUE_SIZE = 64;
fn init() WorkQueue {
return .{
.items = [_]?QueryTask{null} ** QUEUE_SIZE,
.head = 0,
.tail = 0,
.count = 0,
.mutex = .{},
.not_empty = .{},
};
}
fn tryPush(self: *WorkQueue, task: QueryTask) bool {
self.mutex.lock();
defer self.mutex.unlock();
if (self.count >= QUEUE_SIZE) {
return false;
}
self.items[self.tail] = task;
self.tail = (self.tail + 1) % QUEUE_SIZE;
self.count += 1;
self.not_empty.signal();
return true;
}
fn pop(self: *WorkQueue, running: *std.atomic.Value(bool)) ?QueryTask {
self.mutex.lock();
defer self.mutex.unlock();
while (self.count == 0) {
if (!running.load(.acquire)) {
return null;
}
self.not_empty.timedWait(&self.mutex, 100 * std.time.ns_per_ms) catch {};
}
if (self.count == 0) return null;
const task = self.items[self.head];
self.items[self.head] = null;
self.head = (self.head + 1) % QUEUE_SIZE;
self.count -= 1;
return task;
}
fn wakeAll(self: *WorkQueue) void {
self.mutex.lock();
defer self.mutex.unlock();
self.not_empty.broadcast();
}
};
pub const UdpServer = struct {
socket: posix.socket_t,
allocator: Allocator,
handler: *Handler,
running: std.atomic.Value(bool),
work_queue: WorkQueue,
workers: []std.Thread,
num_workers: u32,
dropped_queries: std.atomic.Value(u64),
/// Default number of worker threads
pub const DEFAULT_NUM_WORKERS: u32 = 8;
pub const Handler = struct {
context: *anyopaque,
handleFn: *const fn (*anyopaque, []const u8, std.net.Address, Allocator) ?[]const u8,
pub fn handle(self: Handler, query: []const u8, client_addr: std.net.Address, allocator: Allocator) ?[]const u8 {
return self.handleFn(self.context, query, client_addr, allocator);
}
};
pub const InitError = error{
SocketCreationFailed,
SetSockOptFailed,
BindFailed,
} || posix.SocketError || posix.SetSockOptError;
pub const Config = struct {
num_workers: u32 = DEFAULT_NUM_WORKERS,
};
/// Initialize the UDP server
pub fn init(bind_addr: std.net.Address, handler: *Handler, allocator: Allocator) InitError!UdpServer {
return initWithConfig(bind_addr, handler, allocator, .{});
}
/// Initialize the UDP server with custom configuration
pub fn initWithConfig(bind_addr: std.net.Address, handler: *Handler, allocator: Allocator, config: Config) InitError!UdpServer {
// Create UDP socket
const socket = try posix.socket(
bind_addr.any.family,
posix.SOCK.DGRAM,
0,
);
errdefer posix.close(socket);
// Bind to address (no SO_REUSEADDR - we want bind to fail if another instance is running)
posix.bind(socket, &bind_addr.any, bind_addr.getOsSockLen()) catch {
return error.BindFailed;
};
return UdpServer{
.socket = socket,
.allocator = allocator,
.handler = handler,
.running = std.atomic.Value(bool).init(false),
.work_queue = WorkQueue.init(),
.workers = &[_]std.Thread{},
.num_workers = config.num_workers,
.dropped_queries = std.atomic.Value(u64).init(0),
};
}
/// Get the count of dropped queries (for monitoring)
pub fn getDroppedQueries(self: *UdpServer) u64 {
return self.dropped_queries.load(.monotonic);
}
/// Start the server loop
pub fn run(self: *UdpServer) !void {
self.running.store(true, .release);
// Start worker threads
self.workers = self.allocator.alloc(std.Thread, self.num_workers) catch |err| {
std.log.err("UDP: failed to allocate worker threads: {}", .{err});
return error.OutOfMemory;
};
errdefer self.allocator.free(self.workers);
var started: u32 = 0;
errdefer {
self.running.store(false, .release);
self.work_queue.wakeAll();
for (self.workers[0..started]) |w| w.join();
}
for (self.workers) |*worker| {
worker.* = std.Thread.spawn(.{}, workerLoop, .{self}) catch |err| {
std.log.err("UDP: failed to start worker thread: {}", .{err});
return error.ThreadSpawnFailed;
};
started += 1;
}
std.log.info("UDP: started {} worker threads", .{self.num_workers});
var buffer: [types.EDNS_DEFAULT_SIZE]u8 = undefined;
while (self.running.load(.acquire)) {
// Use poll with timeout to allow checking running flag
var fds = [1]posix.pollfd{
.{
.fd = self.socket,
.events = posix.POLL.IN,
.revents = 0,
},
};
const poll_result = posix.poll(&fds, 100) catch |err| {
std.log.warn("UDP poll error: {}", .{err});
continue;
};
// Timeout - check running flag and continue
if (poll_result == 0) continue;
// No data available
if (fds[0].revents & posix.POLL.IN == 0) continue;
var src_addr: posix.sockaddr = undefined;
var addr_len: posix.socklen_t = @sizeOf(posix.sockaddr);
// Receive query
const recv_len = posix.recvfrom(
self.socket,
&buffer,
0,
&src_addr,
&addr_len,
) catch |err| {
std.log.warn("UDP receive error: {}", .{err});
continue;
};
if (recv_len < types.DNS_HEADER_SIZE) {
continue; // Too small to be valid DNS
}
// Create task and submit to worker pool
var task = QueryTask{
.data = undefined,
.len = recv_len,
.src_addr = src_addr,
.addr_len = addr_len,
};
@memcpy(task.data[0..recv_len], buffer[0..recv_len]);
if (!self.work_queue.tryPush(task)) {
// Backpressure: send SERVFAIL instead of silent drop
_ = self.dropped_queries.fetchAdd(1, .monotonic);
self.sendServfail(task.data[0..task.len], &task.src_addr, task.addr_len);
}
}
// Shutdown: wait for workers
self.work_queue.wakeAll();
for (self.workers) |w| w.join();
self.allocator.free(self.workers);
self.workers = &[_]std.Thread{};
}
/// Worker thread loop
fn workerLoop(self: *UdpServer) void {
while (self.running.load(.acquire)) {
if (self.work_queue.pop(&self.running)) |task| {
self.processQuery(task);
}
}
}
/// Process a single query
fn processQuery(self: *UdpServer, task: QueryTask) void {
const client_addr = std.net.Address{ .any = task.src_addr };
// Handle the query
const response = self.handler.handle(
task.data[0..task.len],
client_addr,
self.allocator,
) orelse return;
defer self.allocator.free(response);
// Determine max response size based on query EDNS support
const max_response_size = getMaxResponseSize(task.data[0..task.len]);
// Send response (truncate if needed)
if (response.len > max_response_size) {
// Set TC (truncation) bit in response header
var truncated_response: [types.EDNS_DEFAULT_SIZE]u8 = undefined;
const safe_max = @min(max_response_size, types.EDNS_DEFAULT_SIZE);
const truncated_len = @min(response.len, safe_max);
@memcpy(truncated_response[0..truncated_len], response[0..truncated_len]);
truncated_response[2] |= 0x02;
_ = posix.sendto(
self.socket,
truncated_response[0..truncated_len],
0,
&task.src_addr,
task.addr_len,
) catch |err| {
std.log.warn("UDP send error: {}", .{err});
};
} else {
_ = posix.sendto(
self.socket,
response,
0,
&task.src_addr,
task.addr_len,
) catch |err| {
std.log.warn("UDP send error: {}", .{err});
};
}
}
/// Determine maximum response size based on EDNS in query
/// Returns 512 (RFC 1035 default) if no EDNS, otherwise client's advertised size
fn getMaxResponseSize(query: []const u8) usize {
// Need at least header + minimal question
if (query.len < types.DNS_HEADER_SIZE) {
return types.DNS_UDP_SIZE;
}
// Check ARCOUNT (additional record count) - bytes 10-11
const arcount = std.mem.readInt(u16, query[10..12], .big);
if (arcount == 0) {
return types.DNS_UDP_SIZE;
}
// Quick scan for OPT record (type 41)
// OPT records have root name (0x00), type 0x0029
// This is a simplified scan - look for the pattern in additional section
var i: usize = types.DNS_HEADER_SIZE;
// Skip questions
const qdcount = std.mem.readInt(u16, query[4..6], .big);
var q: u16 = 0;
while (q < qdcount and i < query.len) : (q += 1) {
// Skip name
while (i < query.len) {
const len = query[i];
if (len == 0) {
i += 1;
break;
} else if ((len & 0xC0) == 0xC0) {
i += 2;
break;
} else {
i += 1 + len;
}
}
i += 4; // Skip QTYPE and QCLASS
}
// Skip answers
const ancount = std.mem.readInt(u16, query[6..8], .big);
var a: u16 = 0;
while (a < ancount and i < query.len) : (a += 1) {
i = skipResourceRecord(query, i);
}
// Skip authority
const nscount = std.mem.readInt(u16, query[8..10], .big);
var n: u16 = 0;
while (n < nscount and i < query.len) : (n += 1) {
i = skipResourceRecord(query, i);
}
// Look for OPT in additional
var ar: u16 = 0;
while (ar < arcount and i + 11 <= query.len) : (ar += 1) {
const name_start = i;
// Skip name
while (i < query.len) {
const len = query[i];
if (len == 0) {
i += 1;
break;
} else if ((len & 0xC0) == 0xC0) {
i += 2;
break;
} else {
i += 1 + len;
}
}
if (i + 10 > query.len) break;
const rtype = std.mem.readInt(u16, query[i..][0..2], .big);
if (rtype == 41 and query[name_start] == 0) {
// Found OPT record - CLASS field contains UDP payload size
const udp_size = std.mem.readInt(u16, query[i + 2 ..][0..2], .big);
// Return client's size, capped at our max
return @min(udp_size, types.EDNS_DEFAULT_SIZE);
}
// Skip to next record
const rdlength = std.mem.readInt(u16, query[i + 8 ..][0..2], .big);
i += 10 + rdlength;
}
return types.DNS_UDP_SIZE;
}
/// Skip a resource record and return new position
fn skipResourceRecord(data: []const u8, start: usize) usize {
var i = start;
// Skip name
while (i < data.len) {
const len = data[i];
if (len == 0) {
i += 1;
break;
} else if ((len & 0xC0) == 0xC0) {
i += 2;
break;
} else {
i += 1 + len;
}
}
// Need TYPE(2) + CLASS(2) + TTL(4) + RDLENGTH(2)
if (i + 10 > data.len) return data.len;
const rdlength = std.mem.readInt(u16, data[i + 8 ..][0..2], .big);
const new_pos = i + 10 + rdlength;
// Clamp to data.len to ensure callers don't need to handle overflow
return @min(new_pos, data.len);
}
/// Stop the server
pub fn stop(self: *UdpServer) void {
self.running.store(false, .release);
self.work_queue.wakeAll();
}
/// Close the server socket
pub fn deinit(self: *UdpServer) void {
posix.close(self.socket);
}
/// Send a SERVFAIL response for backpressure
fn sendServfail(self: *UdpServer, query: []const u8, addr: *const posix.sockaddr, addr_len: posix.socklen_t) void {
if (query.len < types.DNS_HEADER_SIZE) return;
// Build minimal SERVFAIL response (12 bytes - header only)
var response: [12]u8 = undefined;
// Copy transaction ID (bytes 0-1)
response[0] = query[0];
response[1] = query[1];
// Flags: QR=1 (response), OPCODE=copy, AA=0, TC=0, RD=copy, RA=1, Z=0, RCODE=2 (SERVFAIL)
const opcode = query[2] & 0x78; // Extract OPCODE bits
const rd = query[2] & 0x01; // Extract RD bit
response[2] = 0x80 | opcode | rd; // QR=1, copy OPCODE and RD
response[3] = 0x82; // RA=1, RCODE=2 (SERVFAIL)
// Counts: all zeros (no questions/answers in minimal response)
response[4] = 0;
response[5] = 0;
response[6] = 0;
response[7] = 0;
response[8] = 0;
response[9] = 0;
response[10] = 0;
response[11] = 0;
_ = posix.sendto(self.socket, &response, 0, addr, addr_len) catch {};
}
};
/// Create a simple echo handler for testing
/// Caller must call destroyEchoHandler when done to free allocated context
pub fn createEchoHandler(allocator: Allocator) !UdpServer.Handler {
const EchoContext = struct {
allocator: Allocator,
fn handle(ctx: *anyopaque, query: []const u8, _: std.net.Address, alloc: Allocator) ?[]const u8 {
_ = ctx;
// Parse query and create response
var pkt = packet.Packet.parse(query, alloc) catch return null;
defer pkt.deinit();
// Create simple response echoing the query
var response = packet.Packet.createDeniedResponse(&pkt, alloc) catch return null;
defer response.deinit();
var response_buffer: [types.EDNS_DEFAULT_SIZE]u8 = undefined;
const response_len = response.encode(&response_buffer) catch return null;
return alloc.dupe(u8, response_buffer[0..response_len]) catch return null;
}
};
const ctx = try allocator.create(EchoContext);
ctx.* = EchoContext{ .allocator = allocator };
return UdpServer.Handler{
.context = ctx,
.handleFn = EchoContext.handle,
};
}
/// Free the echo handler context allocated by createEchoHandler
pub fn destroyEchoHandler(handler: *UdpServer.Handler, allocator: Allocator) void {
const EchoContext = struct { allocator: Allocator };
const ctx: *EchoContext = @ptrCast(@alignCast(handler.context));
allocator.destroy(ctx);
handler.context = undefined;
}
test "UDP server creation" {
const testing = std.testing;
const allocator = testing.allocator;
var handler = try createEchoHandler(allocator);
defer destroyEchoHandler(&handler, allocator);
// Try to create server on a high port to avoid permission issues
const addr = std.net.Address.initIp4([4]u8{ 127, 0, 0, 1 }, 15353);
var server = UdpServer.init(addr, &handler, allocator) catch |err| {
// Skip test if we can't bind (e.g., in CI)
std.log.warn("Could not create UDP server: {}", .{err});
return;
};
defer server.deinit();
try testing.expect(server.socket != 0);
}