tests: better infrastructure for discovery in tests

iroh-net/Cargo.toml

@@ -16,6 +16,7 @@ rust-version = "1.75"
 workspace = true
 
 [dependencies]
+axum = { version = "0.7.4", optional = true }
 aead = { version = "0.5.2", features = ["bytes"] }
 anyhow = { version = "1" }
 backoff = "0.4.0"
@@ -96,6 +97,7 @@ wmi = "0.13"
 windows = { version = "0.51", features = ["Win32_NetworkManagement_IpHelper", "Win32_Foundation", "Win32_NetworkManagement_Ndis", "Win32_Networking_WinSock"] }
 
 [dev-dependencies]
+axum = { version = "0.7.4" }
 clap = { version = "4", features = ["derive"] }
 criterion = "0.5.1"
 crypto_box = { version = "0.9.1", features = ["serde", "chacha20"] }
@@ -108,7 +110,6 @@ tokio = { version = "1", features = ["io-util", "sync", "rt", "net", "fs", "macr
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 iroh-test = { path = "../iroh-test" }
 serde_json = "1.0.107"
-axum = "0.7.4"
 
 [[bench]]
 name = "key"
@@ -121,7 +122,7 @@ duct = "0.13.6"
 default = ["metrics"]
 iroh-relay = ["clap", "toml", "rustls-pemfile", "regex", "serde_with", "tracing-subscriber"]
 metrics = ["iroh-metrics/metrics"]
-test-utils = []
+test-utils = ["axum"]
 
 [[bin]]
 name = "iroh-relay"

iroh-net/src/discovery.rs

@@ -570,10 +570,9 @@ mod test_dns_pkarr {
         dns::node_info::{lookup_by_id, NodeInfo},
         relay::{RelayMap, RelayMode},
         test_utils::{
-            dns_and_pkarr_servers::run_dns_and_pkarr_servers,
             dns_server::{create_dns_resolver, run_dns_server},
             pkarr_dns_state::State,
-            run_relay_server,
+            run_relay_server, DnsPkarrServer,
         },
         AddrInfo, MagicEndpoint, NodeAddr,
     };
@@ -610,8 +609,7 @@ mod test_dns_pkarr {
         let origin = "testdns.example".to_string();
         let timeout = Duration::from_secs(2);
 
-        let (nameserver, pkarr_url, state, _dns_drop_guard, _pkarr_drop_guard) =
-            run_dns_and_pkarr_servers(origin.clone()).await?;
+        let dns_pkarr_server = DnsPkarrServer::run_with_origin(origin.clone()).await?;
 
         let secret_key = SecretKey::generate();
         let node_id = secret_key.public();
@@ -621,12 +619,12 @@ mod test_dns_pkarr {
             ..Default::default()
         };
 
-        let resolver = create_dns_resolver(nameserver)?;
-        let publisher = PkarrPublisher::new(secret_key, pkarr_url);
+        let resolver = create_dns_resolver(dns_pkarr_server.nameserver)?;
+        let publisher = PkarrPublisher::new(secret_key, dns_pkarr_server.pkarr_url.clone());
         // does not block, update happens in background task
         publisher.update_addr_info(&addr_info);
         // wait until our shared state received the update from pkarr publishing
-        state.on_node(&node_id, timeout).await?;
+        dns_pkarr_server.on_node(&node_id, timeout).await?;
         let resolved = lookup_by_id(&resolver, &node_id, &origin).await?;
 
         let expected = NodeAddr {
@@ -647,15 +645,28 @@ mod test_dns_pkarr {
         let origin = "testdns.example".to_string();
         let timeout = Duration::from_secs(2);
 
-        let (nameserver, pkarr_url, state, _dns_drop_guard, _pkarr_drop_guard) =
-            run_dns_and_pkarr_servers(&origin).await?;
+        let dns_pkarr_server = DnsPkarrServer::run_with_origin(origin.clone()).await?;
         let (relay_map, _relay_url, _relay_guard) = run_relay_server().await?;
 
-        let ep1 = ep_with_discovery(relay_map.clone(), nameserver, &origin, &pkarr_url).await?;
-        let ep2 = ep_with_discovery(relay_map, nameserver, &origin, &pkarr_url).await?;
+        let ep1 = ep_with_discovery(
+            relay_map.clone(),
+            dns_pkarr_server.nameserver,
+            &origin,
+            &dns_pkarr_server.pkarr_url,
+        )
+        .await?;
+        let ep2 = ep_with_discovery(
+            relay_map,
+            dns_pkarr_server.nameserver,
+            &origin,
+            &dns_pkarr_server.pkarr_url,
+        )
+        .await?;
 
         // wait until our shared state received the update from pkarr publishing
-        state.on_node(&ep1.node_id(), timeout).await?;
+        dns_pkarr_server
+            .on_node(&ep1.node_id(), timeout)
+            .await?;
 
         // we connect only by node id!
         let res = ep2.connect(ep1.node_id().into(), TEST_ALPN).await;
@@ -670,15 +681,26 @@ mod test_dns_pkarr {
         let origin = "testdns.example".to_string();
         let timeout = Duration::from_secs(2);
 
-        let (nameserver, pkarr_url, state, _dns_drop_guard, _pkarr_drop_guard) =
-            run_dns_and_pkarr_servers(&origin).await?;
+        let dns_pkarr_server = DnsPkarrServer::run_with_origin(origin.clone()).await?;
         let (relay_map, _relay_url, _relay_guard) = run_relay_server().await?;
 
-        let ep1 = ep_with_discovery(relay_map.clone(), nameserver, &origin, &pkarr_url).await?;
-        let ep2 = ep_with_discovery(relay_map, nameserver, &origin, &pkarr_url).await?;
+        let ep1 = ep_with_discovery(
+            relay_map.clone(),
+            dns_pkarr_server.nameserver,
+            &origin,
+            &dns_pkarr_server.pkarr_url,
+        )
+        .await?;
+        let ep2 = ep_with_discovery(
+            relay_map,
+            dns_pkarr_server.nameserver,
+            &origin,
+            &dns_pkarr_server.pkarr_url,
+        )
+        .await?;
 
         // wait until our shared state received the update from pkarr publishing
-        state.on_node(&ep1.node_id(), timeout).await?;
+        dns_pkarr_server.on_node(&ep1.node_id(), timeout).await?;
 
         let node_addr = NodeAddr::new(ep1.node_id());
 

iroh-net/src/magic_endpoint.rs

@@ -164,7 +164,8 @@ impl MagicEndpointBuilder {
 
     /// Optionally set a custom DNS resolver to use for this endpoint.
     ///
-    /// The DNS resolver is used to resolve relay hostnames.
+    /// The DNS resolver is used to resolve relay hostnames, and node addresses if
+    /// [`DnsDiscovery`] is configured.
     ///
     /// By default, all magic endpoints share a DNS resolver, which is configured to use the
     /// host system's DNS configuration. You can pass a custom instance of [`DnsResolver`]

iroh-net/src/test_utils.rs

@@ -9,6 +9,9 @@ use crate::{
     relay::{RelayMap, RelayNode, RelayUrl},
 };
 
+pub use dns_and_pkarr_servers::DnsPkarrServer;
+pub use dns_server::create_dns_resolver;
+
 /// A drop guard to clean up test infrastructure.
 ///
 /// After dropping the test infrastructure will asynchronously shutdown and release its
@@ -65,35 +68,85 @@ pub async fn run_relay_server() -> Result<(RelayMap, RelayUrl, CleanupDropGuard)
     Ok((m, url, CleanupDropGuard(tx)))
 }
 
-#[cfg(test)]
 pub(crate) mod dns_and_pkarr_servers {
     use anyhow::Result;
-    use std::net::SocketAddr;
+    use iroh_base::key::{NodeId, SecretKey};
+    use std::{net::SocketAddr, time::Duration};
     use url::Url;
 
-    use super::CleanupDropGuard;
+    use super::{create_dns_resolver, CleanupDropGuard};
 
-    use crate::test_utils::{
-        dns_server::run_dns_server, pkarr_dns_state::State, pkarr_relay::run_pkarr_relay,
+    use crate::{
+        discovery::{dns::DnsDiscovery, pkarr_publish::PkarrPublisher, ConcurrentDiscovery},
+        dns::DnsResolver,
+        test_utils::{
+            dns_server::run_dns_server, pkarr_dns_state::State, pkarr_relay::run_pkarr_relay,
+        },
     };
 
-    pub async fn run_dns_and_pkarr_servers(
-        origin: impl ToString,
-    ) -> Result<(SocketAddr, Url, State, CleanupDropGuard, CleanupDropGuard)> {
-        let state = State::new(origin.to_string());
-        let (nameserver, dns_drop_guard) = run_dns_server(state.clone()).await?;
-        let (pkarr_url, pkarr_drop_guard) = run_pkarr_relay(state.clone()).await?;
-        Ok((
-            nameserver,
-            pkarr_url,
-            state,
-            dns_drop_guard,
-            pkarr_drop_guard,
-        ))
+    /// Handle and drop guard for test DNS and Pkarr servers.
+    ///
+    /// Once the struct is dropped the servers will shut down.
+    #[derive(Debug)]
+    pub struct DnsPkarrServer {
+        /// The node origin domain.
+        pub node_origin: String,
+        /// The shared state of the DNS and Pkarr servers.
+        state: State,
+        /// The socket address of the DNS server.
+        pub nameserver: SocketAddr,
+        /// The HTTP URL of the Pkarr server.
+        pub pkarr_url: Url,
+        _dns_drop_guard: CleanupDropGuard,
+        _pkarr_drop_guard: CleanupDropGuard,
+    }
+
+    impl DnsPkarrServer {
+        /// Run DNS and Pkarr servers on localhost.
+        pub async fn run() -> anyhow::Result<Self> {
+            Self::run_with_origin("dns.example".to_string()).await
+        }
+
+        /// Run DNS and Pkarr servers on localhost with the specified `node_origin` domain.
+        pub async fn run_with_origin(node_origin: String) -> anyhow::Result<Self> {
+            let state = State::new(node_origin.clone());
+            let (nameserver, dns_drop_guard) = run_dns_server(state.clone()).await?;
+            let (pkarr_url, pkarr_drop_guard) = run_pkarr_relay(state.clone()).await?;
+            Ok(Self {
+                node_origin,
+                nameserver,
+                pkarr_url,
+                state,
+                _dns_drop_guard: dns_drop_guard,
+                _pkarr_drop_guard: pkarr_drop_guard,
+            })
+        }
+
+        /// Create a [`ConcurrentDiscovery`] with [`DnsDiscovery`] and [`PkarrPublisher`]
+        /// configured to use the test servers.
+        pub fn discovery(&self, secret_key: SecretKey) -> Box<ConcurrentDiscovery> {
+            Box::new(ConcurrentDiscovery::from_services(vec![
+                // Enable DNS discovery by default
+                Box::new(DnsDiscovery::new(self.node_origin.clone())),
+                // Enable pkarr publishing by default
+                Box::new(PkarrPublisher::new(secret_key, self.pkarr_url.clone())),
+            ]))
+        }
+
+        /// Create a [`DnsResolver`] configured to use the test DNS server.
+        pub fn dns_resolver(&self) -> DnsResolver {
+            create_dns_resolver(self.nameserver).expect("failed to create DNS resolver")
+        }
+
+        /// Wait until a Pkarr announce for a node is published to the server.
+        ///
+        /// If `timeout` elapses an error is returned.
+        pub async fn on_node(&self, node_id: &NodeId, timeout: Duration) -> Result<()> {
+            self.state.on_node(node_id, timeout).await
+        }
     }
 }
 
-#[cfg(test)]
 pub(crate) mod dns_server {
     use std::future::Future;
     use std::net::{Ipv4Addr, SocketAddr};
@@ -199,7 +252,6 @@ pub(crate) mod dns_server {
     }
 }
 
-#[cfg(test)]
 pub(crate) mod pkarr_relay {
     use std::future::IntoFuture;
     use std::net::{Ipv4Addr, SocketAddr};
@@ -274,7 +326,6 @@ pub(crate) mod pkarr_relay {
     }
 }
 
-#[cfg(test)]
 pub(crate) mod pkarr_dns_state {
     use anyhow::{bail, Result};
     use parking_lot::{Mutex, MutexGuard};

iroh/src/node.rs

@@ -280,7 +280,7 @@ mod tests {
     use anyhow::{bail, Context};
     use bytes::Bytes;
     use iroh_bytes::provider::AddProgress;
-    use iroh_net::relay::RelayMode;
+    use iroh_net::{relay::RelayMode, test_utils::DnsPkarrServer};
 
     use crate::{
         client::BlobAddOutcome,
@@ -448,4 +448,54 @@ mod tests {
         );
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_download_via_relay_and_discovery() -> Result<()> {
+        let _guard = iroh_test::logging::setup();
+        let (relay_map, _relay_url, _guard) = iroh_net::test_utils::run_relay_server().await?;
+        let dns_pkarr_server = DnsPkarrServer::run().await?;
+
+        let secret1 = SecretKey::generate();
+        let node1 = Node::memory()
+            .secret_key(secret1.clone())
+            .bind_port(0)
+            .relay_mode(RelayMode::Custom(relay_map.clone()))
+            .insecure_skip_relay_cert_verify(true)
+            .dns_resolver(dns_pkarr_server.dns_resolver())
+            .node_discovery(dns_pkarr_server.discovery(secret1).into())
+            .spawn()
+            .await?;
+        let secret2 = SecretKey::generate();
+        let node2 = Node::memory()
+            .secret_key(secret2.clone())
+            .bind_port(0)
+            .relay_mode(RelayMode::Custom(relay_map.clone()))
+            .insecure_skip_relay_cert_verify(true)
+            .dns_resolver(dns_pkarr_server.dns_resolver())
+            .node_discovery(dns_pkarr_server.discovery(secret2).into())
+            .spawn()
+            .await?;
+        let BlobAddOutcome { hash, .. } = node1.blobs.add_bytes(b"foo".to_vec()).await?;
+
+        // create a node addr with node id only
+        let addr = NodeAddr::new(node1.node_id());
+        let req = BlobDownloadRequest {
+            hash,
+            tag: SetTagOption::Auto,
+            format: BlobFormat::Raw,
+            mode: DownloadMode::Direct,
+            nodes: vec![addr],
+        };
+        node2.blobs.download(req).await?.await?;
+        assert_eq!(
+            node2
+                .blobs
+                .read_to_bytes(hash)
+                .await
+                .context("get")?
+                .as_ref(),
+            b"foo"
+        );
+        Ok(())
+    }
 }