Re-enable storing and restoring DHT IP addresses

equalitie · Jan 24, 2024 · c7b4e39 · c7b4e39
1 parent bd47da4
commit c7b4e39
Showing 1 changed file with 29 additions and 33 deletions.
diff --git a/lib/src/network/dht_discovery.rs b/lib/src/network/dht_discovery.rs
@@ -257,20 +257,17 @@ impl MonitoredDht {
         contacts_store: Option<Arc<dyn DhtContactsStoreTrait>>,
     ) -> Self {
         // TODO: load the DHT state from a previous save if it exists.
-        let builder = MainlineDht::builder()
+        let mut builder = MainlineDht::builder()
             .add_routers(DHT_ROUTERS.iter().copied())
             .set_read_only(false);
 
-        // TODO: The reuse of initial contacts is incorrectly implemented, once the issue
-        // https://github.com/equalitie/btdht/issues/8
-        // is fixed, this can be uncommented again.
-        //if let Some(contacts_store) = &contacts_store {
-        //    let initial_contacts = Self::load_initial_contacts(is_v4, &**contacts_store).await;
+        if let Some(contacts_store) = &contacts_store {
+            let initial_contacts = Self::load_initial_contacts(is_v4, &**contacts_store).await;
 
-        //    for contact in initial_contacts {
-        //        builder = builder.add_node(contact);
-        //    }
-        //}
+            for contact in initial_contacts {
+                builder = builder.add_node(contact);
+            }
+        }
 
         let dht = builder
             .start(Socket(socket))
@@ -401,29 +398,28 @@ impl MonitoredDht {
         }
     }
 
-    // TODO: See comment where this function is used (it's also commented out).
-    //async fn load_initial_contacts(
-    //    is_v4: bool,
-    //    contacts_store: &(impl DhtContactsStoreTrait + ?Sized),
-    //) -> HashSet<SocketAddr> {
-    //    if is_v4 {
-    //        match contacts_store.load_v4().await {
-    //            Ok(contacts) => contacts.iter().cloned().map(SocketAddr::V4).collect(),
-    //            Err(error) => {
-    //                tracing::error!("Failed to load DHT IPv4 contacts {:?}", error);
-    //                Default::default()
-    //            }
-    //        }
-    //    } else {
-    //        match contacts_store.load_v6().await {
-    //            Ok(contacts) => contacts.iter().cloned().map(SocketAddr::V6).collect(),
-    //            Err(error) => {
-    //                tracing::error!("Failed to load DHT IPv4 contacts {:?}", error);
-    //                Default::default()
-    //            }
-    //        }
-    //    }
-    //}
+    async fn load_initial_contacts(
+        is_v4: bool,
+        contacts_store: &(impl DhtContactsStoreTrait + ?Sized),
+    ) -> HashSet<SocketAddr> {
+        if is_v4 {
+            match contacts_store.load_v4().await {
+                Ok(contacts) => contacts.iter().cloned().map(SocketAddr::V4).collect(),
+                Err(error) => {
+                    tracing::error!("Failed to load DHT IPv4 contacts {:?}", error);
+                    Default::default()
+                }
+            }
+        } else {
+            match contacts_store.load_v6().await {
+                Ok(contacts) => contacts.iter().cloned().map(SocketAddr::V6).collect(),
+                Err(error) => {
+                    tracing::error!("Failed to load DHT IPv4 contacts {:?}", error);
+                    Default::default()
+                }
+            }
+        }
+    }
 }
 
 type Lookups = HashMap<InfoHash, Lookup>;