fud, dht: move `ping_locks` from fud to dht

bin/fud/fud/src/dht.rs

@@ -16,28 +16,38 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
 
 use async_trait::async_trait;
-use smol::lock::Mutex;
+use rand::{rngs::OsRng, Rng};
 use tinyjson::JsonValue;
-use tracing::debug;
+use tracing::{debug, warn};
 use url::Url;
 
 use darkfi::{
-    dht::{impl_dht_node_defaults, Dht, DhtHandler, DhtLookupReply, DhtNode},
+    dht::{
+        event::DhtEvent, impl_dht_node_defaults, Dht, DhtHandler, DhtLookupReply, DhtNode,
+        HostCacheItem,
+    },
     geode::hash_to_string,
-    net::ChannelPtr,
+    net::{
+        session::{SESSION_DIRECT, SESSION_INBOUND, SESSION_MANUAL, SESSION_OUTBOUND},
+        ChannelPtr,
+    },
     rpc::util::json_map,
+    system::timeout::timeout,
     util::time::Timestamp,
-    Result,
+    Error, Result,
 };
 use darkfi_sdk::crypto::schnorr::{SchnorrPublic, Signature};
 use darkfi_serial::{serialize_async, SerialDecodable, SerialEncodable};
 
 use crate::{
     pow::VerifiableNodeData,
-    proto::{FudAnnounce, FudNodesReply, FudNodesRequest, FudSeedersReply, FudSeedersRequest},
+    proto::{
+        FudAnnounce, FudNodesReply, FudNodesRequest, FudPingReply, FudPingRequest, FudSeedersReply,
+        FudSeedersRequest,
+    },
     util::receive_resource_msg,
     Fud,
 };
@@ -138,30 +148,104 @@ impl DhtHandler for Fud {
     }
 
     async fn ping(&self, channel: ChannelPtr) -> Result<FudNode> {
-        let lock_map = self.ping_locks.clone();
-        let mut locks = lock_map.lock().await;
+        debug!(target: "fud::DhtHandler::ping()", "Sending ping to {}", channel.display_address());
 
-        // Get or create the lock
-        let lock = if let Some(lock) = locks.get(&channel.info.id) {
-            lock.clone()
-        } else {
-            let lock = Arc::new(Mutex::new(None));
-            locks.insert(channel.info.id, lock.clone());
-            lock
-        };
-        drop(locks);
+        // Setup `FudPingReply` subscriber
+        let msg_subscriber = channel.subscribe_msg::<FudPingReply>().await.unwrap();
+
+        // Send `FudPingRequest`
+        let mut rng = OsRng;
+        let request = FudPingRequest { random: rng.gen() };
+        if channel.is_stopped() {
+            return Err(Error::ChannelStopped)
+        }
+        channel.send(&request).await?;
 
-        // Acquire the lock
-        let mut result = lock.lock().await;
+        // Wait for `FudPingReply`
+        let reply = msg_subscriber.receive_with_timeout(self.dht.settings.timeout).await;
+        msg_subscriber.unsubscribe().await;
+        let reply = reply?;
+        let node = &reply.node;
+
+        // Verify the signature
+        if !node.data.public_key.verify(&request.random.to_be_bytes(), &reply.sig) {
+            warn!(target: "fud::DhtHandler::ping()", "Received an invalid signature while pinging {}", channel.display_address());
+            self.dht
+                .event_publisher
+                .notify(DhtEvent::PingReceived {
+                    from: channel.clone(),
+                    result: Err(Error::InvalidSignature),
+                })
+                .await;
+            self.dht.cleanup_channel(channel.clone()).await;
+            channel.ban().await;
+            return Err(Error::InvalidSignature)
+        }
 
-        if let Some(res) = result.clone() {
-            return res
+        // Verify PoW
+        if let Err(e) = self.pow.write().await.verify_node(&node.data).await {
+            warn!(target: "fud::DhtHandler::ping()", "Received an invalid PoW while pinging {}: {e}", channel.display_address());
+            self.dht
+                .event_publisher
+                .notify(DhtEvent::PingReceived { from: channel.clone(), result: Err(e.clone()) })
+                .await;
+            self.dht.cleanup_channel(channel.clone()).await;
+            channel.ban().await;
+            return Err(e)
         }
+        self.dht
+            .event_publisher
+            .notify(DhtEvent::PingReceived { from: channel.clone(), result: Ok(node.id()) })
+            .await;
+
+        if channel.session_type_id() & (SESSION_OUTBOUND | SESSION_DIRECT | SESSION_MANUAL) != 0 {
+            // Wait for the other node to ping us
+            let ping_timeout = Duration::from_secs(10);
+
+            if let Err(e) = timeout(ping_timeout, self.dht.wait_fully_pinged(channel.info.id)).await
+            {
+                self.dht.cleanup_channel(channel).await;
+                return Err(e.into())
+            }
+
+            let mut host_cache = self.dht.host_cache.write().await;
+
+            // If we had another node id for this host in our cache, remove
+            // the old one from the buckets and seeders
+            if let Some(cached) = host_cache.get(channel.address()) {
+                if cached.node_id != node.id() {
+                    self.dht.remove_node(&cached.node_id).await;
+
+                    for (_, seeders) in self.dht.hash_table.write().await.iter_mut() {
+                        seeders.retain(|seeder| seeder.node.id() != cached.node_id);
+                    }
+                }
+            }
+
+            // Update host cache
+            host_cache.insert(
+                channel.address().clone(),
+                HostCacheItem { last_ping: Timestamp::current_time(), node_id: node.id() },
+            );
+
+            drop(host_cache);
+
+            // Update our buckets
+            if !node.addresses().is_empty() {
+                self.dht.update_node(&node.clone(), channel.clone()).await;
+            }
+        } else if channel.session_type_id() & SESSION_INBOUND != 0 {
+            // If it's an inbound connection, verify that we can connect to at
+            // least one of the provided external addresses.
+            // This may try to create a new outbound channel and it will update
+            // our buckets if successful.
+            let _ = self.verify_node_tx.send(node.clone()).await;
+        }
+
+        // Update the channel cache
+        self.dht.add_channel_to_cache(channel.info.id, node).await;
 
-        // Do the actual pinging process
-        let ping_result = self.do_ping(channel.clone()).await;
-        *result = Some(ping_result.clone());
-        ping_result
+        Ok(node.clone())
     }
 
     async fn store(

bin/fud/fud/src/lib.rs

@@ -21,35 +21,25 @@ use std::{
     io::ErrorKind,
     path::{Path, PathBuf},
     sync::Arc,
-    time::Duration,
 };
 
-use rand::{rngs::OsRng, Rng};
 use sled_overlay::sled;
 use smol::{
     channel,
     fs::{self, OpenOptions},
-    lock::{Mutex, RwLock},
+    lock::RwLock,
 };
-use tracing::{debug, error, info, warn};
+use tracing::{error, info, warn};
 
 use darkfi::{
-    dht::{
-        event::DhtEvent, tasks as dht_tasks, Dht, DhtHandler, DhtNode, DhtSettings, HostCacheItem,
-    },
+    dht::{tasks as dht_tasks, Dht, DhtHandler, DhtSettings},
     geode::{hash_to_string, ChunkedStorage, FileSequence, Geode, MAX_CHUNK_SIZE},
-    net::{
-        session::{SESSION_DIRECT, SESSION_INBOUND, SESSION_MANUAL, SESSION_OUTBOUND},
-        ChannelPtr, P2pPtr,
-    },
-    system::{timeout::timeout, ExecutorPtr, PublisherPtr, StoppableTask},
+    net::P2pPtr,
+    system::{ExecutorPtr, PublisherPtr, StoppableTask},
     util::{path::expand_path, time::Timestamp},
     Error, Result,
 };
-use darkfi_sdk::crypto::{
-    schnorr::{SchnorrPublic, SchnorrSecret},
-    SecretKey,
-};
+use darkfi_sdk::crypto::{schnorr::SchnorrSecret, SecretKey};
 use darkfi_serial::{deserialize_async, serialize_async};
 
 /// P2P protocols
@@ -101,18 +91,12 @@ use download::{fetch_chunks, fetch_metadata};
 pub mod dht;
 use dht::FudSeeder;
 
-use crate::{
-    dht::FudNode,
-    pow::PowSettings,
-    proto::{FudPingReply, FudPingRequest},
-};
+use crate::{dht::FudNode, pow::PowSettings};
 
 const SLED_PATH_TREE: &[u8] = b"_fud_paths";
 const SLED_FILE_SELECTION_TREE: &[u8] = b"_fud_file_selections";
 const SLED_SCRAP_TREE: &[u8] = b"_fud_scraps";
 
-type PingLock = Arc<Mutex<Option<Result<FudNode>>>>;
-
 pub struct Fud {
     /// Our own [`VerifiableNodeData`]
     pub node_data: Arc<RwLock<VerifiableNodeData>>,
@@ -143,8 +127,6 @@ pub struct Fud {
     /// is not saved to the filesystem in the downloaded files.
     /// "chunk/scrap hash -> chunk content"
     scrap_tree: sled::Tree,
-    /// Locks that prevent pinging the same channel multiple times at once.
-    ping_locks: Arc<Mutex<HashMap<u32, PingLock>>>,
     /// Get requests sender
     get_tx: channel::Sender<(blake3::Hash, PathBuf, FileSelection)>,
     /// Get requests receiver
@@ -234,7 +216,6 @@ impl Fud {
             file_selection_tree: sled_db.open_tree(SLED_FILE_SELECTION_TREE)?,
             scrap_tree: sled_db.open_tree(SLED_SCRAP_TREE)?,
             resources: Arc::new(RwLock::new(HashMap::new())),
-            ping_locks: Arc::new(Mutex::new(HashMap::new())),
             get_tx,
             get_rx,
             put_tx,
@@ -419,110 +400,6 @@ impl Fud {
         }
     }
 
-    async fn do_ping(&self, channel: ChannelPtr) -> Result<FudNode> {
-        debug!(target: "fud::DhtHandler::do_ping()", "Sending ping to {}", channel.display_address());
-
-        let dht = self.dht();
-
-        // Setup `FudPingReply` subscriber
-        let msg_subscriber = channel.subscribe_msg::<FudPingReply>().await.unwrap();
-
-        // Send `FudPingRequest`
-        let mut rng = OsRng;
-        let request = FudPingRequest { random: rng.gen() };
-        if channel.is_stopped() {
-            return Err(Error::ChannelStopped)
-        }
-        channel.send(&request).await?;
-
-        // Wait for `FudPingReply`
-        let reply = msg_subscriber.receive_with_timeout(dht.settings.timeout).await;
-        msg_subscriber.unsubscribe().await;
-        let reply = reply?;
-        let node = &reply.node;
-
-        // Verify the signature
-        if !node.data.public_key.verify(&request.random.to_be_bytes(), &reply.sig) {
-            warn!(target: "fud::do_ping()", "Received an invalid signature while pinging {}", channel.display_address());
-            self.dht
-                .event_publisher
-                .notify(DhtEvent::PingReceived {
-                    from: channel.clone(),
-                    result: Err(Error::InvalidSignature),
-                })
-                .await;
-            self.dht.cleanup_channel(channel.clone()).await;
-            channel.ban().await;
-            return Err(Error::InvalidSignature)
-        }
-
-        // Verify PoW
-        if let Err(e) = self.pow.write().await.verify_node(&node.data).await {
-            warn!(target: "fud::do_ping()", "Received an invalid PoW while pinging {}: {e}", channel.display_address());
-            self.dht
-                .event_publisher
-                .notify(DhtEvent::PingReceived { from: channel.clone(), result: Err(e.clone()) })
-                .await;
-            self.dht.cleanup_channel(channel.clone()).await;
-            channel.ban().await;
-            return Err(e)
-        }
-        self.dht
-            .event_publisher
-            .notify(DhtEvent::PingReceived { from: channel.clone(), result: Ok(node.id()) })
-            .await;
-
-        if channel.session_type_id() & (SESSION_OUTBOUND | SESSION_DIRECT | SESSION_MANUAL) != 0 {
-            // Wait for the other node to ping us
-            let ping_timeout = Duration::from_secs(10);
-
-            if let Err(e) = timeout(ping_timeout, dht.wait_fully_pinged(channel.info.id)).await {
-                dht.cleanup_channel(channel).await;
-                return Err(e.into())
-            }
-
-            let mut host_cache = dht.host_cache.write().await;
-
-            // If we had another node id for this host in our cache, remove
-            // the old one from the buckets and seeders
-            if let Some(cached) = host_cache.get(channel.address()) {
-                if cached.node_id != node.id() {
-                    dht.remove_node(&cached.node_id).await;
-
-                    for (_, seeders) in dht.hash_table.write().await.iter_mut() {
-                        seeders.retain(|seeder| seeder.node.id() != cached.node_id);
-                    }
-                }
-            }
-
-            // Update host cache
-            host_cache.insert(
-                channel.address().clone(),
-                HostCacheItem { last_ping: Timestamp::current_time(), node_id: node.id() },
-            );
-
-            drop(host_cache);
-
-            // Update our buckets
-            if !node.addresses().is_empty() {
-                dht.update_node(&node.clone(), channel.clone()).await;
-            }
-        } else if channel.session_type_id() & SESSION_INBOUND != 0 {
-            // If it's an inbound connection, verify that we can connect to at
-            // least one of the provided external addresses.
-            // This may try to create a new outbound channel and it will update
-            // our buckets if successful.
-            if let Ok((channel, _)) = dht.create_channel_to_node(node).await {
-                dht.cleanup_channel(channel).await;
-            }
-        }
-
-        // Update the channel cache
-        dht.add_channel_to_cache(channel.info.id, node).await;
-
-        Ok(node.clone())
-    }
-
     /// Verify if resources are complete and uncorrupted.
     /// If a resource is incomplete or corrupted, its status is changed to Incomplete.
     /// If a resource is complete, its status is changed to Seeding.

src/dht/mod.rs

@@ -99,6 +99,8 @@ pub struct DhtBucket<N: DhtNode> {
 /// Our local hash table, storing DHT keys and values
 pub type DhtHashTable<V> = Arc<RwLock<HashMap<blake3::Hash, V>>>;
 
+type PingLock<N> = Arc<Mutex<Option<Result<N>>>>;
+
 #[derive(Clone, Debug)]
 pub struct ChannelCacheItem<N: DhtNode> {
     /// The DHT node the channel is connected to.
@@ -135,6 +137,8 @@ pub struct Dht<H: DhtHandler> {
     pub channel_cache: Arc<RwLock<HashMap<u32, ChannelCacheItem<H::Node>>>>,
     /// Host address -> ChannelCacheItem
     pub host_cache: Arc<RwLock<HashMap<Url, HostCacheItem>>>,
+    /// Locks that prevent pinging the same channel multiple times at once.
+    ping_locks: Arc<Mutex<HashMap<u32, PingLock<H::Node>>>>,
     /// Add node sender
     pub add_node_tx: channel::Sender<(H::Node, ChannelPtr)>,
     /// Add node receiver
@@ -172,6 +176,7 @@ impl<H: DhtHandler> Dht<H> {
             bootstrapped: Arc::new(RwLock::new(false)),
             channel_cache: Arc::new(RwLock::new(HashMap::new())),
             host_cache: Arc::new(RwLock::new(HashMap::new())),
+            ping_locks: Arc::new(Mutex::new(HashMap::new())),
             add_node_tx,
             add_node_rx,
 
@@ -376,6 +381,35 @@ impl<H: DhtHandler> Dht<H> {
         bucket.nodes.retain(|node| node.id() != *node_id);
     }
 
+    /// Send a DHT ping to `channel` using the handler's ping method.
+    /// Prevents sending multiple pings at once to the same channel.
+    pub async fn ping(&self, channel: ChannelPtr) -> Result<H::Node> {
+        let lock_map = self.ping_locks.clone();
+        let mut locks = lock_map.lock().await;
+
+        // Get or create the lock
+        let lock = if let Some(lock) = locks.get(&channel.info.id) {
+            lock.clone()
+        } else {
+            let lock = Arc::new(Mutex::new(None));
+            locks.insert(channel.info.id, lock.clone());
+            lock
+        };
+        drop(locks);
+
+        // Acquire the lock
+        let mut result = lock.lock().await;
+
+        if let Some(res) = result.clone() {
+            return res
+        }
+
+        // Do the actual pinging process as defined by the handler
+        let ping_result = self.handler().await.ping(channel.clone()).await;
+        *result = Some(ping_result.clone());
+        ping_result
+    }
+
     /// Lookup algorithm for both nodes lookup and value lookup.
     async fn lookup(
         &self,
@@ -632,7 +666,7 @@ impl<H: DhtHandler> Dht<H> {
         }
         drop(channel_cache);
 
-        let node = self.handler().await.ping(channel.clone()).await;
+        let node = self.ping(channel.clone()).await;
         // If ping failed, cleanup the channel and abort
         if let Err(e) = node {
             self.cleanup_channel(channel).await;
@@ -698,8 +732,10 @@ impl<H: DhtHandler> Dht<H> {
     pub async fn cleanup_channel(&self, channel: ChannelPtr) {
         let channel_cache_lock = self.channel_cache.clone();
         let mut channel_cache = channel_cache_lock.write().await;
+        let mut ping_locks = self.ping_locks.lock().await;
         if self.p2p.session_direct().cleanup_channel(channel.clone()).await {
             channel_cache.remove(&channel.info.id);
+            ping_locks.remove(&channel.info.id);
         }
     }
 }

src/dht/tasks.rs

@@ -62,7 +62,9 @@ pub async fn events_task<H: DhtHandler>(handler: Arc<H>) -> Result<()> {
 /// Send a DHT ping request when there is a new channel, to know the node id of the new peer,
 /// Then fill the channel cache and the buckets
 pub async fn channel_task<H: DhtHandler>(handler: Arc<H>) -> Result<()> {
-    let channel_sub = handler.dht().p2p.hosts().subscribe_channel().await;
+    let dht = handler.dht();
+    let p2p = dht.p2p.clone();
+    let channel_sub = p2p.hosts().subscribe_channel().await;
     loop {
         let res = channel_sub.receive().await;
         if res.is_err() {
@@ -70,7 +72,7 @@ pub async fn channel_task<H: DhtHandler>(handler: Arc<H>) -> Result<()> {
         }
         let channel = res.unwrap();
 
-        let channel_cache_lock = handler.dht().channel_cache.clone();
+        let channel_cache_lock = dht.channel_cache.clone();
         let mut channel_cache = channel_cache_lock.write().await;
 
         // Skip this channel if it's not new
@@ -91,7 +93,7 @@ pub async fn channel_task<H: DhtHandler>(handler: Arc<H>) -> Result<()> {
 
         // It's a manual connection
         if channel.session_type_id() & SESSION_MANUAL != 0 {
-            let ping_res = handler.ping(channel.clone()).await;
+            let ping_res = dht.ping(channel.clone()).await;
 
             if let Err(e) = ping_res {
                 warn!(target: "dht::channel_task()", "Error while pinging manual connection (requesting node id) {}: {e}", channel.display_address());
@@ -101,10 +103,7 @@ pub async fn channel_task<H: DhtHandler>(handler: Arc<H>) -> Result<()> {
 
         // It's an outbound connection
         if channel.session_type_id() & SESSION_OUTBOUND != 0 {
-            let node = handler.ping(channel.clone()).await;
-            if node.is_err() {
-                continue;
-            }
+            let _ = dht.ping(channel.clone()).await;
 
             continue;
         }