feat: implement M2 — multiple upstreams & config
- Failover strategy: try resolvers in order until one succeeds - Env var overrides: DOH_LISTEN, DOH_UPSTREAM_STRATEGY, DOH_UPSTREAM_RESOLVERS - Graceful shutdown on SIGINT/SIGTERM - Strategy parsed from config string to enum - 14 tests passing, clippy clean, fmt clean
This commit is contained in:
@@ -86,7 +86,7 @@ See [docs/plan.md](docs/plan.md) for the full project plan and milestone checkli
|
||||
| Milestone | Status |
|
||||
|-----------|--------|
|
||||
| M1 — Core Forwarding | ✅ Done |
|
||||
| M2 — Multiple Upstreams & Config | 🔲 Not started |
|
||||
| M2 — Multiple Upstreams & Config | ✅ Done |
|
||||
| M3 — Caching | 🔲 Not started |
|
||||
| M4 — Observability | 🔲 Not started |
|
||||
| M5 — Blocking & Filtering | 🔲 Not started |
|
||||
|
||||
10
docs/plan.md
10
docs/plan.md
@@ -46,11 +46,11 @@ A lightweight DNS-over-HTTPS (DoH) forwarder that accepts DNS queries over HTTPS
|
||||
|
||||
### M2 — Multiple Upstreams & Config
|
||||
|
||||
- [ ] TOML config file support
|
||||
- [ ] Multiple upstream resolver support
|
||||
- [ ] Upstream selection strategy (round-robin / failover)
|
||||
- [ ] Environment variable overrides
|
||||
- [ ] Graceful shutdown
|
||||
- [x] TOML config file support
|
||||
- [x] Multiple upstream resolver support
|
||||
- [x] Upstream selection strategy (round-robin / failover)
|
||||
- [x] Environment variable overrides
|
||||
- [x] Graceful shutdown
|
||||
|
||||
### M3 — Caching
|
||||
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
use serde::Deserialize;
|
||||
use std::path::Path;
|
||||
|
||||
use crate::upstream::doh_client::Strategy;
|
||||
|
||||
/// Top-level application configuration.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct AppConfig {
|
||||
@@ -51,13 +53,57 @@ pub struct CacheConfig {
|
||||
}
|
||||
|
||||
impl AppConfig {
|
||||
/// Load configuration from a TOML file, with optional env var overrides.
|
||||
/// Load configuration from a TOML file, then apply env var overrides.
|
||||
pub fn load(path: &Path) -> anyhow::Result<Self> {
|
||||
let content = std::fs::read_to_string(path)
|
||||
.map_err(|e| anyhow::anyhow!("failed to read config file {:?}: {}", path, e))?;
|
||||
let config: AppConfig = toml::from_str(&content)?;
|
||||
let mut config: AppConfig = toml::from_str(&content)?;
|
||||
config.apply_env_overrides();
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// Override config values from environment variables.
|
||||
///
|
||||
/// Supported variables:
|
||||
/// - `DOH_LISTEN` — override server listen address
|
||||
/// - `DOH_UPSTREAM_STRATEGY` — override upstream selection strategy
|
||||
/// - `DOH_UPSTREAM_RESOLVERS` — comma-separated `name=url` pairs
|
||||
fn apply_env_overrides(&mut self) {
|
||||
if let Ok(listen) = std::env::var("DOH_LISTEN") {
|
||||
self.server.listen = listen;
|
||||
}
|
||||
if let Ok(strategy) = std::env::var("DOH_UPSTREAM_STRATEGY") {
|
||||
self.upstream.strategy = strategy;
|
||||
}
|
||||
if let Ok(resolvers) = std::env::var("DOH_UPSTREAM_RESOLVERS") {
|
||||
let entries: Vec<ResolverEntry> = resolvers
|
||||
.split(',')
|
||||
.filter_map(|entry| {
|
||||
let (name, url) = entry.trim().split_once('=')?;
|
||||
Some(ResolverEntry {
|
||||
name: name.trim().to_string(),
|
||||
url: url.trim().to_string(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
if !entries.is_empty() {
|
||||
self.upstream.resolvers = entries;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl UpstreamConfig {
|
||||
/// Parse the strategy string into a [`Strategy`] enum.
|
||||
pub fn strategy(&self) -> anyhow::Result<Strategy> {
|
||||
match self.strategy.as_str() {
|
||||
"round-robin" => Ok(Strategy::RoundRobin),
|
||||
"failover" => Ok(Strategy::Failover),
|
||||
other => Err(anyhow::anyhow!(
|
||||
"unknown upstream strategy: {other:?} (expected \"round-robin\" or \"failover\")"
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -114,4 +160,53 @@ max_ttl_secs = 0
|
||||
assert_eq!(config.upstream.strategy, "failover");
|
||||
assert!(!config.cache.enabled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_env_override_listen() {
|
||||
let config_str = r#"
|
||||
[server]
|
||||
listen = "127.0.0.1:8080"
|
||||
|
||||
[upstream]
|
||||
resolvers = [{ name = "cf", url = "https://cf/dns-query" }]
|
||||
strategy = "round-robin"
|
||||
|
||||
[cache]
|
||||
enabled = false
|
||||
max_entries = 0
|
||||
max_ttl_secs = 0
|
||||
"#;
|
||||
let mut config: AppConfig = toml::from_str(config_str).unwrap();
|
||||
std::env::set_var("DOH_LISTEN", "0.0.0.0:9999");
|
||||
config.apply_env_overrides();
|
||||
assert_eq!(config.server.listen, "0.0.0.0:9999");
|
||||
std::env::remove_var("DOH_LISTEN");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_env_override_resolvers() {
|
||||
let config_str = r#"
|
||||
[server]
|
||||
listen = "127.0.0.1:8080"
|
||||
|
||||
[upstream]
|
||||
resolvers = [{ name = "cf", url = "https://cf/dns-query" }]
|
||||
strategy = "round-robin"
|
||||
|
||||
[cache]
|
||||
enabled = false
|
||||
max_entries = 0
|
||||
max_ttl_secs = 0
|
||||
"#;
|
||||
let mut config: AppConfig = toml::from_str(config_str).unwrap();
|
||||
std::env::set_var(
|
||||
"DOH_UPSTREAM_RESOLVERS",
|
||||
"google=https://dns.google/dns-query,quad9=https://dns.quad9.net/dns-query",
|
||||
);
|
||||
config.apply_env_overrides();
|
||||
assert_eq!(config.upstream.resolvers.len(), 2);
|
||||
assert_eq!(config.upstream.resolvers[0].name, "google");
|
||||
assert_eq!(config.upstream.resolvers[1].name, "quad9");
|
||||
std::env::remove_var("DOH_UPSTREAM_RESOLVERS");
|
||||
}
|
||||
}
|
||||
|
||||
33
src/main.rs
33
src/main.rs
@@ -41,11 +41,13 @@ async fn main() -> anyhow::Result<()> {
|
||||
"upstream configuration"
|
||||
);
|
||||
|
||||
let strategy = config.upstream.strategy()?;
|
||||
let doh_client = DohClient::new(config.upstream.resolvers.clone());
|
||||
|
||||
let state = AppState {
|
||||
doh_client,
|
||||
rr_counter: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
|
||||
strategy,
|
||||
};
|
||||
|
||||
let app = Router::new()
|
||||
@@ -59,8 +61,12 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
info!(addr = %listener.local_addr()?, "listening");
|
||||
|
||||
axum::serve(listener, app).await.context("server error")?;
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(shutdown_signal())
|
||||
.await
|
||||
.context("server error")?;
|
||||
|
||||
info!("shutdown complete");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -69,6 +75,31 @@ async fn health() -> &'static str {
|
||||
"ok"
|
||||
}
|
||||
|
||||
/// Wait for a shutdown signal (SIGINT or SIGTERM).
|
||||
async fn shutdown_signal() {
|
||||
let ctrl_c = async {
|
||||
tokio::signal::ctrl_c()
|
||||
.await
|
||||
.expect("failed to install Ctrl+C handler");
|
||||
};
|
||||
|
||||
#[cfg(unix)]
|
||||
let terminate = async {
|
||||
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
|
||||
.expect("failed to install signal handler")
|
||||
.recv()
|
||||
.await;
|
||||
};
|
||||
|
||||
#[cfg(not(unix))]
|
||||
let terminate = std::future::pending::<()>();
|
||||
|
||||
tokio::select! {
|
||||
_ = ctrl_c => info!("received SIGINT, shutting down"),
|
||||
_ = terminate => info!("received SIGTERM, shutting down"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse the config file path from CLI args, defaulting to `config/default.toml`.
|
||||
fn parse_config_path() -> PathBuf {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
|
||||
@@ -8,15 +8,17 @@ use tracing::{error, info, warn};
|
||||
|
||||
use crate::dns::query::parse_query;
|
||||
use crate::dns::response::build_servfail;
|
||||
use crate::upstream::doh_client::DohClient;
|
||||
use crate::upstream::doh_client::{DohClient, Strategy};
|
||||
|
||||
/// Shared application state passed to route handlers.
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
/// The DoH upstream client.
|
||||
pub doh_client: DohClient,
|
||||
/// Round-robin counter for upstream selection.
|
||||
/// Counter for upstream selection (round-robin index / failover start).
|
||||
pub rr_counter: Arc<std::sync::atomic::AtomicUsize>,
|
||||
/// Upstream selection strategy.
|
||||
pub strategy: Strategy,
|
||||
}
|
||||
|
||||
/// Query parameters for the GET variant of `/dns-query`.
|
||||
@@ -105,12 +107,15 @@ async fn forward_and_respond(
|
||||
query_bytes: &[u8],
|
||||
query: &hickory_proto::op::Message,
|
||||
) -> Response {
|
||||
// Pick upstream via round-robin
|
||||
let idx = state
|
||||
.rr_counter
|
||||
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
|
||||
match state.doh_client.forward(query_bytes, idx).await {
|
||||
match state
|
||||
.doh_client
|
||||
.forward(query_bytes, state.strategy, idx)
|
||||
.await
|
||||
{
|
||||
Ok(response_bytes) => {
|
||||
info!(bytes = response_bytes.len(), "forwarding upstream response");
|
||||
(
|
||||
|
||||
@@ -17,6 +17,15 @@ pub enum UpstreamError {
|
||||
BadStatus { name: String, status: u16 },
|
||||
}
|
||||
|
||||
/// Upstream selection strategy.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Strategy {
|
||||
/// Distribute queries across resolvers in order.
|
||||
RoundRobin,
|
||||
/// Try resolvers in order until one succeeds.
|
||||
Failover,
|
||||
}
|
||||
|
||||
/// A DoH client that forwards DNS wire-format queries to upstream resolvers.
|
||||
#[derive(Clone)]
|
||||
pub struct DohClient {
|
||||
@@ -34,15 +43,63 @@ impl DohClient {
|
||||
Self { http, resolvers }
|
||||
}
|
||||
|
||||
/// Forward a DNS wire-format query to the upstream at the given index.
|
||||
/// Forward a DNS wire-format query using the given strategy.
|
||||
///
|
||||
/// Returns the raw DNS wire-format response bytes.
|
||||
/// For round-robin, `index` is used to select the resolver.
|
||||
/// For failover, resolvers are tried in order starting from `index`.
|
||||
pub async fn forward(
|
||||
&self,
|
||||
query_bytes: &[u8],
|
||||
strategy: Strategy,
|
||||
index: usize,
|
||||
) -> Result<Vec<u8>, UpstreamError> {
|
||||
match strategy {
|
||||
Strategy::RoundRobin => self.forward_to(query_bytes, index).await,
|
||||
Strategy::Failover => self.forward_failover(query_bytes, index).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward to a specific resolver by index (wraps around).
|
||||
async fn forward_to(
|
||||
&self,
|
||||
query_bytes: &[u8],
|
||||
resolver_index: usize,
|
||||
) -> Result<Vec<u8>, UpstreamError> {
|
||||
let resolver = &self.resolvers[resolver_index % self.resolvers.len()];
|
||||
self.send_to(query_bytes, resolver).await
|
||||
}
|
||||
|
||||
/// Try resolvers in order starting from `start_index`, returning the
|
||||
/// first successful response.
|
||||
async fn forward_failover(
|
||||
&self,
|
||||
query_bytes: &[u8],
|
||||
start_index: usize,
|
||||
) -> Result<Vec<u8>, UpstreamError> {
|
||||
let n = self.resolvers.len();
|
||||
let mut last_err = None;
|
||||
|
||||
for i in 0..n {
|
||||
let idx = (start_index + i) % n;
|
||||
let resolver = &self.resolvers[idx];
|
||||
match self.send_to(query_bytes, resolver).await {
|
||||
Ok(bytes) => return Ok(bytes),
|
||||
Err(e) => {
|
||||
warn!(resolver = %resolver.name, error = %e, "failover: resolver failed, trying next");
|
||||
last_err = Some(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(last_err.expect("failover: no resolvers configured"))
|
||||
}
|
||||
|
||||
/// Send a query to a single resolver.
|
||||
async fn send_to(
|
||||
&self,
|
||||
query_bytes: &[u8],
|
||||
resolver: &ResolverEntry,
|
||||
) -> Result<Vec<u8>, UpstreamError> {
|
||||
debug!(resolver = %resolver.name, url = %resolver.url, "forwarding query to upstream");
|
||||
|
||||
let response = self
|
||||
@@ -80,7 +137,6 @@ impl DohClient {
|
||||
}
|
||||
|
||||
/// Return the number of configured resolvers.
|
||||
#[allow(dead_code)] // Used by health/stats endpoints in M4
|
||||
pub fn resolver_count(&self) -> usize {
|
||||
self.resolvers.len()
|
||||
}
|
||||
|
||||
@@ -11,11 +11,10 @@ use base64::Engine;
|
||||
use hickory_proto::op::{Message, Query, ResponseCode};
|
||||
use hickory_proto::rr::{Name, RecordType};
|
||||
use std::sync::Arc;
|
||||
use tokio::net::TcpListener;
|
||||
use tower::ServiceExt;
|
||||
|
||||
use doh_forwarder::routes::{dns_query_get, dns_query_post, AppState};
|
||||
use doh_forwarder::upstream::doh_client::DohClient;
|
||||
use doh_forwarder::upstream::doh_client::{DohClient, Strategy};
|
||||
|
||||
/// Build a minimal DNS query for the given domain and record type.
|
||||
fn build_dns_query(name: &str, rtype: RecordType) -> Vec<u8> {
|
||||
@@ -36,6 +35,7 @@ fn test_state() -> AppState {
|
||||
AppState {
|
||||
doh_client: DohClient::new(resolvers),
|
||||
rr_counter: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
|
||||
strategy: Strategy::RoundRobin,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user