fix(portmapper): enforce timeouts for upnp (#2877)

Closes #2876 --------- Co-authored-by: Philipp Krüger <[email protected]> Co-authored-by: Divma <[email protected]>
n0-computer · Nov 1, 2024 · 00a3f88 · 00a3f88
1 parent 2507e62
commit 00a3f88
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 20 deletions.
diff --git a/iroh-net/src/endpoint.rs b/iroh-net/src/endpoint.rs
@@ -1388,7 +1388,7 @@ mod tests {
     use tracing::{error_span, info, info_span, Instrument};
 
     use super::*;
-    use crate::test_utils::run_relay_server;
+    use crate::test_utils::{run_relay_server, run_relay_server_with};
 
     const TEST_ALPN: &[u8] = b"n0/iroh/test";
 
@@ -1854,4 +1854,23 @@ mod tests {
         r1.expect("ep1 timeout").unwrap();
         r2.expect("ep2 timeout").unwrap();
     }
+
+    #[tokio::test]
+    async fn test_direct_addresses_no_stun_relay() {
+        let _guard = iroh_test::logging::setup();
+        let (relay_map, _, _guard) = run_relay_server_with(None).await.unwrap();
+
+        let ep = Endpoint::builder()
+            .alpns(vec![TEST_ALPN.to_vec()])
+            .relay_mode(RelayMode::Custom(relay_map))
+            .insecure_skip_relay_cert_verify(true)
+            .bind()
+            .await
+            .unwrap();
+
+        tokio::time::timeout(Duration::from_secs(10), ep.direct_addresses().next())
+            .await
+            .unwrap()
+            .unwrap();
+    }
 }
diff --git a/iroh-net/src/test_utils.rs b/iroh-net/src/test_utils.rs
@@ -6,9 +6,12 @@ pub use dns_and_pkarr_servers::DnsPkarrServer;
 pub use dns_server::create_dns_resolver;
 use tokio::sync::oneshot;
 
-use crate::relay::{
-    server::{CertConfig, RelayConfig, Server, ServerConfig, StunConfig, TlsConfig},
-    RelayMap, RelayNode, RelayUrl,
+use crate::{
+    defaults::DEFAULT_STUN_PORT,
+    relay::{
+        server::{CertConfig, RelayConfig, Server, ServerConfig, StunConfig, TlsConfig},
+        RelayMap, RelayNode, RelayUrl,
+    },
 };
 
 /// A drop guard to clean up test infrastructure.
@@ -26,6 +29,21 @@ pub struct CleanupDropGuard(pub(crate) oneshot::Sender<()>);
 /// The returned `Url` is the url of the relay server in the returned [`RelayMap`].
 /// When dropped, the returned [`Server`] does will stop running.
 pub async fn run_relay_server() -> Result<(RelayMap, RelayUrl, Server)> {
+    run_relay_server_with(Some(StunConfig {
+        bind_addr: (Ipv4Addr::LOCALHOST, 0).into(),
+    }))
+    .await
+}
+
+/// Runs a relay server.
+///
+/// `stun` can be set to `None` to disable stun, or set to `Some` `StunConfig`,
+/// to enable stun on a specific socket.
+///
+/// The return value is similar to [`run_relay_server`].
+pub async fn run_relay_server_with(
+    stun: Option<StunConfig>,
+) -> Result<(RelayMap, RelayUrl, Server)> {
     let cert = rcgen::generate_simple_self_signed(vec!["localhost".to_string()]).unwrap();
     let rustls_cert = rustls::pki_types::CertificateDer::from(cert.serialize_der().unwrap());
     let private_key =
@@ -44,9 +62,7 @@ pub async fn run_relay_server() -> Result<(RelayMap, RelayUrl, Server)> {
             }),
             limits: Default::default(),
         }),
-        stun: Some(StunConfig {
-            bind_addr: (Ipv4Addr::LOCALHOST, 0).into(),
-        }),
+        stun,
         #[cfg(feature = "metrics")]
         metrics_addr: None,
     };
@@ -57,7 +73,7 @@ pub async fn run_relay_server() -> Result<(RelayMap, RelayUrl, Server)> {
     let m = RelayMap::from_nodes([RelayNode {
         url: url.clone(),
         stun_only: false,
-        stun_port: server.stun_addr().unwrap().port(),
+        stun_port: server.stun_addr().map_or(DEFAULT_STUN_PORT, |s| s.port()),
     }])
     .unwrap();
     Ok((m, url, server))

diff --git a/net-tools/portmapper/src/upnp.rs b/net-tools/portmapper/src/upnp.rs
@@ -49,11 +49,15 @@ impl Mapping {
         let gateway = if let Some(known_gateway) = gateway {
             known_gateway
         } else {
-            aigd::tokio::search_gateway(igd_next::SearchOptions {
-                timeout: Some(SEARCH_TIMEOUT),
-                ..Default::default()
-            })
-            .await?
+            // Wrap in manual timeout, because igd_next doesn't respect the set timeout
+            tokio::time::timeout(
+                SEARCH_TIMEOUT,
+                aigd::tokio::search_gateway(igd_next::SearchOptions {
+                    timeout: Some(SEARCH_TIMEOUT),
+                    ..Default::default()
+                }),
+            )
+            .await??
         };
 
         let std::net::IpAddr::V4(external_ip) = gateway.get_external_ip().await? else {
@@ -126,14 +130,25 @@ impl Mapping {
 /// Searches for UPnP gateways.
 pub async fn probe_available() -> Option<Gateway> {
     inc!(Metrics, upnp_probes);
-    match aigd::tokio::search_gateway(igd_next::SearchOptions {
-        timeout: Some(SEARCH_TIMEOUT),
-        ..Default::default()
-    })
-    .await
-    {
-        Ok(gateway) => Some(gateway),
+
+    // Wrap in manual timeout, because igd_next doesn't respect the set timeout
+    let res = tokio::time::timeout(
+        SEARCH_TIMEOUT,
+        aigd::tokio::search_gateway(igd_next::SearchOptions {
+            timeout: Some(SEARCH_TIMEOUT),
+            ..Default::default()
+        }),
+    )
+    .await;
+
+    match res {
+        Ok(Ok(gateway)) => Some(gateway),
         Err(e) => {
+            inc!(Metrics, upnp_probes_failed);
+            debug!("upnp probe timed out: {e}");
+            None
+        }
+        Ok(Err(e)) => {
             inc!(Metrics, upnp_probes_failed);
             debug!("upnp probe failed: {e}");
             None