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/upstream/pool.zig
+413
View File
@@ -0,0 +1,413 @@
const std = @import("std");
const posix = std.posix;
const net = std.net;
const Allocator = std.mem.Allocator;
const handler = @import("../server/handler.zig");
const doh = @import("doh.zig");
const dot = @import("dot.zig");
const connection_pool = @import("connection_pool.zig");
/// Upstream protocol type
pub const Protocol = enum {
udp, // Plain UDP (IP:port)
doh, // DNS-over-HTTPS (https://...)
dot, // DNS-over-TLS (tls://...)
};
/// Persistent UDP socket for upstream queries
/// Uses mutex to prevent response mixing between concurrent queries
const UdpUpstream = struct {
socket: posix.socket_t,
mutex: std.Thread.Mutex,
timeout_ms: u32,
fn init(ip: [4]u8, port: u16, timeout_ms: u32) !UdpUpstream {
const sock = try posix.socket(posix.AF.INET, posix.SOCK.DGRAM, 0);
errdefer posix.close(sock);
const addr = net.Address.initIp4(ip, port);
// Connect the socket to the upstream - allows send/recv and ICMP errors
try posix.connect(sock, &addr.any, addr.getOsSockLen());
return UdpUpstream{
.socket = sock,
.mutex = .{},
.timeout_ms = timeout_ms,
};
}
fn deinit(self: *UdpUpstream) void {
posix.close(self.socket);
}
fn query(self: *UdpUpstream, dns_packet: []const u8, allocator: Allocator) ![]const u8 {
self.mutex.lock();
defer self.mutex.unlock();
// Set timeout for this query
const timeout = posix.timeval{
.sec = @intCast(self.timeout_ms / 1000),
.usec = @intCast((self.timeout_ms % 1000) * 1000),
};
posix.setsockopt(self.socket, posix.SOL.SOCKET, posix.SO.RCVTIMEO, std.mem.asBytes(&timeout)) catch {};
posix.setsockopt(self.socket, posix.SOL.SOCKET, posix.SO.SNDTIMEO, std.mem.asBytes(&timeout)) catch {};
// Send using connected socket (no address needed)
_ = try posix.send(self.socket, dns_packet, 0);
// Receive response
var response_buf: [4096]u8 = undefined;
const n = try posix.recv(self.socket, &response_buf, 0);
if (n < 12) return error.InvalidResponse;
const response = try allocator.alloc(u8, n);
@memcpy(response, response_buf[0..n]);
return response;
}
};
/// Upstream DNS server configuration
pub const UpstreamConfig = struct {
/// Server address/URL
/// - "8.8.8.8" or "8.8.8.8:53" for plain UDP
/// - "https://cloudflare-dns.com/dns-query" for DoH
/// - "tls://cloudflare-dns.com" for DoT
address: []const u8,
port: u16 = 53,
enabled: bool = true,
timeout_ms: u32 = 1000, // 1 second default (reduced from 2s for faster failover)
/// Detect protocol from address
pub fn getProtocol(self: UpstreamConfig) Protocol {
if (std.mem.startsWith(u8, self.address, "https://")) {
return .doh;
} else if (std.mem.startsWith(u8, self.address, "tls://")) {
return .dot;
} else {
return .udp;
}
}
};
/// Health state for an upstream server
const UpstreamHealth = struct {
failures: u32 = 0,
last_failure: i64 = 0,
const FAILURE_THRESHOLD: u32 = 3;
const COOLDOWN_SECONDS: i64 = 30;
fn isHealthy(self: *const UpstreamHealth) bool {
if (self.failures < FAILURE_THRESHOLD) return true;
// Allow retry after cooldown
const now = std.time.timestamp();
return now - self.last_failure > COOLDOWN_SECONDS;
}
fn recordFailure(self: *UpstreamHealth) void {
self.failures +|= 1; // Saturating add
self.last_failure = std.time.timestamp();
}
fn recordSuccess(self: *UpstreamHealth) void {
self.failures = 0;
}
};
/// Pool of upstream DNS servers with failover
pub const UpstreamPool = struct {
configs: []UpstreamConfig,
allocator: Allocator,
/// Connection pool for DoT - maintains persistent TLS connections
dot_pool: connection_pool.DotConnectionPool,
/// Connection pool for DoH - maintains persistent HTTP connections
doh_pool: doh.DohConnectionPool,
/// Persistent UDP sockets - one per UDP upstream config (null for non-UDP)
udp_upstreams: []?UdpUpstream,
/// Health state per upstream for fast failover
health: []UpstreamHealth,
pub fn init(configs: []const UpstreamConfig, allocator: Allocator) !UpstreamPool {
const configs_copy = try allocator.alloc(UpstreamConfig, configs.len);
@memcpy(configs_copy, configs);
const udp_upstreams = try allocator.alloc(?UdpUpstream, configs.len);
for (configs_copy, 0..) |config, i| {
if (config.getProtocol() == .udp) {
if (parseIpv4(config.address)) |ip| {
udp_upstreams[i] = UdpUpstream.init(ip, config.port, config.timeout_ms) catch |err| {
std.log.warn("Failed to create UDP socket for {s}:{d}: {} (will use per-query fallback)", .{ config.address, config.port, err });
udp_upstreams[i] = null;
continue;
};
} else {
udp_upstreams[i] = null;
}
} else {
udp_upstreams[i] = null;
}
}
const health = try allocator.alloc(UpstreamHealth, configs.len);
for (health) |*h| {
h.* = UpstreamHealth{};
}
return UpstreamPool{
.configs = configs_copy,
.allocator = allocator,
.dot_pool = connection_pool.DotConnectionPool.init(allocator),
.doh_pool = doh.DohConnectionPool.init(allocator),
.udp_upstreams = udp_upstreams,
.health = health,
};
}
pub fn deinit(self: *UpstreamPool) void {
for (self.udp_upstreams) |*maybe_udp| {
if (maybe_udp.*) |*udp| {
udp.deinit();
}
}
self.allocator.free(self.udp_upstreams);
self.allocator.free(self.health);
self.dot_pool.deinit();
self.doh_pool.deinit();
self.allocator.free(self.configs);
}
/// Convert to handler-compatible Upstream interface
pub fn toHandlerUpstream(self: *UpstreamPool) handler.Upstream {
return handler.Upstream{
.context = self,
.queryFn = queryWrapper,
};
}
fn queryWrapper(ctx: *anyopaque, dns_packet: []const u8, allocator: Allocator) ?[]const u8 {
const self: *UpstreamPool = @ptrCast(@alignCast(ctx));
return self.query(dns_packet, allocator);
}
/// Query upstream DNS servers, trying each until one succeeds.
/// Skips unhealthy upstreams (>3 consecutive failures) for 30 seconds.
pub fn query(self: *UpstreamPool, dns_packet: []const u8, allocator: Allocator) ?[]const u8 {
// First pass: try healthy upstreams only
for (self.configs, 0..) |config, i| {
if (!config.enabled) continue;
if (!self.health[i].isHealthy()) continue;
const result = self.queryUpstream(config, i, dns_packet, allocator);
if (result) |response| {
self.health[i].recordSuccess();
return response;
}
self.health[i].recordFailure();
}
// Second pass: try unhealthy upstreams as last resort
for (self.configs, 0..) |config, i| {
if (!config.enabled) continue;
if (self.health[i].isHealthy()) continue; // Already tried
const result = self.queryUpstream(config, i, dns_packet, allocator);
if (result) |response| {
self.health[i].recordSuccess();
return response;
}
self.health[i].recordFailure();
}
std.log.err("All upstream DNS servers failed", .{});
return null;
}
/// Query a single upstream based on its protocol
fn queryUpstream(self: *UpstreamPool, config: UpstreamConfig, index: usize, dns_packet: []const u8, allocator: Allocator) ?[]const u8 {
const protocol = config.getProtocol();
switch (protocol) {
.doh => {
// Use connection pool with slot-based sharding for parallelism
return self.doh_pool.query(config.address, dns_packet) catch |err| {
std.log.warn("DoH query failed for {s}: {}", .{ config.address, err });
return null;
};
},
.dot => {
// Use persistent connection pool for DoT
const parsed = parseDotUrl(config.address) orelse {
std.log.warn("DoT invalid URL: {s}", .{config.address});
return null;
};
const conn = self.dot_pool.getConnection(parsed.host, parsed.port) catch |err| {
std.log.warn("DoT pool failed for {s}: {}", .{ config.address, err });
return null;
};
conn.setTimeout(config.timeout_ms);
return conn.query(dns_packet) catch |err| {
std.log.warn("DoT query failed for {s}: {}", .{ config.address, err });
return null;
};
},
.udp => {
// Use pooled socket if available, fall back to per-query socket
if (self.udp_upstreams[index]) |*udp| {
return udp.query(dns_packet, allocator) catch |err| {
std.log.warn("UDP upstream {s}:{d} failed: {}", .{ config.address, config.port, err });
return null;
};
} else if (parseIpv4(config.address)) |ip| {
// Fallback: per-query socket (socket creation failed at init)
return queryUdp(ip, config.port, dns_packet, config.timeout_ms, allocator) catch |err| {
std.log.warn("UDP upstream {s}:{d} failed: {}", .{ config.address, config.port, err });
return null;
};
} else {
std.log.warn("Invalid IPv4 address: {s}", .{config.address});
return null;
}
},
}
}
};
/// Parse a tls:// URL into host and port
fn parseDotUrl(url: []const u8) ?struct { host: []const u8, port: u16 } {
const prefix = "tls://";
if (!std.mem.startsWith(u8, url, prefix)) {
return null;
}
const host_port = url[prefix.len..];
// Check for port
if (std.mem.lastIndexOfScalar(u8, host_port, ':')) |colon_idx| {
const host = host_port[0..colon_idx];
const port_str = host_port[colon_idx + 1 ..];
const port = std.fmt.parseInt(u16, port_str, 10) catch return null;
return .{ .host = host, .port = port };
}
return .{ .host = host_port, .port = 853 };
}
/// Query a DNS server using plain UDP
fn queryUdp(ip: [4]u8, port: u16, dns_packet: []const u8, timeout_ms: u32, allocator: Allocator) ![]const u8 {
const sock = try posix.socket(posix.AF.INET, posix.SOCK.DGRAM, 0);
defer posix.close(sock);
// Set receive timeout - this is critical to avoid indefinite blocking
const timeout = posix.timeval{
.sec = @intCast(timeout_ms / 1000),
.usec = @intCast((timeout_ms % 1000) * 1000),
};
posix.setsockopt(sock, posix.SOL.SOCKET, posix.SO.RCVTIMEO, std.mem.asBytes(&timeout)) catch |err| {
std.log.warn("Failed to set socket receive timeout (queries may hang): {}", .{err});
// Continue anyway - the query might still work, just without timeout protection
};
// Also set send timeout
posix.setsockopt(sock, posix.SOL.SOCKET, posix.SO.SNDTIMEO, std.mem.asBytes(&timeout)) catch |err| {
std.log.warn("Failed to set socket send timeout: {}", .{err});
};
const addr = net.Address.initIp4(ip, port);
_ = try posix.sendto(sock, dns_packet, 0, &addr.any, addr.getOsSockLen());
var response_buf: [4096]u8 = undefined;
const n = try posix.recvfrom(sock, &response_buf, 0, null, null);
// DNS header is 12 bytes minimum
if (n < 12) return error.InvalidResponse;
const response = try allocator.alloc(u8, n);
@memcpy(response, response_buf[0..n]);
return response;
}
/// Parse an IPv4 address string like "8.8.8.8" into bytes
fn parseIpv4(addr: []const u8) ?[4]u8 {
// Strip any port suffix
const host = if (std.mem.indexOf(u8, addr, ":")) |idx| addr[0..idx] else addr;
var result: [4]u8 = undefined;
var parts = std.mem.splitScalar(u8, host, '.');
var i: usize = 0;
while (parts.next()) |part| {
if (i >= 4) return null;
result[i] = std.fmt.parseInt(u8, part, 10) catch return null;
i += 1;
}
if (i != 4) return null;
return result;
}
test "parseIpv4" {
const testing = std.testing;
try testing.expectEqual([4]u8{ 8, 8, 8, 8 }, parseIpv4("8.8.8.8").?);
try testing.expectEqual([4]u8{ 1, 1, 1, 1 }, parseIpv4("1.1.1.1").?);
try testing.expectEqual([4]u8{ 192, 168, 1, 1 }, parseIpv4("192.168.1.1:53").?);
try testing.expect(parseIpv4("invalid") == null);
try testing.expect(parseIpv4("256.0.0.1") == null);
}
test "UpstreamConfig protocol detection" {
const testing = std.testing;
const doh_config = UpstreamConfig{ .address = "https://cloudflare-dns.com/dns-query" };
try testing.expectEqual(Protocol.doh, doh_config.getProtocol());
const dot_config = UpstreamConfig{ .address = "tls://cloudflare-dns.com" };
try testing.expectEqual(Protocol.dot, dot_config.getProtocol());
const udp_config = UpstreamConfig{ .address = "8.8.8.8" };
try testing.expectEqual(Protocol.udp, udp_config.getProtocol());
const udp_with_port = UpstreamConfig{ .address = "1.1.1.1:53" };
try testing.expectEqual(Protocol.udp, udp_with_port.getProtocol());
}
test "UpstreamPool initialization" {
const testing = std.testing;
const allocator = testing.allocator;
const configs = [_]UpstreamConfig{
.{ .address = "https://cloudflare-dns.com/dns-query" },
.{ .address = "tls://1.1.1.1" },
.{ .address = "8.8.8.8" },
};
var pool = try UpstreamPool.init(&configs, allocator);
defer pool.deinit();
try testing.expectEqual(@as(usize, 3), pool.configs.len);
}
test "parseDotUrl" {
const testing = std.testing;
// Valid URLs
const r1 = parseDotUrl("tls://cloudflare-dns.com").?;
try testing.expectEqualStrings("cloudflare-dns.com", r1.host);
try testing.expectEqual(@as(u16, 853), r1.port);
const r2 = parseDotUrl("tls://1.1.1.1:853").?;
try testing.expectEqualStrings("1.1.1.1", r2.host);
try testing.expectEqual(@as(u16, 853), r2.port);
const r3 = parseDotUrl("tls://dns.google:8853").?;
try testing.expectEqualStrings("dns.google", r3.host);
try testing.expectEqual(@as(u16, 8853), r3.port);
// Invalid URLs
try testing.expect(parseDotUrl("https://example.com") == null);
try testing.expect(parseDotUrl("example.com") == null);
}