Add support for multiple upstream servers

Author: jedisct1

example-encrypted-dns.toml

@@ -27,9 +27,14 @@ listen_addrs = [
 ]
 
 
-## Upstream DNS server and port
+## Upstream DNS server(s) and port(s)
+## The server tries each address in order and fails over on timeout/error.
 
-upstream_addr = "9.9.9.9:53"
+# Single upstream:
+# upstream_addrs = ["9.9.9.9:53"]
+
+# Multiple upstreams with failover:
+upstream_addrs = ["9.9.9.9:53", "149.112.112.112:53"]
 
 
 ## File name to save the state to

src/config.rs

@@ -59,11 +59,34 @@ pub struct FilteringConfig {
     pub ignore_unqualified_hostnames: Option<bool>,
 }
 
+fn deserialize_upstream_addrs<'de, D>(deserializer: D) -> Result<Vec<SocketAddr>, D::Error>
+where
+    D: serde::Deserializer<'de>,
+{
+    use serde::Deserialize;
+
+    #[derive(Deserialize)]
+    #[serde(untagged)]
+    enum SingleOrVec {
+        Single(SocketAddr),
+        Vec(Vec<SocketAddr>),
+    }
+
+    match SingleOrVec::deserialize(deserializer)? {
+        SingleOrVec::Single(addr) => Ok(vec![addr]),
+        SingleOrVec::Vec(addrs) => Ok(addrs),
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct Config {
     pub listen_addrs: Vec<ListenAddrConfig>,
     pub external_addr: Option<IpAddr>,
-    pub upstream_addr: SocketAddr,
+    #[serde(
+        alias = "upstream_addr",
+        deserialize_with = "deserialize_upstream_addrs"
+    )]
+    pub upstream_addrs: Vec<SocketAddr>,
     pub state_file: PathBuf,
     pub udp_timeout: u32,
     pub tcp_timeout: u32,
@@ -91,17 +114,20 @@ pub struct Config {
 }
 
 impl Config {
-    pub fn from_string(toml: &str) -> Result<Config, Error> {
-        let config: Config = match toml::from_str(toml) {
+    pub fn from_string(toml_str: &str) -> Result<Config, Error> {
+        let config: Config = match toml::from_str(toml_str) {
             Ok(config) => config,
             Err(e) => bail!("Parse error in the configuration file: {}", e),
         };
+        if config.upstream_addrs.is_empty() {
+            bail!("At least one upstream address must be specified");
+        }
         Ok(config)
     }
 
     pub fn from_path(path: impl AsRef<Path>) -> Result<Config, Error> {
-        let toml = fs::read_to_string(path)?;
-        Config::from_string(&toml)
+        let toml_str = fs::read_to_string(path)?;
+        Config::from_string(&toml_str)
     }
 }
 

src/globals.rs

@@ -27,7 +27,7 @@ pub struct Globals {
     pub provider_kp: SignKeyPair,
     pub listen_addrs: Vec<SocketAddr>,
     pub external_addr: Option<SocketAddr>,
-    pub upstream_addr: SocketAddr,
+    pub upstream_addrs: Vec<SocketAddr>,
     pub tls_upstream_addr: Option<SocketAddr>,
     pub udp_timeout: Duration,
     pub tcp_timeout: Duration,

src/main.rs

@@ -153,7 +153,7 @@ async fn encrypt_and_respond_to_query(
         (Some(_), None) => {
             warn!("Shared key provided without nonce");
             bail!("Internal error: shared key without nonce");
-        },
+        }
         (Some(shared_key), Some(nonce)) => dnscrypt::encrypt(
             maybe_truncate_response(&client_ctx, packet, response, original_packet_size)?,
             shared_key,
@@ -309,11 +309,11 @@ async fn tcp_acceptor(globals: Arc<Globals>, tcp_listener: TcpListener) -> Resul
                 // Rate limit repeated errors to avoid spinning
                 let is_resource_error = matches!(
                     e.kind(),
-                    std::io::ErrorKind::ConnectionRefused |
-                    std::io::ErrorKind::ConnectionReset |
-                    std::io::ErrorKind::ConnectionAborted |
-                    std::io::ErrorKind::AddrInUse |
-                    std::io::ErrorKind::AddrNotAvailable
+                    std::io::ErrorKind::ConnectionRefused
+                        | std::io::ErrorKind::ConnectionReset
+                        | std::io::ErrorKind::ConnectionAborted
+                        | std::io::ErrorKind::AddrInUse
+                        | std::io::ErrorKind::AddrNotAvailable
                 );
 
                 if is_resource_error {
@@ -845,7 +845,7 @@ fn main() -> Result<(), Error> {
         provider_name,
         provider_kp,
         listen_addrs,
-        upstream_addr: config.upstream_addr,
+        upstream_addrs: config.upstream_addrs,
         tls_upstream_addr: config.tls.upstream_addr,
         external_addr,
         tcp_timeout: Duration::from_secs(u64::from(config.tcp_timeout)),

src/resolver.rs

@@ -1,9 +1,10 @@
 use std::cmp;
 use std::hash::Hasher;
 use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
+use std::time::Duration;
 
 use byteorder::{BigEndian, ByteOrder};
-use rand::{random, Rng, rng};
+use rand::{random, rng, Rng};
 use siphasher::sip128::Hasher128;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::{TcpSocket, UdpSocket};
@@ -14,16 +15,17 @@ use crate::errors::*;
 use crate::globals::*;
 use crate::ClientCtx;
 
-pub async fn resolve_udp(
-    globals: &Globals,
-    packet: &mut Vec<u8>,
+async fn resolve_udp_single(
+    upstream_addr: SocketAddr,
+    external_addr: Option<SocketAddr>,
+    packet: &[u8],
     packet_qname: &[u8],
     tid: u16,
-    has_cached_response: bool,
+    timeout: Duration,
 ) -> Result<Vec<u8>, Error> {
-    let ext_socket = match globals.external_addr {
+    let ext_socket = match external_addr {
         Some(x) => UdpSocket::bind(x).await?,
-        None => match globals.upstream_addr {
+        None => match upstream_addr {
             SocketAddr::V4(_) => {
                 UdpSocket::bind(&SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0)))
                     .await?
@@ -39,49 +41,80 @@ pub async fn resolve_udp(
             }
         },
     };
-    ext_socket.connect(globals.upstream_addr).await?;
+    ext_socket.connect(upstream_addr).await?;
+    ext_socket.send(packet).await?;
+    let mut response = vec![0u8; DNS_MAX_PACKET_SIZE];
+    let fut = tokio::time::timeout(timeout, ext_socket.recv_from(&mut response[..]));
+    match fut.await {
+        Ok(Ok((response_len, response_addr))) => {
+            response.truncate(response_len);
+            if response_addr == upstream_addr
+                && response_len >= DNS_HEADER_SIZE
+                && dns::tid(&response) == tid
+                && packet_qname.eq_ignore_ascii_case(dns::qname(&response)?.as_slice())
+            {
+                return Ok(response);
+            }
+            bail!("Invalid response from upstream");
+        }
+        Ok(Err(e)) => bail!("UDP receive error: {}", e),
+        Err(_) => bail!("UDP timeout"),
+    }
+}
+
+pub async fn resolve_udp(
+    globals: &Globals,
+    packet: &mut Vec<u8>,
+    packet_qname: &[u8],
+    tid: u16,
+    has_cached_response: bool,
+) -> Result<Vec<u8>, Error> {
     dns::set_edns_max_payload_size(packet, DNS_MAX_PACKET_SIZE as u16)?;
-    let mut response;
     let timeout = if has_cached_response {
         globals.udp_timeout / 2
     } else {
         globals.udp_timeout
     };
-    loop {
-        ext_socket.send(packet).await?;
-        response = vec![0u8; DNS_MAX_PACKET_SIZE];
-        dns::set_rcode_servfail(&mut response);
-        let fut = tokio::time::timeout(timeout, ext_socket.recv_from(&mut response[..]));
-        match fut.await {
-            Ok(Ok((response_len, response_addr))) => {
-                response.truncate(response_len);
-                if response_addr == globals.upstream_addr
-                    && response_len >= DNS_HEADER_SIZE
-                    && dns::tid(&response) == tid
-                    && packet_qname.eq_ignore_ascii_case(dns::qname(&response)?.as_slice())
-                {
-                    break;
-                }
-            }
-            _ => {
-                if has_cached_response {
-                    trace!("Timeout, but cached response is present");
-                    break;
-                }
-                trace!("Timeout, no cached response");
+
+    let mut last_error = None;
+    for upstream_addr in &globals.upstream_addrs {
+        match resolve_udp_single(
+            *upstream_addr,
+            globals.external_addr,
+            packet,
+            packet_qname,
+            tid,
+            timeout,
+        )
+        .await
+        {
+            Ok(response) => return Ok(response),
+            Err(e) => {
+                trace!("Upstream {} failed: {}", upstream_addr, e);
+                last_error = Some(e);
             }
         }
     }
-    Ok(response)
+
+    if has_cached_response {
+        trace!("All upstreams failed, but cached response is present");
+        let mut response = vec![0u8; DNS_MAX_PACKET_SIZE];
+        dns::set_rcode_servfail(&mut response);
+        return Ok(response);
+    }
+
+    Err(last_error.unwrap_or_else(|| anyhow!("No upstream servers configured")))
 }
 
-pub async fn resolve_tcp(
-    globals: &Globals,
-    packet: &mut [u8],
+async fn resolve_tcp_single(
+    upstream_addr: SocketAddr,
+    external_addr: Option<SocketAddr>,
+    packet: &[u8],
     packet_qname: &[u8],
     tid: u16,
+    timeout: Duration,
 ) -> Result<Vec<u8>, Error> {
-    let socket = match globals.external_addr {
+    let socket = match external_addr {
         Some(x @ SocketAddr::V4(_)) => {
             let socket = TcpSocket::new_v4()?;
             socket.set_reuseaddr(true).ok();
@@ -94,26 +127,53 @@ pub async fn resolve_tcp(
             socket.bind(x)?;
             socket
         }
-        None => match globals.upstream_addr {
+        None => match upstream_addr {
             SocketAddr::V4(_) => TcpSocket::new_v4()?,
             SocketAddr::V6(_) => TcpSocket::new_v6()?,
         },
     };
-    let mut ext_socket = socket.connect(globals.upstream_addr).await?;
+
+    let connect_fut = tokio::time::timeout(timeout, socket.connect(upstream_addr));
+    let mut ext_socket = match connect_fut.await {
+        Ok(Ok(s)) => s,
+        Ok(Err(e)) => bail!("TCP connect error: {}", e),
+        Err(_) => bail!("TCP connect timeout"),
+    };
+
     ext_socket.set_nodelay(true)?;
     let mut binlen = [0u8, 0];
     BigEndian::write_u16(&mut binlen[..], packet.len() as u16);
-    ext_socket.write_all(&binlen).await?;
-    ext_socket.write_all(packet).await?;
-    ext_socket.flush().await?;
-    ext_socket.read_exact(&mut binlen).await?;
-    let response_len = BigEndian::read_u16(&binlen) as usize;
-    ensure!(
-        (DNS_HEADER_SIZE..=DNS_MAX_PACKET_SIZE).contains(&response_len),
-        "Unexpected response size"
-    );
-    let mut response = vec![0u8; response_len];
-    ext_socket.read_exact(&mut response).await?;
+
+    let write_fut = async {
+        ext_socket.write_all(&binlen).await?;
+        ext_socket.write_all(packet).await?;
+        ext_socket.flush().await?;
+        Ok::<_, std::io::Error>(())
+    };
+    tokio::time::timeout(timeout, write_fut)
+        .await
+        .map_err(|_| anyhow!("TCP write timeout"))?
+        .map_err(|e| anyhow!("TCP write error: {}", e))?;
+
+    let read_fut = async {
+        ext_socket.read_exact(&mut binlen).await?;
+        let response_len = BigEndian::read_u16(&binlen) as usize;
+        if !(DNS_HEADER_SIZE..=DNS_MAX_PACKET_SIZE).contains(&response_len) {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidData,
+                "Unexpected response size",
+            ));
+        }
+        let mut response = vec![0u8; response_len];
+        ext_socket.read_exact(&mut response).await?;
+        Ok::<_, std::io::Error>(response)
+    };
+
+    let response = tokio::time::timeout(timeout, read_fut)
+        .await
+        .map_err(|_| anyhow!("TCP read timeout"))?
+        .map_err(|e| anyhow!("TCP read error: {}", e))?;
+
     ensure!(dns::tid(&response) == tid, "Unexpected transaction ID");
     ensure!(
         packet_qname.eq_ignore_ascii_case(dns::qname(&response)?.as_slice()),
@@ -122,6 +182,34 @@ pub async fn resolve_tcp(
     Ok(response)
 }
 
+pub async fn resolve_tcp(
+    globals: &Globals,
+    packet: &mut [u8],
+    packet_qname: &[u8],
+    tid: u16,
+) -> Result<Vec<u8>, Error> {
+    let mut last_error = None;
+    for upstream_addr in &globals.upstream_addrs {
+        match resolve_tcp_single(
+            *upstream_addr,
+            globals.external_addr,
+            packet,
+            packet_qname,
+            tid,
+            globals.tcp_timeout,
+        )
+        .await
+        {
+            Ok(response) => return Ok(response),
+            Err(e) => {
+                trace!("Upstream {} TCP failed: {}", upstream_addr, e);
+                last_error = Some(e);
+            }
+        }
+    }
+    Err(last_error.unwrap_or_else(|| anyhow!("No upstream servers configured")))
+}
+
 pub async fn resolve(
     globals: &Globals,
     packet: &mut Vec<u8>,