start dealing with dns, again

divagant-martian · divagant-martian · commit adc58e794e3a · 2024-11-15T17:28:32.000-05:00
diff --git a/iroh-netcheck/src/dns.rs b/iroh-netcheck/src/dns.rs
@@ -0,0 +1,228 @@
+use std::fmt::Write;
+use std::net::IpAddr;
+
+use anyhow::Result;
+use futures_lite::{Future, StreamExt};
+use hickory_resolver::{IntoName, TokioAsyncResolver};
+
+use crate::defaults::timeouts::DNS_TIMEOUT;
+
+/// Delay used to perform staggered dns queries.
+const DNS_STAGGERING_MS: &[u64] = &[200, 300];
+
+/// Extension trait to [`TokioAsyncResolver`].
+pub trait ResolverExt {
+    /// Perform an ipv4 lookup with a timeout.
+    fn lookup_ipv4<N: IntoName>(
+        &self,
+        host: N,
+    ) -> impl Future<Output = Result<impl Iterator<Item = IpAddr>>>;
+
+    /// Perform an ipv6 lookup with a timeout.
+    fn lookup_ipv6<N: IntoName>(
+        &self,
+        host: N,
+    ) -> impl Future<Output = Result<impl Iterator<Item = IpAddr>>>;
+
+    /// Race an ipv4 and ipv6 lookup with a timeout.
+    fn lookup_ipv4_ipv6<N: IntoName + Clone>(
+        &self,
+        host: N,
+    ) -> impl Future<Output = Result<impl Iterator<Item = IpAddr>>>;
+
+    /// Perform an ipv4 lookup with a timeout in a staggered fashion.
+    ///
+    /// From the moment this function is called, each lookup is scheduled after the delays in
+    /// `delays_ms` with the first call being done immediately. `[200ms, 300ms]` results in calls
+    /// at T+0ms, T+200ms and T+300ms. The `timeout` is applied to each call individually. The
+    /// result of the first successful call is returned, or a summary of all errors otherwise.
+    fn lookup_ipv4_staggered<N: IntoName + Clone>(
+        &self,
+        host: N,
+    ) -> impl Future<Output = Result<impl Iterator<Item = IpAddr>>>;
+
+    /// Perform an ipv6 lookup with a timeout in a staggered fashion.
+    ///
+    /// From the moment this function is called, each lookup is scheduled after the delays in
+    /// `delays_ms` with the first call being done immediately. `[200ms, 300ms]` results in calls
+    /// at T+0ms, T+200ms and T+300ms. The `timeout` is applied to each call individually. The
+    /// result of the first successful call is returned, or a summary of all errors otherwise.
+    fn lookup_ipv6_staggered<N: IntoName + Clone>(
+        &self,
+        host: N,
+    ) -> impl Future<Output = Result<impl Iterator<Item = IpAddr>>>;
+
+    /// Race an ipv4 and ipv6 lookup with a timeout in a staggered fashion.
+    ///
+    /// From the moment this function is called, each lookup is scheduled after the delays in
+    /// `delays_ms` with the first call being done immediately. `[200ms, 300ms]` results in calls
+    /// at T+0ms, T+200ms and T+300ms. The `timeout` is applied as stated in
+    /// [`Self::lookup_ipv4_ipv6`]. The result of the first successful call is returned, or a
+    /// summary of all errors otherwise.
+    // TODO(@divma): adjust docs
+    fn lookup_ipv4_ipv6_staggered<N: IntoName + Clone>(
+        &self,
+        host: N,
+    ) -> impl Future<Output = Result<impl Iterator<Item = IpAddr>>>;
+}
+
+impl ResolverExt for TokioAsyncResolver {
+    async fn lookup_ipv4<N: IntoName>(&self, host: N) -> Result<impl Iterator<Item = IpAddr>> {
+        let addrs = tokio::time::timeout(DNS_TIMEOUT, self.ipv4_lookup(host)).await??;
+        Ok(addrs.into_iter().map(|ip| IpAddr::V4(ip.0)))
+    }
+
+    async fn lookup_ipv6<N: IntoName>(&self, host: N) -> Result<impl Iterator<Item = IpAddr>> {
+        let addrs = tokio::time::timeout(DNS_TIMEOUT, self.ipv6_lookup(host)).await??;
+        Ok(addrs.into_iter().map(|ip| IpAddr::V6(ip.0)))
+    }
+
+    /// Resolve IPv4 and IPv6 in parallel.
+    ///
+    /// `LookupIpStrategy::Ipv4AndIpv6` will wait for ipv6 resolution timeout, even if it is
+    /// not usable on the stack, so we manually query both lookups concurrently and time them out
+    /// individually.
+    async fn lookup_ipv4_ipv6<N: IntoName + Clone>(
+        &self,
+        host: N,
+    ) -> Result<impl Iterator<Item = IpAddr>> {
+        let res = tokio::join!(self.lookup_ipv4(host.clone()), self.lookup_ipv6(host));
+
+        match res {
+            (Ok(ipv4), Ok(ipv6)) => Ok(LookupIter::Both(ipv4.chain(ipv6))),
+            (Ok(ipv4), Err(_)) => Ok(LookupIter::Ipv4(ipv4)),
+            (Err(_), Ok(ipv6)) => Ok(LookupIter::Ipv6(ipv6)),
+            (Err(ipv4_err), Err(ipv6_err)) => {
+                anyhow::bail!("Ipv4: {:?}, Ipv6: {:?}", ipv4_err, ipv6_err)
+            }
+        }
+    }
+
+    /// Perform an ipv4 lookup with a timeout in a staggered fashion.
+    ///
+    /// From the moment this function is called, each lookup is scheduled after the delays in
+    /// `delays_ms` with the first call being done immediately. `[200ms, 300ms]` results in calls
+    /// at T+0ms, T+200ms and T+300ms. The `timeout` is applied to each call individually. The
+    /// result of the first successful call is returned, or a summary of all errors otherwise.
+    async fn lookup_ipv4_staggered<N: IntoName + Clone>(
+        &self,
+        host: N,
+    ) -> Result<impl Iterator<Item = IpAddr>> {
+        let f = || self.lookup_ipv4(host.clone());
+        stagger_call(f, DNS_STAGGERING_MS).await
+    }
+
+    /// Perform an ipv6 lookup with a timeout in a staggered fashion.
+    ///
+    /// From the moment this function is called, each lookup is scheduled after the delays in
+    /// `delays_ms` with the first call being done immediately. `[200ms, 300ms]` results in calls
+    /// at T+0ms, T+200ms and T+300ms. The `timeout` is applied to each call individually. The
+    /// result of the first successful call is returned, or a summary of all errors otherwise.
+    async fn lookup_ipv6_staggered<N: IntoName + Clone>(
+        &self,
+        host: N,
+    ) -> Result<impl Iterator<Item = IpAddr>> {
+        let f = || self.lookup_ipv6(host.clone());
+        stagger_call(f, DNS_STAGGERING_MS).await
+    }
+
+    /// Race an ipv4 and ipv6 lookup with a timeout in a staggered fashion.
+    ///
+    /// From the moment this function is called, each lookup is scheduled after the delays in
+    /// `delays_ms` with the first call being done immediately. `[200ms, 300ms]` results in calls
+    /// at T+0ms, T+200ms and T+300ms. The `timeout` is applied as stated in
+    /// [`Self::lookup_ipv4_ipv6`]. The result of the first successful call is returned, or a
+    /// summary of all errors otherwise.
+    async fn lookup_ipv4_ipv6_staggered<N: IntoName + Clone>(
+        &self,
+        host: N,
+    ) -> Result<impl Iterator<Item = IpAddr>> {
+        let f = || self.lookup_ipv4_ipv6(host.clone());
+        stagger_call(f, DNS_STAGGERING_MS).await
+    }
+}
+
+/// Helper enum to give a unified type to the iterators of [`ResolverExt::lookup_ipv4_ipv6`].
+enum LookupIter<A, B> {
+    Ipv4(A),
+    Ipv6(B),
+    Both(std::iter::Chain<A, B>),
+}
+
+impl<A: Iterator<Item = IpAddr>, B: Iterator<Item = IpAddr>> Iterator for LookupIter<A, B> {
+    type Item = IpAddr;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self {
+            LookupIter::Ipv4(iter) => iter.next(),
+            LookupIter::Ipv6(iter) => iter.next(),
+            LookupIter::Both(iter) => iter.next(),
+        }
+    }
+}
+
+/// Staggers calls to the future F with the given delays.
+///
+/// The first call is performed immediately. The first call to succeed generates an Ok result
+/// ignoring any previous error. If all calls fail, an error summarizing all errors is returned.
+async fn stagger_call<T, F: Fn() -> Fut, Fut: Future<Output = Result<T>>>(
+    f: F,
+    delays_ms: &[u64],
+) -> Result<T> {
+    let mut calls = futures_buffered::FuturesUnorderedBounded::new(delays_ms.len() + 1);
+    // NOTE: we add the 0 delay here to have a uniform set of futures. This is more performant than
+    // using alternatives that allow futures of different types.
+    for delay in std::iter::once(&0u64).chain(delays_ms) {
+        let delay = std::time::Duration::from_millis(*delay);
+        let fut = f();
+        let staggered_fut = async move {
+            tokio::time::sleep(delay).await;
+            fut.await
+        };
+        calls.push(staggered_fut)
+    }
+
+    let mut errors = vec![];
+    while let Some(call_result) = calls.next().await {
+        match call_result {
+            Ok(t) => return Ok(t),
+            Err(e) => errors.push(e),
+        }
+    }
+
+    anyhow::bail!(
+        "no calls succeed: [ {}]",
+        errors.into_iter().fold(String::new(), |mut summary, e| {
+            write!(summary, "{e} ").expect("infallible");
+            summary
+        })
+    )
+}
+
+#[cfg(test)]
+pub(crate) mod tests {
+    use std::sync::atomic::AtomicUsize;
+    use std::time::Duration;
+
+    use super::*;
+    const TIMEOUT: Duration = Duration::from_secs(5);
+    const STAGGERING_DELAYS: &[u64] = &[200, 300];
+
+    #[tokio::test]
+    async fn stagger_basic() {
+        let _logging = iroh_test::logging::setup();
+        const CALL_RESULTS: &[Result<u8, u8>] = &[Err(2), Ok(3), Ok(5), Ok(7)];
+        static DONE_CALL: AtomicUsize = AtomicUsize::new(0);
+        let f = || {
+            let r_pos = DONE_CALL.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+            async move {
+                tracing::info!(r_pos, "call");
+                CALL_RESULTS[r_pos].map_err(|e| anyhow::anyhow!("{e}"))
+            }
+        };
+
+        let delays = [1000, 15];
+        let result = stagger_call(f, &delays).await.unwrap();
+        assert_eq!(result, 5)
+    }
+}
diff --git a/iroh-netcheck/src/lib.rs b/iroh-netcheck/src/lib.rs
@@ -29,6 +29,7 @@ use tracing::{debug, error, info_span, trace, warn, Instrument};
 pub(crate) use relay_map::{RelayMap, RelayMode, RelayNode, RelayUrl};
 
 pub mod defaults;
+mod dns;
 mod metrics;
 mod ping;
 mod relay_map;
diff --git a/iroh-netcheck/src/reportgen.rs b/iroh-netcheck/src/reportgen.rs
@@ -21,10 +21,11 @@ use std::{
     net::{IpAddr, SocketAddr},
     pin::Pin,
     sync::Arc,
+    task::{Context, Poll},
     time::Duration,
 };
 
-use anyhow::{anyhow, bail, Context, Result};
+use anyhow::{anyhow, bail, Context as _, Result};
 use iroh_metrics::inc;
 use iroh_relay::{http::RELAY_PROBE_PATH, protos::stun};
 use netwatch::{interfaces, UdpSocket};
@@ -42,13 +43,9 @@ use super::NetcheckMetrics;
 use crate::{
     self as netcheck,
     defaults::DEFAULT_STUN_PORT,
-    // dns::{DnsResolver, ResolverExt},
+    dns::ResolverExt,
     ping::{PingError, Pinger},
-    util::MaybeFuture,
-    RelayMap,
-    RelayNode,
-    RelayUrl,
-    Report,
+    RelayMap, RelayNode, RelayUrl, Report,
 };
 use hickory_resolver::TokioAsyncResolver as DnsResolver;
 
@@ -58,15 +55,11 @@ mod probes;
 use probes::{Probe, ProbePlan, ProbeProto};
 
 use crate::defaults::timeouts::{
-    CAPTIVE_PORTAL_DELAY, CAPTIVE_PORTAL_TIMEOUT, DNS_TIMEOUT, OVERALL_REPORT_TIMEOUT,
-    PROBES_TIMEOUT,
+    CAPTIVE_PORTAL_DELAY, CAPTIVE_PORTAL_TIMEOUT, OVERALL_REPORT_TIMEOUT, PROBES_TIMEOUT,
 };
 
 const ENOUGH_NODES: usize = 3;
 
-/// Delay used to perform staggered dns queries.
-const DNS_STAGGERING_MS: &[u64] = &[200, 300];
-
 /// Holds the state for a single invocation of [`netcheck::Client::get_report`].
 ///
 /// Dropping this will cancel the actor and stop the report generation.
@@ -900,7 +893,7 @@ async fn check_captive_portal(
         // Ideally we would try to resolve **both** IPv4 and IPv6 rather than purely race
         // them.  But our resolver doesn't support that yet.
         let addrs: Vec<_> = dns_resolver
-            .lookup_ipv4_ipv6_staggered(domain, DNS_TIMEOUT, DNS_STAGGERING_MS)
+            .lookup_ipv4_ipv6_staggered(domain)
             .await?
             .map(|ipaddr| SocketAddr::new(ipaddr, 0))
             .collect();
@@ -963,10 +956,7 @@ async fn get_relay_addr(
         ProbeProto::StunIpv4 | ProbeProto::IcmpV4 => match relay_node.url.host() {
             Some(url::Host::Domain(hostname)) => {
                 debug!(?proto, %hostname, "Performing DNS A lookup for relay addr");
-                match dns_resolver
-                    .lookup_ipv4_staggered(hostname, DNS_TIMEOUT, DNS_STAGGERING_MS)
-                    .await
-                {
+                match dns_resolver.lookup_ipv4_staggered(hostname).await {
                     Ok(mut addrs) => addrs
                         .next()
                         .map(|ip| ip.to_canonical())
@@ -983,10 +973,7 @@ async fn get_relay_addr(
         ProbeProto::StunIpv6 | ProbeProto::IcmpV6 => match relay_node.url.host() {
             Some(url::Host::Domain(hostname)) => {
                 debug!(?proto, %hostname, "Performing DNS AAAA lookup for relay addr");
-                match dns_resolver
-                    .lookup_ipv6_staggered(hostname, DNS_TIMEOUT, DNS_STAGGERING_MS)
-                    .await
-                {
+                match dns_resolver.lookup_ipv6_staggered(hostname).await {
                     Ok(mut addrs) => addrs
                         .next()
                         .map(|ip| ip.to_canonical())
@@ -1071,7 +1058,7 @@ async fn measure_https_latency(
         // but staggered for reliability.  Ideally this tries to resolve **both** IPv4 and
         // IPv6 though.  But our resolver does not have a function for that yet.
         let addrs: Vec<_> = dns_resolver
-            .lookup_ipv4_ipv6_staggered(domain, DNS_TIMEOUT, DNS_STAGGERING_MS)
+            .lookup_ipv4_ipv6_staggered(domain)
             .await?
             .map(|ipaddr| SocketAddr::new(ipaddr, 0))
             .collect();
@@ -1175,6 +1162,31 @@ fn update_report(report: &mut Report, probe_report: ProbeReport) {
         .or(probe_report.icmpv6);
 }
 
+/// Resolves to pending if the inner is `None`.
+#[derive(Debug)]
+pub(crate) struct MaybeFuture<T> {
+    /// Future to be polled.
+    pub inner: Option<T>,
+}
+
+// NOTE: explicit implementation to bypass derive unnecessary bounds
+impl<T> Default for MaybeFuture<T> {
+    fn default() -> Self {
+        MaybeFuture { inner: None }
+    }
+}
+
+impl<T: Future + Unpin> Future for MaybeFuture<T> {
+    type Output = T::Output;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match self.inner {
+            Some(ref mut t) => Pin::new(t).poll(cx),
+            None => Poll::Pending,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::net::{Ipv4Addr, Ipv6Addr};
