feat: DNS resolver trait

Author: Frando

iroh-dns-server/examples/resolve.rs

@@ -54,7 +54,7 @@ async fn main() -> anyhow::Result<()> {
         DnsResolver::with_nameserver(addr)
     } else {
         match args.env {
-            Env::Staging | Env::Prod => DnsResolver::new(),
+            Env::Staging | Env::Prod => DnsResolver::default(),
             Env::Dev => {
                 DnsResolver::with_nameserver(DEV_DNS_SERVER.parse().expect("valid address"))
             }

iroh-dns-server/src/lib.rs

@@ -116,25 +116,25 @@ mod tests {
         // resolve root record
         let name = Name::from_utf8(format!("{pubkey}."))?;
         let res = resolver.lookup_txt(name, DNS_TIMEOUT).await?;
-        let records = res.into_iter().map(|t| t.to_string()).collect::<Vec<_>>();
+        let records = res.map(|t| t.to_string()).collect::<Vec<_>>();
         assert_eq!(records, vec!["hi0".to_string()]);
 
         // resolve level one record
         let name = Name::from_utf8(format!("_hello.{pubkey}."))?;
         let res = resolver.lookup_txt(name, DNS_TIMEOUT).await?;
-        let records = res.into_iter().map(|t| t.to_string()).collect::<Vec<_>>();
+        let records = res.map(|t| t.to_string()).collect::<Vec<_>>();
         assert_eq!(records, vec!["hi1".to_string()]);
 
         // resolve level two record
         let name = Name::from_utf8(format!("_hello.world.{pubkey}."))?;
         let res = resolver.lookup_txt(name, DNS_TIMEOUT).await?;
-        let records = res.into_iter().map(|t| t.to_string()).collect::<Vec<_>>();
+        let records = res.map(|t| t.to_string()).collect::<Vec<_>>();
         assert_eq!(records, vec!["hi2".to_string()]);
 
         // resolve multiple records for same name
         let name = Name::from_utf8(format!("multiple.{pubkey}."))?;
         let res = resolver.lookup_txt(name, DNS_TIMEOUT).await?;
-        let records = res.into_iter().map(|t| t.to_string()).collect::<Vec<_>>();
+        let records = res.map(|t| t.to_string()).collect::<Vec<_>>();
         assert_eq!(records, vec!["hi3".to_string(), "hi4".to_string()]);
 
         // resolve A record

iroh-relay/src/dns.rs

@@ -3,36 +3,41 @@
 use std::{
     fmt::{self, Write},
     future::Future,
-    net::{IpAddr, Ipv6Addr, SocketAddr},
+    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
+    sync::Arc,
 };
 
 use anyhow::{bail, Context, Result};
 use hickory_resolver::{name_server::TokioConnectionProvider, TokioResolver};
 use iroh_base::NodeId;
 use n0_future::{
+    boxed::BoxFuture,
     time::{self, Duration},
     StreamExt,
 };
 use url::Url;
 
-use crate::node_info::NodeInfo;
+use crate::{
+    defaults::timeouts::DNS_TIMEOUT,
+    node_info::{self, NodeInfo},
+};
 
 /// The n0 testing DNS node origin, for production.
 pub const N0_DNS_NODE_ORIGIN_PROD: &str = "dns.iroh.link";
 /// The n0 testing DNS node origin, for testing.
 pub const N0_DNS_NODE_ORIGIN_STAGING: &str = "staging-dns.iroh.link";
 
-/// The DNS resolver used throughout `iroh`.
+/// The default DNS resolver used throughout `iroh`.
 #[derive(Debug, Clone)]
-pub struct DnsResolver(TokioResolver);
+struct HickoryResolver(TokioResolver);
 
-impl DnsResolver {
+impl HickoryResolver {
     /// Create a new DNS resolver with sensible cross-platform defaults.
     ///
     /// We first try to read the system's resolver from `/etc/resolv.conf`.
     /// This does not work at least on some Androids, therefore we fallback
     /// to the default `ResolverConfig` which uses eg. to google's `8.8.8.8` or `8.8.4.4`.
-    pub fn new() -> Self {
+    fn new() -> Self {
         let (system_config, mut options) =
             hickory_resolver::system_conf::read_system_conf().unwrap_or_default();
 
@@ -57,11 +62,11 @@ impl DnsResolver {
         let mut builder =
             TokioResolver::builder_with_config(config, TokioConnectionProvider::default());
         *builder.options_mut() = options;
-        DnsResolver(builder.build())
+        HickoryResolver(builder.build())
     }
 
     /// Create a new DNS resolver configured with a single UDP DNS nameserver.
-    pub fn with_nameserver(nameserver: SocketAddr) -> Self {
+    fn with_nameserver(nameserver: SocketAddr) -> Self {
         let mut config = hickory_resolver::config::ResolverConfig::new();
         let nameserver_config = hickory_resolver::config::NameServerConfig::new(
             nameserver,
@@ -71,19 +76,115 @@ impl DnsResolver {
 
         let builder =
             TokioResolver::builder_with_config(config, TokioConnectionProvider::default());
-        DnsResolver(builder.build())
+        HickoryResolver(builder.build())
+    }
+}
+
+impl Default for HickoryResolver {
+    fn default() -> Self {
+        Self::new()
     }
+}
 
-    /// Removes all entries from the cache.
-    pub fn clear_cache(&self) {
+/// Trait for DNS resolvers used in iroh.
+pub trait Resolver: std::fmt::Debug + Send + Sync + 'static {
+    /// Lookup an IPv4 address.
+    fn lookup_ipv4(&self, host: String) -> BoxFuture<Result<BoxIter<Ipv4Addr>>>;
+    /// Lookup an IPv6 address.
+    fn lookup_ipv6(&self, host: String) -> BoxFuture<Result<BoxIter<Ipv6Addr>>>;
+    /// Lookup TXT records.
+    fn lookup_txt(&self, host: String) -> BoxFuture<Result<BoxIter<TxtRecord>>>;
+    /// Clear the internal cache.
+    fn clear_cache(&self);
+}
+
+/// Boxed iterator alias.
+pub type BoxIter<T> = Box<dyn Iterator<Item = T> + Send + 'static>;
+
+impl Resolver for HickoryResolver {
+    fn lookup_ipv4(&self, host: String) -> BoxFuture<Result<BoxIter<Ipv4Addr>>> {
+        let this = self.0.clone();
+        Box::pin(async move {
+            let addrs = this.ipv4_lookup(host).await?;
+            let iter: BoxIter<Ipv4Addr> = Box::new(addrs.into_iter().map(Ipv4Addr::from));
+            Ok(iter)
+        })
+    }
+
+    fn lookup_ipv6(&self, host: String) -> BoxFuture<Result<BoxIter<Ipv6Addr>>> {
+        let this = self.0.clone();
+        Box::pin(async move {
+            let addrs = this.ipv6_lookup(host).await?;
+            let iter: BoxIter<Ipv6Addr> = Box::new(addrs.into_iter().map(Ipv6Addr::from));
+            Ok(iter)
+        })
+    }
+
+    fn lookup_txt(&self, host: String) -> BoxFuture<Result<BoxIter<TxtRecord>>> {
+        let this = self.0.clone();
+        Box::pin(async move {
+            let lookup = this.txt_lookup(host).await?;
+            let iter: BoxIter<TxtRecord> = Box::new(
+                lookup
+                    .into_iter()
+                    .map(|txt| TxtRecord::from_iter(txt.iter().cloned())),
+            );
+            Ok(iter)
+        })
+    }
+
+    fn clear_cache(&self) {
         self.0.clear_cache();
     }
+}
+
+impl From<TokioResolver> for HickoryResolver {
+    fn from(resolver: TokioResolver) -> Self {
+        HickoryResolver(resolver)
+    }
+}
+
+///
+#[derive(Debug, Clone)]
+pub struct DnsResolver(Arc<dyn Resolver>);
+
+impl std::ops::Deref for DnsResolver {
+    type Target = dyn Resolver;
+    fn deref(&self) -> &Self::Target {
+        &*self.0
+    }
+}
+
+impl Default for DnsResolver {
+    fn default() -> Self {
+        Self::new_with_system_defaults()
+    }
+}
+
+impl DnsResolver {
+    ///
+    pub fn new(resolver: impl Resolver) -> Self {
+        Self(Arc::new(resolver))
+    }
+
+    ///
+    pub fn new_with_system_defaults() -> Self {
+        Self::new(HickoryResolver::new())
+    }
+
+    ///
+    pub fn with_nameserver(nameserver: SocketAddr) -> Self {
+        Self::new(HickoryResolver::with_nameserver(nameserver))
+    }
 
     /// Lookup a TXT record.
-    pub async fn lookup_txt(&self, host: impl ToString, timeout: Duration) -> Result<TxtLookup> {
-        let host = host.to_string();
-        let res = time::timeout(timeout, self.0.txt_lookup(host)).await??;
-        Ok(TxtLookup(res))
+    pub async fn lookup_txt(
+        &self,
+        host: impl ToString,
+        timeout: Duration,
+    ) -> Result<impl Iterator<Item = TxtRecord>> {
+        let res = time::timeout(timeout, self.0.lookup_txt(host.to_string())).await??;
+        Ok(res)
     }
 
     /// Perform an ipv4 lookup with a timeout.
@@ -92,9 +193,8 @@ impl DnsResolver {
         host: impl ToString,
         timeout: Duration,
     ) -> Result<impl Iterator<Item = IpAddr>> {
-        let host = host.to_string();
-        let addrs = time::timeout(timeout, self.0.ipv4_lookup(host)).await??;
-        Ok(addrs.into_iter().map(|ip| IpAddr::V4(ip.0)))
+        let addrs = time::timeout(timeout, self.0.lookup_ipv4(host.to_string())).await??;
+        Ok(addrs.map(|ip| IpAddr::V4(ip)))
     }
 
     /// Perform an ipv6 lookup with a timeout.
@@ -103,9 +203,8 @@ impl DnsResolver {
         host: impl ToString,
         timeout: Duration,
     ) -> Result<impl Iterator<Item = IpAddr>> {
-        let host = host.to_string();
-        let addrs = time::timeout(timeout, self.0.ipv6_lookup(host)).await??;
-        Ok(addrs.into_iter().map(|ip| IpAddr::V6(ip.0)))
+        let addrs = time::timeout(timeout, self.0.lookup_ipv6(host.to_string())).await??;
+        Ok(addrs.map(|ip| IpAddr::V6(ip)))
     }
 
     /// Resolve IPv4 and IPv6 in parallel with a timeout.
@@ -221,21 +320,19 @@ impl DnsResolver {
     /// To lookup nodes that published their node info to the DNS servers run by n0,
     /// pass [`N0_DNS_NODE_ORIGIN_PROD`] as `origin`.
     pub async fn lookup_node_by_id(&self, node_id: &NodeId, origin: &str) -> Result<NodeInfo> {
-        let attrs = crate::node_info::TxtAttrs::<crate::node_info::IrohAttr>::lookup_by_id(
-            self, node_id, origin,
-        )
-        .await?;
-        let info = attrs.into();
-        Ok(info)
+        let name = node_info::node_domain(node_id, origin);
+        let name = node_info::ensure_iroh_txt_label(name);
+        let lookup = self.lookup_txt(name.clone(), DNS_TIMEOUT).await?;
+        let attrs = node_info::TxtAttrs::from_txt_lookup(name, lookup)?;
+        Ok(attrs.into())
     }
 
     /// Looks up node info by DNS name.
     pub async fn lookup_node_by_domain_name(&self, name: &str) -> Result<NodeInfo> {
-        let attrs =
-            crate::node_info::TxtAttrs::<crate::node_info::IrohAttr>::lookup_by_name(self, name)
-                .await?;
-        let info = attrs.into();
-        Ok(info)
+        let name = node_info::ensure_iroh_txt_label(name.to_string());
+        let lookup = self.lookup_txt(name.clone(), DNS_TIMEOUT).await?;
+        let attrs = node_info::TxtAttrs::from_txt_lookup(name, lookup)?;
+        Ok(attrs.into())
     }
 
     /// Looks up node info by DNS name in a staggered fashion.
@@ -270,52 +367,35 @@ impl DnsResolver {
     }
 }
 
-impl Default for DnsResolver {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl From<TokioResolver> for DnsResolver {
-    fn from(resolver: TokioResolver) -> Self {
-        DnsResolver(resolver)
-    }
-}
-
-/// TXT records returned from [`DnsResolver::lookup_txt`]
+/// Record data for a TXT record.
 #[derive(Debug, Clone)]
-pub struct TxtLookup(pub(crate) hickory_resolver::lookup::TxtLookup);
+pub struct TxtRecord(Vec<Box<[u8]>>);
 
-impl From<hickory_resolver::lookup::TxtLookup> for TxtLookup {
-    fn from(value: hickory_resolver::lookup::TxtLookup) -> Self {
-        Self(value)
+impl TxtRecord {
+    ///
+    pub fn iter(&self) -> impl Iterator<Item = &[u8]> {
+        self.0.iter().map(|x| x.as_ref())
     }
-}
-
-impl IntoIterator for TxtLookup {
-    type Item = TXT;
-
-    type IntoIter = Box<dyn Iterator<Item = TXT>>;
 
-    fn into_iter(self) -> Self::IntoIter {
-        Box::new(self.0.into_iter().map(TXT))
+    ///
+    pub fn to_strings(&self) -> impl Iterator<Item = String> + '_ {
+        self.iter()
+            .map(|cstr| String::from_utf8_lossy(&cstr).to_string())
     }
 }
 
-/// Record data for a TXT record
-#[derive(Debug, Clone)]
-pub struct TXT(hickory_resolver::proto::rr::rdata::TXT);
-
-impl TXT {
-    /// Returns the raw character strings of this TXT record.
-    pub fn txt_data(&self) -> &[Box<[u8]>] {
-        self.0.txt_data()
+impl std::fmt::Display for TxtRecord {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        for s in self.iter() {
+            write!(f, "{}", &String::from_utf8_lossy(s))?
+        }
+        Ok(())
     }
 }
 
-impl fmt::Display for TXT {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}", self.0)
+impl FromIterator<Box<[u8]>> for TxtRecord {
+    fn from_iter<T: IntoIterator<Item = Box<[u8]>>>(iter: T) -> Self {
+        Self(iter.into_iter().collect())
     }
 }