feat: override to staging relays

iroh-cli/src/config.rs

@@ -10,10 +10,7 @@ use std::{
 };
 
 use anyhow::{anyhow, bail, Context, Result};
-use iroh::net::{
-    defaults,
-    relay::{RelayMap, RelayNode},
-};
+use iroh::net::relay::{RelayMap, RelayNode};
 use iroh::node::GcPolicy;
 use iroh::{
     client::Iroh,
@@ -68,24 +65,13 @@ pub(crate) struct NodeConfig {
 
 impl Default for NodeConfig {
     fn default() -> Self {
-        #[cfg(not(test))]
-        let relay_nodes = {
-            use defaults::prod::{
-                default_ap_relay_node, default_eu_relay_node, default_na_relay_node,
-            };
-            [
-                default_na_relay_node(),
-                default_eu_relay_node(),
-                default_ap_relay_node(),
-            ]
-        };
-        #[cfg(test)]
-        let relay_nodes = {
-            use defaults::staging::{default_eu_relay_node, default_na_relay_node};
-            [default_na_relay_node(), default_eu_relay_node()]
-        };
+        let relay_map = iroh::net::endpoint::default_relay_mode().relay_map();
+        let relay_nodes = relay_map
+            .nodes()
+            .map(|v| Arc::unwrap_or_clone(v.clone()))
+            .collect();
         Self {
-            relay_nodes: relay_nodes.into(),
+            relay_nodes,
             gc_policy: GcPolicyConfig::default(),
             metrics_addr: Some(([127, 0, 0, 1], 9090).into()),
             file_logs: Default::default(),

iroh-net/src/endpoint.rs

@@ -19,20 +19,19 @@ use std::sync::Arc;
 use std::task::Poll;
 use std::time::Duration;
 
-use anyhow::{anyhow, bail, ensure, Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 use derive_more::Debug;
 use futures_lite::{Stream, StreamExt};
 use tokio_util::sync::{CancellationToken, WaitForCancellationFuture};
 use tracing::{debug, info_span, trace, warn};
 use url::Url;
 
 use crate::{
-    defaults,
     discovery::{Discovery, DiscoveryTask},
     dns::{default_resolver, DnsResolver},
     key::{PublicKey, SecretKey},
     magicsock::{self, Handle},
-    relay::{RelayMap, RelayMode, RelayUrl},
+    relay::{RelayMode, RelayUrl},
     tls, NodeId,
 };
 
@@ -59,6 +58,10 @@ pub use iroh_base::node_addr::{AddrInfo, NodeAddr};
 /// is still no connection the configured [`Discovery`] will be used however.
 const DISCOVERY_WAIT_PERIOD: Duration = Duration::from_millis(500);
 
+/// Environment variable to force the use of staging relays.
+#[cfg(not(any(test, feature = "test-utils")))]
+const ENV_FORCE_STAGING_RELAYS: &str = "IROH_FORCE_STAGING_RELAYS";
+
 /// Builder for [`Endpoint`].
 ///
 /// By default the endpoint will generate a new random [`SecretKey`], which will result in a
@@ -84,15 +87,9 @@ pub struct Builder {
 
 impl Default for Builder {
     fn default() -> Self {
-        // Use staging in testing
-        #[cfg(not(any(test, feature = "test-utils")))]
-        let relay_mode = RelayMode::Default;
-        #[cfg(any(test, feature = "test-utils"))]
-        let relay_mode = RelayMode::Staging;
-
         Self {
             secret_key: Default::default(),
-            relay_mode,
+            relay_mode: default_relay_mode(),
             alpn_protocols: Default::default(),
             transport_config: Default::default(),
             concurrent_connections: Default::default(),
@@ -121,15 +118,7 @@ impl Builder {
     ///
     /// NOTE: This will be improved soon to add support for binding on specific addresses.
     pub async fn bind(self, bind_port: u16) -> Result<Endpoint> {
-        let relay_map = match self.relay_mode {
-            RelayMode::Disabled => RelayMap::empty(),
-            RelayMode::Default => defaults::prod::default_relay_map(),
-            RelayMode::Staging => defaults::staging::default_relay_map(),
-            RelayMode::Custom(relay_map) => {
-                ensure!(!relay_map.is_empty(), "Empty custom relay server map",);
-                relay_map
-            }
-        };
+        let relay_map = self.relay_mode.relay_map();
         let secret_key = self.secret_key.unwrap_or_else(SecretKey::generate);
         let static_config = StaticConfig {
             transport_config: Arc::new(self.transport_config.unwrap_or_default()),
@@ -1083,6 +1072,26 @@ fn proxy_url_from_env() -> Option<Url> {
     None
 }
 
+/// Returns the default relay mode.
+///
+/// If the `IROH_FORCE_STAGING_RELAYS` environment variable is set to `1`, it will return `RelayMode::Staging`.
+/// Otherwise, it will return `RelayMode::Default`.
+pub fn default_relay_mode() -> RelayMode {
+    // Use staging in testing
+    #[cfg(not(any(test, feature = "test-utils")))]
+    let force_staging_relays = match std::env::var(ENV_FORCE_STAGING_RELAYS) {
+        Ok(value) => value == "1",
+        Err(_) => false,
+    };
+    #[cfg(any(test, feature = "test-utils"))]
+    let force_staging_relays = true;
+
+    match force_staging_relays {
+        true => RelayMode::Staging,
+        false => RelayMode::Default,
+    }
+}
+
 /// Check if we are being executed in a CGI context.
 ///
 /// If so, a malicious client can send the `Proxy:` header, and it will

iroh-net/src/relay/map.rs

@@ -22,6 +22,18 @@ pub enum RelayMode {
     Custom(RelayMap),
 }
 
+impl RelayMode {
+    /// Returns the relay map for this mode.
+    pub fn relay_map(&self) -> RelayMap {
+        match self {
+            RelayMode::Disabled => RelayMap::empty(),
+            RelayMode::Default => crate::defaults::prod::default_relay_map(),
+            RelayMode::Staging => crate::defaults::staging::default_relay_map(),
+            RelayMode::Custom(relay_map) => relay_map.clone(),
+        }
+    }
+}
+
 /// Configuration of all the relay servers that can be used.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct RelayMap {