From 3a16647522ba0b7ae9d3d5df9d728595d863fa29 Mon Sep 17 00:00:00 2001
From: Nazar Mokrynskyi <nazar@mokrynskyi.com>
Date: Mon, 20 Nov 2023 12:45:25 +0200
Subject: [PATCH] Use bootstrap nodes for probing too alongside connected peers

---
 crates/subspace-networking/src/constructor.rs |  1 +
 .../src/protocols/autonat_wrapper.rs          | 28 +++++++++++++------
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/crates/subspace-networking/src/constructor.rs b/crates/subspace-networking/src/constructor.rs
index cd380e38fd..3f0d3dcc81 100644
--- a/crates/subspace-networking/src/constructor.rs
+++ b/crates/subspace-networking/src/constructor.rs
@@ -538,6 +538,7 @@ where
                 ..Default::default()
             },
             local_peer_id,
+            servers: bootstrap_addresses.clone(),
         },
     });
 
diff --git a/crates/subspace-networking/src/protocols/autonat_wrapper.rs b/crates/subspace-networking/src/protocols/autonat_wrapper.rs
index 7c19764d82..0af6f1e822 100644
--- a/crates/subspace-networking/src/protocols/autonat_wrapper.rs
+++ b/crates/subspace-networking/src/protocols/autonat_wrapper.rs
@@ -14,27 +14,39 @@ use tracing::debug;
 pub(crate) struct Config {
     pub(crate) inner_config: AutonatConfig,
     pub(crate) local_peer_id: PeerId,
+    pub(crate) servers: Vec<Multiaddr>,
 }
 
 pub(crate) struct Behaviour {
     inner: Autonat,
-    config: Config,
+    private_ips_enabled: bool,
     listen_addresses: HashSet<Multiaddr>,
 }
 
 impl Behaviour {
     pub(crate) fn new(config: Config) -> Self {
+        let mut inner = Autonat::new(config.local_peer_id, config.inner_config.clone());
+
+        for server in config.servers {
+            let maybe_peer_id = server.iter().find_map(|protocol| {
+                if let Protocol::P2p(peer_id) = protocol {
+                    Some(peer_id)
+                } else {
+                    None
+                }
+            });
+            if let Some(peer_id) = maybe_peer_id {
+                inner.add_server(peer_id, Some(server));
+            }
+        }
+
         Self {
-            inner: Autonat::new(config.local_peer_id, config.inner_config.clone()),
-            config,
+            inner,
+            private_ips_enabled: !config.inner_config.only_global_ips,
             listen_addresses: Default::default(),
         }
     }
 
-    fn private_ips_enabled(&self) -> bool {
-        !self.config.inner_config.only_global_ips
-    }
-
     fn address_corresponds_to_listening_addresses(&self, addr: &Multiaddr) -> bool {
         let Some(candidate_protocol) = addr.iter().find_map(|protocol| match protocol {
             udp @ Protocol::Udp(_) => Some(udp),
@@ -105,7 +117,7 @@ impl NetworkBehaviour for Behaviour {
                     //TODO: handle listener address change
                     self.listen_addresses.insert(addr.addr.clone());
 
-                    if self.private_ips_enabled() || is_global_address_or_dns(addr.addr) {
+                    if self.private_ips_enabled || is_global_address_or_dns(addr.addr) {
                         self.inner.on_swarm_event(new_listen_addr_event);
                     } else {
                         debug!(addr=?addr.addr, "Skipped listening address in AutonatWrapper.");