/* This file is part of DarkFi (https://dark.fi) * * Copyright (C) 2020-2026 Dyne.org foundation * * This program is free software: you can redistribute it and/or modify * it under the terms of the GNU Affero General Public License as * published by the Free Software Foundation, either version 3 of the * License, or (at your option) any later version. * * This program is distributed in the hope that it will be useful, * but WITHOUT ANY WARRANTY; without even the implied warranty of * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the * GNU Affero General Public License for more details. * * You should have received a copy of the GNU Affero General Public License * along with this program. If not, see . */ use std::{ cmp::Eq, collections::{HashMap, HashSet}, fmt::Debug, hash::Hash, marker::{Send, Sync}, sync::{Arc, Weak}, }; use futures::stream::FuturesUnordered; use num_bigint::BigUint; use smol::{ channel, lock::{Mutex, RwLock, Semaphore}, stream::StreamExt, }; use tracing::{info, warn}; use url::Url; use crate::{ dht::event::DhtEvent, net::{ connector::Connector, session::{SESSION_DIRECT, SESSION_MANUAL}, ChannelPtr, Message, P2pPtr, }, system::{msleep, ExecutorPtr, Publisher, PublisherPtr, Subscription}, util::time::Timestamp, Error, Result, }; pub mod settings; pub use settings::{DhtSettings, DhtSettingsOpt}; pub mod handler; pub use handler::DhtHandler; pub mod tasks; pub mod event; pub trait DhtNode: Debug + Clone + Send + Sync + PartialEq + Eq + Hash { fn id(&self) -> blake3::Hash; fn addresses(&self) -> Vec; } /// Implements default Hash, PartialEq, and Eq for a struct implementing [`DhtNode`] #[macro_export] macro_rules! impl_dht_node_defaults { ($t:ty) => { impl std::hash::Hash for $t { fn hash(&self, state: &mut H) { self.id().hash(state); } } impl std::cmp::PartialEq for $t { fn eq(&self, other: &Self) -> bool { self.id() == other.id() } } impl std::cmp::Eq for $t {} }; } pub use impl_dht_node_defaults; enum DhtLookupType { Nodes, Value, } pub enum DhtLookupReply { Nodes(Vec), Value(V), NodesAndValue(Vec, V), } pub struct DhtBucket { pub nodes: Vec, } /// Our local hash table, storing DHT keys and values pub type DhtHashTable = Arc>>; type PingLock = Arc>>>; #[derive(Clone, Debug)] pub struct ChannelCacheItem { /// The DHT node the channel is connected to. pub node: Option, /// The last time this channel was used by the [`DhtHandler`]. It's used /// to stop inbound connections in [`crate::dht::tasks::disconnect_inbounds_task()`]. pub last_used: Timestamp, /// Have we already received a DHT ping from this channel? pub ping_received: bool, /// Have we already sent a DHT ping to this channel? pub ping_sent: bool, } #[derive(Clone, Debug)] pub struct HostCacheItem { /// The last time we tried to send a DHT ping to this host. pub last_ping: Timestamp, /// The last known node id for this host. pub node_id: blake3::Hash, } pub struct Dht { /// [`DhtHandler`] that implements application-specific behaviors over a [`Dht`] pub handler: RwLock>, /// Are we bootstrapped? pub bootstrapped: Arc>, /// Vec of buckets pub buckets: Arc>>>, /// Our local hash table, storing a part of the full DHT keys/values pub hash_table: DhtHashTable, /// Number of buckets pub n_buckets: usize, /// Channel ID -> ChannelCacheItem pub channel_cache: Arc>>>, /// Host address -> ChannelCacheItem pub host_cache: Arc>>, /// Locks that prevent pinging the same channel multiple times at once. ping_locks: Arc>>>, /// Add node sender pub add_node_tx: channel::Sender<(H::Node, ChannelPtr)>, /// Add node receiver pub add_node_rx: channel::Receiver<(H::Node, ChannelPtr)>, /// DHT settings pub settings: DhtSettings, /// DHT event publisher pub event_publisher: PublisherPtr>, /// P2P network pointer pub p2p: P2pPtr, /// Connector to create manual connections pub connector: Connector, /// Global multithreaded executor reference pub executor: ExecutorPtr, } impl Dht { pub async fn new(settings: &DhtSettings, p2p: P2pPtr, ex: ExecutorPtr) -> Self { // Create empty buckets let mut buckets = vec![]; for _ in 0..256 { buckets.push(DhtBucket { nodes: vec![] }) } let (add_node_tx, add_node_rx) = smol::channel::unbounded(); let session_weak = Arc::downgrade(&p2p.session_manual()); let connector = Connector::new(p2p.settings(), session_weak); Self { handler: RwLock::new(Weak::new()), buckets: Arc::new(RwLock::new(buckets)), hash_table: Arc::new(RwLock::new(HashMap::new())), n_buckets: 256, 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, event_publisher: Publisher::new(), settings: settings.clone(), p2p: p2p.clone(), connector, executor: ex, } } pub async fn handler(&self) -> Arc { self.handler.read().await.upgrade().unwrap() } pub async fn is_bootstrapped(&self) -> bool { let bootstrapped = self.bootstrapped.read().await; *bootstrapped } pub async fn set_bootstrapped(&self, value: bool) { let mut bootstrapped = self.bootstrapped.write().await; *bootstrapped = value; } pub async fn subscribe(&self) -> Subscription> { self.event_publisher.clone().subscribe().await } /// Get the distance between `key_1` and `key_2` pub fn distance(&self, key_1: &blake3::Hash, key_2: &blake3::Hash) -> [u8; 32] { let bytes1 = key_1.as_bytes(); let bytes2 = key_2.as_bytes(); let mut result_bytes = [0u8; 32]; for i in 0..32 { result_bytes[i] = bytes1[i] ^ bytes2[i]; } result_bytes } /// Sort `nodes` by distance from `key` pub fn sort_by_distance(&self, nodes: &mut [H::Node], key: &blake3::Hash) { nodes.sort_by(|a, b| { let distance_a = BigUint::from_bytes_be(&self.distance(key, &a.id())); let distance_b = BigUint::from_bytes_be(&self.distance(key, &b.id())); distance_a.cmp(&distance_b) }); } /// `key` -> bucket index pub async fn get_bucket_index(&self, self_node_id: &blake3::Hash, key: &blake3::Hash) -> usize { if key == self_node_id { return 0; } let distance = self.distance(self_node_id, key); let mut leading_zeros = 0; for &byte in &distance { if byte == 0 { leading_zeros += 8; } else { leading_zeros += byte.leading_zeros() as usize; break; } } let bucket_index = self.n_buckets - leading_zeros; std::cmp::min(bucket_index, self.n_buckets - 1) } /// Get `n` closest known nodes to a key /// TODO: Can be optimized pub async fn find_neighbors(&self, key: &blake3::Hash, n: usize) -> Vec { let buckets_lock = self.buckets.clone(); let buckets = buckets_lock.read().await; let mut neighbors = Vec::new(); for i in 0..self.n_buckets { if let Some(bucket) = buckets.get(i) { neighbors.extend(bucket.nodes.iter().cloned()); } } self.sort_by_distance(&mut neighbors, key); neighbors.truncate(n); neighbors } /// Channel ID -> [`DhtNode`] pub async fn get_node_from_channel(&self, channel_id: u32) -> Option { let channel_cache_lock = self.channel_cache.clone(); let channel_cache = channel_cache_lock.read().await; if let Some(cached) = channel_cache.get(&channel_id) { return cached.node.clone(); } None } /// Reset the DHT state (nodes and hash table) pub async fn reset(&self) { let mut bootstrapped = self.bootstrapped.write().await; *bootstrapped = false; let mut buckets = vec![]; for _ in 0..256 { buckets.push(DhtBucket { nodes: vec![] }) } *self.buckets.write().await = buckets; *self.hash_table.write().await = HashMap::new(); } /// Add `value` to our hash table and send `message` for a `key` to the closest nodes found pub async fn announce( &self, key: &blake3::Hash, value: &H::Value, message: &M, ) -> Result<()> { let self_node = self.handler().await.node().await?; if self_node.addresses().is_empty() { return Err(().into()); // TODO } self.handler().await.add_value(key, value).await; let nodes = self.lookup_nodes(key).await; info!(target: "dht::announce", "[DHT] Announcing {} to {} nodes", H::key_to_string(key), nodes.len()); for node in nodes { if let Ok((channel, _)) = self.get_channel(&node).await { let _ = channel.send(message).await; self.cleanup_channel(channel).await; } } Ok(()) } /// Lookup our own node id pub async fn bootstrap(&self) { let self_node = self.handler().await.node().await; if self_node.is_err() { return; } let self_node = self_node.unwrap(); self.set_bootstrapped(true).await; info!(target: "dht::bootstrap", "[DHT] Bootstrapping"); self.event_publisher.notify(DhtEvent::BootstrapStarted).await; let _nodes = self.lookup_nodes(&self_node.id()).await; // if nodes.is_empty() { // self.set_bootstrapped(false).await; // } else { // } self.event_publisher.notify(DhtEvent::BootstrapCompleted).await; } // TODO: Optimize this async fn on_new_node(&self, node: &H::Node, channel: ChannelPtr) { info!(target: "dht::on_new_node", "[DHT] Found new node {}", H::key_to_string(&node.id())); // If this is the first node we know about then bootstrap if !self.is_bootstrapped().await { self.bootstrap().await; } // Send keys that are closer to this node than we are let self_node = self.handler().await.node().await; if self_node.is_err() { return; } let self_id = self_node.unwrap().id(); for (key, value) in self.hash_table.read().await.iter() { let node_distance = BigUint::from_bytes_be(&self.distance(key, &node.id())); let self_distance = BigUint::from_bytes_be(&self.distance(key, &self_id)); if node_distance <= self_distance { let _ = self.handler().await.store(channel.clone(), key, value).await; } } } /// Move a node to the tail in its bucket, /// to show that it is the most recently seen in the bucket. /// If the node is not in a bucket it will be added using `add_node`. pub async fn update_node(&self, node: &H::Node, channel: ChannelPtr) { self.p2p.session_direct().inc_channel_usage(&channel, 1).await; if let Err(e) = self.add_node_tx.send((node.clone(), channel.clone())).await { warn!(target: "dht::update_node", "[DHT] Cannot add node {}: {e}", H::key_to_string(&node.id())) } } /// Remove a node from the buckets. pub async fn remove_node(&self, node_id: &blake3::Hash) { let handler = self.handler().await; let self_node = handler.node().await; if self_node.is_err() { return; } let bucket_index = handler.dht().get_bucket_index(&self_node.unwrap().id(), node_id).await; let buckets_lock = handler.dht().buckets.clone(); let mut buckets = buckets_lock.write().await; let bucket = &mut buckets[bucket_index]; 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 { 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, key: blake3::Hash, lookup_type: DhtLookupType, ) -> (Vec, Vec) { let net_settings = self.p2p.settings().read_arc().await; let active_profiles = net_settings.active_profiles.clone(); drop(net_settings); let external_addrs = self.p2p.hosts().external_addrs().await; let (k, a) = (self.settings.k, self.settings.alpha); let semaphore = Arc::new(Semaphore::new(self.settings.concurrency)); let queried_addrs = Arc::new(Mutex::new(HashSet::new())); let mut seen_nodes = HashSet::new(); let mut nodes_to_visit = self.find_neighbors(&key, k).await; let mut result = Vec::new(); let mut futures = FuturesUnordered::new(); let mut consecutive_stalls = 0; let mut values = Vec::new(); let distance_check = |(furthest, next): (&H::Node, &H::Node)| { BigUint::from_bytes_be(&self.distance(&key, &furthest.id())) < BigUint::from_bytes_be(&self.distance(&key, &next.id())) }; // Create a channel if necessary and send a FIND NODES or FIND VALUE // request to `addr` let lookup = async |node: H::Node, key, addrs: Vec| { let _permit = semaphore.acquire().await; // Try all valid addresses for the node let mut last_err = None; for addr in addrs { let mut queried_addrs_set = queried_addrs.lock().await; // Skip if this address has already been queried if queried_addrs_set.contains(&addr) { continue; } queried_addrs_set.insert(addr.clone()); drop(queried_addrs_set); // Try to create or find an existing channel let channel = self.create_channel(&addr).await.map(|(ch, _)| ch); if let Err(e) = channel { last_err = Some(e); continue } let channel = channel.unwrap(); let handler = self.handler().await; let res = match &lookup_type { DhtLookupType::Nodes => { info!(target: "dht::lookup", "[DHT] [LOOKUP] Querying node {} for nodes lookup of key {}", H::key_to_string(&node.id()), H::key_to_string(key)); handler.find_nodes(channel.clone(), key).await.map(DhtLookupReply::Nodes) } DhtLookupType::Value => { info!(target: "dht::lookup", "[DHT] [LOOKUP] Querying node {} for value lookup of key {}", H::key_to_string(&node.id()), H::key_to_string(key)); handler.find_value(channel.clone(), key).await } }; self.cleanup_channel(channel).await; if res.is_ok() { return (node, res) } last_err = res.err(); } if let Some(e) = last_err { return (node, Err(e)) } (node, Err(Error::Custom("All node's addresses failed".to_string()))) }; // Spawn up to `alpha` futures for lookup() let spawn_futures = async |nodes_to_visit: &mut Vec, futures: &mut FuturesUnordered<_>| { for _ in 0..a { if !nodes_to_visit.is_empty() { let node = nodes_to_visit.remove(0); let valid_addrs: Vec = node .addresses() .iter() .filter(|addr| { active_profiles.contains(&addr.scheme().to_string()) && !external_addrs.contains(addr) }) .cloned() .collect(); if !valid_addrs.is_empty() { futures.push(Box::pin(lookup(node, &key, valid_addrs))); } } } }; // Initial futures spawn_futures(&mut nodes_to_visit, &mut futures).await; // Process lookup responses while let Some((queried_node, res)) = futures.next().await { if let Err(e) = res { warn!(target: "dht::lookup", "[DHT] [LOOKUP] Error in lookup: {e}"); // Spawn next `alpha` futures if there are no more futures but // we still have nodes to visit if futures.is_empty() { spawn_futures(&mut nodes_to_visit, &mut futures).await; } continue; } let (nodes, value) = match res.unwrap() { DhtLookupReply::Nodes(nodes) => (Some(nodes), None), DhtLookupReply::Value(value) => (None, Some(value)), DhtLookupReply::NodesAndValue(nodes, value) => (Some(nodes), Some(value)), }; // Send the value we found to the publisher if let Some(value) = value { info!(target: "dht::lookup", "[DHT] [LOOKUP] Found value for {} from {}", H::key_to_string(&key), H::key_to_string(&queried_node.id())); values.push(value.clone()); self.event_publisher.notify(DhtEvent::ValueFound { key, value }).await; } // Update nodes_to_visit if let Some(mut nodes) = nodes { if !nodes.is_empty() { info!(target: "dht::lookup", "[DHT] [LOOKUP] Found {} nodes from {}", nodes.len(), H::key_to_string(&queried_node.id())); self.event_publisher .notify(DhtEvent::NodesFound { key, nodes: nodes.clone() }) .await; // Remove our own node and duplicates if let Ok(self_node) = self.handler().await.node().await { let self_id = self_node.id(); nodes.retain(|node: &H::Node| { node.id() != self_id && seen_nodes.insert(node.id()) }); } // Add new nodes to the list of nodes to visit nodes_to_visit.extend(nodes.clone()); self.sort_by_distance(&mut nodes_to_visit, &key); } } result.push(queried_node); self.sort_by_distance(&mut result, &key); // Early termination logic: // The closest node to visit must be further than the furthest // queried node, 3 consecutive times if result.len() >= k && result.last().zip(nodes_to_visit.first()).is_some_and(distance_check) { consecutive_stalls += 1; if consecutive_stalls >= 3 { break; } } else { consecutive_stalls = 0; } // Spawn next `alpha` futures spawn_futures(&mut nodes_to_visit, &mut futures).await; } info!(target: "dht::lookup", "[DHT] [LOOKUP] Lookup for {} completed", H::key_to_string(&key)); let nodes: Vec<_> = result.into_iter().take(k).collect(); (nodes, values) } /// Find `k` nodes closest to a key pub async fn lookup_nodes(&self, key: &blake3::Hash) -> Vec { info!(target: "dht::lookup_nodes", "[DHT] [LOOKUP] Starting node lookup for key {}", H::key_to_string(key)); self.event_publisher.notify(DhtEvent::NodesLookupStarted { key: *key }).await; let (nodes, _) = self.lookup(*key, DhtLookupType::Nodes).await; self.event_publisher .notify(DhtEvent::NodesLookupCompleted { key: *key, nodes: nodes.clone() }) .await; nodes } /// Find value for `key` pub async fn lookup_value(&self, key: &blake3::Hash) -> (Vec, Vec) { info!(target: "dht::lookup_value", "[DHT] [LOOKUP] Starting value lookup for key {}", H::key_to_string(key)); self.event_publisher.notify(DhtEvent::ValueLookupStarted { key: *key }).await; let (nodes, values) = self.lookup(*key, DhtLookupType::Value).await; self.event_publisher .notify(DhtEvent::ValueLookupCompleted { key: *key, nodes: nodes.clone(), values: values.clone(), }) .await; (nodes, values) } /// Update a channel's `last_used` field in the channel cache. pub async fn update_channel(&self, channel_id: u32) { let channel_cache_lock = self.channel_cache.clone(); let mut channel_cache = channel_cache_lock.write().await; if let Some(cached) = channel_cache.get_mut(&channel_id) { cached.last_used = Timestamp::current_time(); } } /// Get a channel (existing or create a new one) to `node`. /// Don't forget to call `cleanup_channel()` once you are done with it. pub async fn get_channel(&self, node: &H::Node) -> Result<(ChannelPtr, H::Node)> { let node_id = node.id(); // Look in the channel cache for a channel connected to this node. // We skip direct session channels, for those we will call // `create_channel()` which increments the sessions's usage counter. let channel_cache = self.channel_cache.read().await.clone(); if let Some((channel_id, cached)) = channel_cache .clone() .iter() .find(|(_, cached)| cached.node.clone().is_some_and(|n| n.id() == node_id)) { if let Some(channel) = self.p2p.get_channel(*channel_id) { if channel.session_type_id() & SESSION_DIRECT == 0 { if channel.is_stopped() { self.cleanup_channel(channel).await; } else { return Ok((channel, cached.node.clone().unwrap())) } } } } self.create_channel_to_node(node).await } /// Create a channel in the direct session, ping the peer, add the /// DHT node to our buckets and the channel to our channel cache. pub async fn create_channel(&self, addr: &Url) -> Result<(ChannelPtr, H::Node)> { let external_addrs = self.p2p.hosts().external_addrs().await; if external_addrs.contains(addr) { return Err(Error::Custom( "Can't create a channel to our own external address".to_string(), )) } let channel = self.p2p.session_direct().get_channel(addr).await?; let channel_cache = self.channel_cache.read().await; if let Some(cached) = channel_cache.get(&channel.info.id) { if let Some(node) = &cached.node { return Ok((channel, node.clone())) } } drop(channel_cache); 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; return Err(e); } let node = node.unwrap(); self.add_channel_to_cache(channel.info.id, &node).await; Ok((channel, node)) } pub async fn create_channel_to_node(&self, node: &H::Node) -> Result<(ChannelPtr, H::Node)> { let net_settings = self.p2p.settings().read_arc().await; let active_profiles = net_settings.active_profiles.clone(); drop(net_settings); // Create a channel let mut addrs = node.addresses().clone(); addrs.retain(|addr| active_profiles.contains(&addr.scheme().to_string())); for addr in addrs { let res = self.create_channel(&addr).await; if res.is_err() { continue; } let (channel, node) = res.unwrap(); return Ok((channel, node)); } Err(Error::Custom("Could not create channel".to_string())) } /// Insert a channel to the DHT's channel cache. If the channel is already /// in the cache, `last_used` is updated. pub async fn add_channel_to_cache(&self, channel_id: u32, node: &H::Node) { let mut channel_cache = self.channel_cache.write().await; channel_cache .entry(channel_id) .and_modify(|c| c.last_used = Timestamp::current_time()) .or_insert(ChannelCacheItem { node: Some(node.clone()), last_used: Timestamp::current_time(), ping_received: false, ping_sent: false, }); } /// Wait until we received a DHT ping and sent a DHT ping on a channel. pub async fn wait_fully_pinged(&self, channel_id: u32) -> Result<()> { loop { let channel_cache = self.channel_cache.read().await; let cached = channel_cache .get(&channel_id) .ok_or(Error::Custom("Missing channel".to_string()))?; if cached.ping_received && cached.ping_sent { return Ok(()) } drop(channel_cache); msleep(100).await; } } /// Call [`crate::net::session::DirectSession::cleanup_channel()`] and cleanup the DHT caches. 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); } } }