diff --git a/DEVELOP.md b/DEVELOP.md index 9aa3e8d..86f1860 100644 --- a/DEVELOP.md +++ b/DEVELOP.md @@ -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 | diff --git a/docs/plan.md b/docs/plan.md index b5239c3..4ad8bad 100644 --- a/docs/plan.md +++ b/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 diff --git a/src/config.rs b/src/config.rs index e185f52..e955574 100644 --- a/src/config.rs +++ b/src/config.rs @@ -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 { 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 = 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 { + 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"); + } } diff --git a/src/main.rs b/src/main.rs index ca0d131..e7d65a6 100644 --- a/src/main.rs +++ b/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 = std::env::args().collect(); diff --git a/src/routes.rs b/src/routes.rs index 23e00f5..687b982 100644 --- a/src/routes.rs +++ b/src/routes.rs @@ -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, + /// 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"); ( diff --git a/src/upstream/doh_client.rs b/src/upstream/doh_client.rs index 94028b9..7fd2834 100644 --- a/src/upstream/doh_client.rs +++ b/src/upstream/doh_client.rs @@ -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, 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, 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, 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, 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() } diff --git a/tests/integration.rs b/tests/integration.rs index 3b8227c..83bc425 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -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 { @@ -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, } }