mod.rs 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770
  1. /* This file is part of DarkFi (https://dark.fi)
  2. *
  3. * Copyright (C) 2020-2026 Dyne.org foundation
  4. *
  5. * This program is free software: you can redistribute it and/or modify
  6. * it under the terms of the GNU Affero General Public License as
  7. * published by the Free Software Foundation, either version 3 of the
  8. * License, or (at your option) any later version.
  9. *
  10. * This program is distributed in the hope that it will be useful,
  11. * but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  13. * GNU Affero General Public License for more details.
  14. *
  15. * You should have received a copy of the GNU Affero General Public License
  16. * along with this program. If not, see <https://www.gnu.org/licenses/>.
  17. */
  18. use std::{
  19. cmp::Eq,
  20. collections::{HashMap, HashSet},
  21. fmt::Debug,
  22. hash::Hash,
  23. marker::{Send, Sync},
  24. sync::{Arc, Weak},
  25. };
  26. use futures::stream::FuturesUnordered;
  27. use num_bigint::BigUint;
  28. use smol::{
  29. channel,
  30. lock::{Mutex, RwLock, Semaphore},
  31. stream::StreamExt,
  32. };
  33. use tracing::{info, warn};
  34. use url::Url;
  35. use crate::{
  36. dht::event::DhtEvent,
  37. net::{
  38. connector::Connector,
  39. session::{SESSION_DIRECT, SESSION_MANUAL},
  40. ChannelPtr, Message, P2pPtr,
  41. },
  42. system::{msleep, ExecutorPtr, Publisher, PublisherPtr, Subscription},
  43. util::time::Timestamp,
  44. Error, Result,
  45. };
  46. pub mod settings;
  47. pub use settings::{DhtSettings, DhtSettingsOpt};
  48. pub mod handler;
  49. pub use handler::DhtHandler;
  50. pub mod tasks;
  51. pub mod event;
  52. pub trait DhtNode: Debug + Clone + Send + Sync + PartialEq + Eq + Hash {
  53. fn id(&self) -> blake3::Hash;
  54. fn addresses(&self) -> Vec<Url>;
  55. }
  56. /// Implements default Hash, PartialEq, and Eq for a struct implementing [`DhtNode`]
  57. #[macro_export]
  58. macro_rules! impl_dht_node_defaults {
  59. ($t:ty) => {
  60. impl std::hash::Hash for $t {
  61. fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
  62. self.id().hash(state);
  63. }
  64. }
  65. impl std::cmp::PartialEq for $t {
  66. fn eq(&self, other: &Self) -> bool {
  67. self.id() == other.id()
  68. }
  69. }
  70. impl std::cmp::Eq for $t {}
  71. };
  72. }
  73. pub use impl_dht_node_defaults;
  74. enum DhtLookupType {
  75. Nodes,
  76. Value,
  77. }
  78. pub enum DhtLookupReply<N: DhtNode, V> {
  79. Nodes(Vec<N>),
  80. Value(V),
  81. NodesAndValue(Vec<N>, V),
  82. }
  83. pub struct DhtBucket<N: DhtNode> {
  84. pub nodes: Vec<N>,
  85. }
  86. /// Our local hash table, storing DHT keys and values
  87. pub type DhtHashTable<V> = Arc<RwLock<HashMap<blake3::Hash, V>>>;
  88. type PingLock<N> = Arc<Mutex<Option<Result<N>>>>;
  89. #[derive(Clone, Debug)]
  90. pub struct ChannelCacheItem<N: DhtNode> {
  91. /// The DHT node the channel is connected to.
  92. pub node: Option<N>,
  93. /// The last time this channel was used by the [`DhtHandler`]. It's used
  94. /// to stop inbound connections in [`crate::dht::tasks::disconnect_inbounds_task()`].
  95. pub last_used: Timestamp,
  96. /// Have we already received a DHT ping from this channel?
  97. pub ping_received: bool,
  98. /// Have we already sent a DHT ping to this channel?
  99. pub ping_sent: bool,
  100. }
  101. #[derive(Clone, Debug)]
  102. pub struct HostCacheItem {
  103. /// The last time we tried to send a DHT ping to this host.
  104. pub last_ping: Timestamp,
  105. /// The last known node id for this host.
  106. pub node_id: blake3::Hash,
  107. }
  108. pub struct Dht<H: DhtHandler> {
  109. /// [`DhtHandler`] that implements application-specific behaviors over a [`Dht`]
  110. pub handler: RwLock<Weak<H>>,
  111. /// Are we bootstrapped?
  112. pub bootstrapped: Arc<RwLock<bool>>,
  113. /// Vec of buckets
  114. pub buckets: Arc<RwLock<Vec<DhtBucket<H::Node>>>>,
  115. /// Our local hash table, storing a part of the full DHT keys/values
  116. pub hash_table: DhtHashTable<H::Value>,
  117. /// Number of buckets
  118. pub n_buckets: usize,
  119. /// Channel ID -> ChannelCacheItem
  120. pub channel_cache: Arc<RwLock<HashMap<u32, ChannelCacheItem<H::Node>>>>,
  121. /// Host address -> ChannelCacheItem
  122. pub host_cache: Arc<RwLock<HashMap<Url, HostCacheItem>>>,
  123. /// Locks that prevent pinging the same channel multiple times at once.
  124. ping_locks: Arc<Mutex<HashMap<u32, PingLock<H::Node>>>>,
  125. /// Add node sender
  126. pub add_node_tx: channel::Sender<(H::Node, ChannelPtr)>,
  127. /// Add node receiver
  128. pub add_node_rx: channel::Receiver<(H::Node, ChannelPtr)>,
  129. /// DHT settings
  130. pub settings: DhtSettings,
  131. /// DHT event publisher
  132. pub event_publisher: PublisherPtr<DhtEvent<H::Node, H::Value>>,
  133. /// P2P network pointer
  134. pub p2p: P2pPtr,
  135. /// Connector to create manual connections
  136. pub connector: Connector,
  137. /// Global multithreaded executor reference
  138. pub executor: ExecutorPtr,
  139. }
  140. impl<H: DhtHandler> Dht<H> {
  141. pub async fn new(settings: &DhtSettings, p2p: P2pPtr, ex: ExecutorPtr) -> Self {
  142. // Create empty buckets
  143. let mut buckets = vec![];
  144. for _ in 0..256 {
  145. buckets.push(DhtBucket { nodes: vec![] })
  146. }
  147. let (add_node_tx, add_node_rx) = smol::channel::unbounded();
  148. let session_weak = Arc::downgrade(&p2p.session_manual());
  149. let connector = Connector::new(p2p.settings(), session_weak);
  150. Self {
  151. handler: RwLock::new(Weak::new()),
  152. buckets: Arc::new(RwLock::new(buckets)),
  153. hash_table: Arc::new(RwLock::new(HashMap::new())),
  154. n_buckets: 256,
  155. bootstrapped: Arc::new(RwLock::new(false)),
  156. channel_cache: Arc::new(RwLock::new(HashMap::new())),
  157. host_cache: Arc::new(RwLock::new(HashMap::new())),
  158. ping_locks: Arc::new(Mutex::new(HashMap::new())),
  159. add_node_tx,
  160. add_node_rx,
  161. event_publisher: Publisher::new(),
  162. settings: settings.clone(),
  163. p2p: p2p.clone(),
  164. connector,
  165. executor: ex,
  166. }
  167. }
  168. pub async fn handler(&self) -> Arc<H> {
  169. self.handler.read().await.upgrade().unwrap()
  170. }
  171. pub async fn is_bootstrapped(&self) -> bool {
  172. let bootstrapped = self.bootstrapped.read().await;
  173. *bootstrapped
  174. }
  175. pub async fn set_bootstrapped(&self, value: bool) {
  176. let mut bootstrapped = self.bootstrapped.write().await;
  177. *bootstrapped = value;
  178. }
  179. pub async fn subscribe(&self) -> Subscription<DhtEvent<H::Node, H::Value>> {
  180. self.event_publisher.clone().subscribe().await
  181. }
  182. /// Get the distance between `key_1` and `key_2`
  183. pub fn distance(&self, key_1: &blake3::Hash, key_2: &blake3::Hash) -> [u8; 32] {
  184. let bytes1 = key_1.as_bytes();
  185. let bytes2 = key_2.as_bytes();
  186. let mut result_bytes = [0u8; 32];
  187. for i in 0..32 {
  188. result_bytes[i] = bytes1[i] ^ bytes2[i];
  189. }
  190. result_bytes
  191. }
  192. /// Sort `nodes` by distance from `key`
  193. pub fn sort_by_distance(&self, nodes: &mut [H::Node], key: &blake3::Hash) {
  194. nodes.sort_by(|a, b| {
  195. let distance_a = BigUint::from_bytes_be(&self.distance(key, &a.id()));
  196. let distance_b = BigUint::from_bytes_be(&self.distance(key, &b.id()));
  197. distance_a.cmp(&distance_b)
  198. });
  199. }
  200. /// `key` -> bucket index
  201. pub async fn get_bucket_index(&self, self_node_id: &blake3::Hash, key: &blake3::Hash) -> usize {
  202. if key == self_node_id {
  203. return 0;
  204. }
  205. let distance = self.distance(self_node_id, key);
  206. let mut leading_zeros = 0;
  207. for &byte in &distance {
  208. if byte == 0 {
  209. leading_zeros += 8;
  210. } else {
  211. leading_zeros += byte.leading_zeros() as usize;
  212. break;
  213. }
  214. }
  215. let bucket_index = self.n_buckets - leading_zeros;
  216. std::cmp::min(bucket_index, self.n_buckets - 1)
  217. }
  218. /// Get `n` closest known nodes to a key
  219. /// TODO: Can be optimized
  220. pub async fn find_neighbors(&self, key: &blake3::Hash, n: usize) -> Vec<H::Node> {
  221. let buckets_lock = self.buckets.clone();
  222. let buckets = buckets_lock.read().await;
  223. let mut neighbors = Vec::new();
  224. for i in 0..self.n_buckets {
  225. if let Some(bucket) = buckets.get(i) {
  226. neighbors.extend(bucket.nodes.iter().cloned());
  227. }
  228. }
  229. self.sort_by_distance(&mut neighbors, key);
  230. neighbors.truncate(n);
  231. neighbors
  232. }
  233. /// Channel ID -> [`DhtNode`]
  234. pub async fn get_node_from_channel(&self, channel_id: u32) -> Option<H::Node> {
  235. let channel_cache_lock = self.channel_cache.clone();
  236. let channel_cache = channel_cache_lock.read().await;
  237. if let Some(cached) = channel_cache.get(&channel_id) {
  238. return cached.node.clone();
  239. }
  240. None
  241. }
  242. /// Reset the DHT state (nodes and hash table)
  243. pub async fn reset(&self) {
  244. let mut bootstrapped = self.bootstrapped.write().await;
  245. *bootstrapped = false;
  246. let mut buckets = vec![];
  247. for _ in 0..256 {
  248. buckets.push(DhtBucket { nodes: vec![] })
  249. }
  250. *self.buckets.write().await = buckets;
  251. *self.hash_table.write().await = HashMap::new();
  252. }
  253. /// Add `value` to our hash table and send `message` for a `key` to the closest nodes found
  254. pub async fn announce<M: Message>(
  255. &self,
  256. key: &blake3::Hash,
  257. value: &H::Value,
  258. message: &M,
  259. ) -> Result<()> {
  260. let self_node = self.handler().await.node().await?;
  261. if self_node.addresses().is_empty() {
  262. return Err(().into()); // TODO
  263. }
  264. self.handler().await.add_value(key, value).await;
  265. let nodes = self.lookup_nodes(key).await;
  266. info!(target: "dht::announce", "[DHT] Announcing {} to {} nodes", H::key_to_string(key), nodes.len());
  267. for node in nodes {
  268. if let Ok((channel, _)) = self.get_channel(&node).await {
  269. let _ = channel.send(message).await;
  270. self.cleanup_channel(channel).await;
  271. }
  272. }
  273. Ok(())
  274. }
  275. /// Lookup our own node id
  276. pub async fn bootstrap(&self) {
  277. let self_node = self.handler().await.node().await;
  278. if self_node.is_err() {
  279. return;
  280. }
  281. let self_node = self_node.unwrap();
  282. self.set_bootstrapped(true).await;
  283. info!(target: "dht::bootstrap", "[DHT] Bootstrapping");
  284. self.event_publisher.notify(DhtEvent::BootstrapStarted).await;
  285. let _nodes = self.lookup_nodes(&self_node.id()).await;
  286. // if nodes.is_empty() {
  287. // self.set_bootstrapped(false).await;
  288. // } else {
  289. // }
  290. self.event_publisher.notify(DhtEvent::BootstrapCompleted).await;
  291. }
  292. // TODO: Optimize this
  293. async fn on_new_node(&self, node: &H::Node, channel: ChannelPtr) {
  294. info!(target: "dht::on_new_node", "[DHT] Found new node {}", H::key_to_string(&node.id()));
  295. // If this is the first node we know about then bootstrap
  296. if !self.is_bootstrapped().await {
  297. self.bootstrap().await;
  298. }
  299. // Send keys that are closer to this node than we are
  300. let self_node = self.handler().await.node().await;
  301. if self_node.is_err() {
  302. return;
  303. }
  304. let self_id = self_node.unwrap().id();
  305. for (key, value) in self.hash_table.read().await.iter() {
  306. let node_distance = BigUint::from_bytes_be(&self.distance(key, &node.id()));
  307. let self_distance = BigUint::from_bytes_be(&self.distance(key, &self_id));
  308. if node_distance <= self_distance {
  309. let _ = self.handler().await.store(channel.clone(), key, value).await;
  310. }
  311. }
  312. }
  313. /// Move a node to the tail in its bucket,
  314. /// to show that it is the most recently seen in the bucket.
  315. /// If the node is not in a bucket it will be added using `add_node`.
  316. pub async fn update_node(&self, node: &H::Node, channel: ChannelPtr) {
  317. self.p2p.session_direct().inc_channel_usage(&channel, 1).await;
  318. if let Err(e) = self.add_node_tx.send((node.clone(), channel.clone())).await {
  319. warn!(target: "dht::update_node", "[DHT] Cannot add node {}: {e}", H::key_to_string(&node.id()))
  320. }
  321. }
  322. /// Remove a node from the buckets.
  323. pub async fn remove_node(&self, node_id: &blake3::Hash) {
  324. let handler = self.handler().await;
  325. let self_node = handler.node().await;
  326. if self_node.is_err() {
  327. return;
  328. }
  329. let bucket_index = handler.dht().get_bucket_index(&self_node.unwrap().id(), node_id).await;
  330. let buckets_lock = handler.dht().buckets.clone();
  331. let mut buckets = buckets_lock.write().await;
  332. let bucket = &mut buckets[bucket_index];
  333. bucket.nodes.retain(|node| node.id() != *node_id);
  334. }
  335. /// Send a DHT ping to `channel` using the handler's ping method.
  336. /// Prevents sending multiple pings at once to the same channel.
  337. pub async fn ping(&self, channel: ChannelPtr) -> Result<H::Node> {
  338. let lock_map = self.ping_locks.clone();
  339. let mut locks = lock_map.lock().await;
  340. // Get or create the lock
  341. let lock = if let Some(lock) = locks.get(&channel.info.id) {
  342. lock.clone()
  343. } else {
  344. let lock = Arc::new(Mutex::new(None));
  345. locks.insert(channel.info.id, lock.clone());
  346. lock
  347. };
  348. drop(locks);
  349. // Acquire the lock
  350. let mut result = lock.lock().await;
  351. if let Some(res) = result.clone() {
  352. return res
  353. }
  354. // Do the actual pinging process as defined by the handler
  355. let ping_result = self.handler().await.ping(channel.clone()).await;
  356. *result = Some(ping_result.clone());
  357. ping_result
  358. }
  359. /// Lookup algorithm for both nodes lookup and value lookup.
  360. async fn lookup(
  361. &self,
  362. key: blake3::Hash,
  363. lookup_type: DhtLookupType,
  364. ) -> (Vec<H::Node>, Vec<H::Value>) {
  365. let net_settings = self.p2p.settings().read_arc().await;
  366. let active_profiles = net_settings.active_profiles.clone();
  367. drop(net_settings);
  368. let external_addrs = self.p2p.hosts().external_addrs().await;
  369. let (k, a) = (self.settings.k, self.settings.alpha);
  370. let semaphore = Arc::new(Semaphore::new(self.settings.concurrency));
  371. let queried_addrs = Arc::new(Mutex::new(HashSet::new()));
  372. let mut seen_nodes = HashSet::new();
  373. let mut nodes_to_visit = self.find_neighbors(&key, k).await;
  374. let mut result = Vec::new();
  375. let mut futures = FuturesUnordered::new();
  376. let mut consecutive_stalls = 0;
  377. let mut values = Vec::new();
  378. let distance_check = |(furthest, next): (&H::Node, &H::Node)| {
  379. BigUint::from_bytes_be(&self.distance(&key, &furthest.id())) <
  380. BigUint::from_bytes_be(&self.distance(&key, &next.id()))
  381. };
  382. // Create a channel if necessary and send a FIND NODES or FIND VALUE
  383. // request to `addr`
  384. let lookup = async |node: H::Node, key, addrs: Vec<Url>| {
  385. let _permit = semaphore.acquire().await;
  386. // Try all valid addresses for the node
  387. let mut last_err = None;
  388. for addr in addrs {
  389. let mut queried_addrs_set = queried_addrs.lock().await;
  390. // Skip if this address has already been queried
  391. if queried_addrs_set.contains(&addr) {
  392. continue;
  393. }
  394. queried_addrs_set.insert(addr.clone());
  395. drop(queried_addrs_set);
  396. // Try to create or find an existing channel
  397. let channel = self.create_channel(&addr).await.map(|(ch, _)| ch);
  398. if let Err(e) = channel {
  399. last_err = Some(e);
  400. continue
  401. }
  402. let channel = channel.unwrap();
  403. let handler = self.handler().await;
  404. let res = match &lookup_type {
  405. DhtLookupType::Nodes => {
  406. info!(target: "dht::lookup", "[DHT] [LOOKUP] Querying node {} for nodes lookup of key {}", H::key_to_string(&node.id()), H::key_to_string(key));
  407. handler.find_nodes(channel.clone(), key).await.map(DhtLookupReply::Nodes)
  408. }
  409. DhtLookupType::Value => {
  410. info!(target: "dht::lookup", "[DHT] [LOOKUP] Querying node {} for value lookup of key {}", H::key_to_string(&node.id()), H::key_to_string(key));
  411. handler.find_value(channel.clone(), key).await
  412. }
  413. };
  414. self.cleanup_channel(channel).await;
  415. if res.is_ok() {
  416. return (node, res)
  417. }
  418. last_err = res.err();
  419. }
  420. if let Some(e) = last_err {
  421. return (node, Err(e))
  422. }
  423. (node, Err(Error::Custom("All node's addresses failed".to_string())))
  424. };
  425. // Spawn up to `alpha` futures for lookup()
  426. let spawn_futures = async |nodes_to_visit: &mut Vec<H::Node>,
  427. futures: &mut FuturesUnordered<_>| {
  428. for _ in 0..a {
  429. if !nodes_to_visit.is_empty() {
  430. let node = nodes_to_visit.remove(0);
  431. let valid_addrs: Vec<Url> = node
  432. .addresses()
  433. .iter()
  434. .filter(|addr| {
  435. active_profiles.contains(&addr.scheme().to_string()) &&
  436. !external_addrs.contains(addr)
  437. })
  438. .cloned()
  439. .collect();
  440. if !valid_addrs.is_empty() {
  441. futures.push(Box::pin(lookup(node, &key, valid_addrs)));
  442. }
  443. }
  444. }
  445. };
  446. // Initial futures
  447. spawn_futures(&mut nodes_to_visit, &mut futures).await;
  448. // Process lookup responses
  449. while let Some((queried_node, res)) = futures.next().await {
  450. if let Err(e) = res {
  451. warn!(target: "dht::lookup", "[DHT] [LOOKUP] Error in lookup: {e}");
  452. // Spawn next `alpha` futures if there are no more futures but
  453. // we still have nodes to visit
  454. if futures.is_empty() {
  455. spawn_futures(&mut nodes_to_visit, &mut futures).await;
  456. }
  457. continue;
  458. }
  459. let (nodes, value) = match res.unwrap() {
  460. DhtLookupReply::Nodes(nodes) => (Some(nodes), None),
  461. DhtLookupReply::Value(value) => (None, Some(value)),
  462. DhtLookupReply::NodesAndValue(nodes, value) => (Some(nodes), Some(value)),
  463. };
  464. // Send the value we found to the publisher
  465. if let Some(value) = value {
  466. info!(target: "dht::lookup", "[DHT] [LOOKUP] Found value for {} from {}", H::key_to_string(&key), H::key_to_string(&queried_node.id()));
  467. values.push(value.clone());
  468. self.event_publisher.notify(DhtEvent::ValueFound { key, value }).await;
  469. }
  470. // Update nodes_to_visit
  471. if let Some(mut nodes) = nodes {
  472. if !nodes.is_empty() {
  473. info!(target: "dht::lookup", "[DHT] [LOOKUP] Found {} nodes from {}", nodes.len(), H::key_to_string(&queried_node.id()));
  474. self.event_publisher
  475. .notify(DhtEvent::NodesFound { key, nodes: nodes.clone() })
  476. .await;
  477. // Remove our own node and duplicates
  478. if let Ok(self_node) = self.handler().await.node().await {
  479. let self_id = self_node.id();
  480. nodes.retain(|node: &H::Node| {
  481. node.id() != self_id && seen_nodes.insert(node.id())
  482. });
  483. }
  484. // Add new nodes to the list of nodes to visit
  485. nodes_to_visit.extend(nodes.clone());
  486. self.sort_by_distance(&mut nodes_to_visit, &key);
  487. }
  488. }
  489. result.push(queried_node);
  490. self.sort_by_distance(&mut result, &key);
  491. // Early termination logic:
  492. // The closest node to visit must be further than the furthest
  493. // queried node, 3 consecutive times
  494. if result.len() >= k &&
  495. result.last().zip(nodes_to_visit.first()).is_some_and(distance_check)
  496. {
  497. consecutive_stalls += 1;
  498. if consecutive_stalls >= 3 {
  499. break;
  500. }
  501. } else {
  502. consecutive_stalls = 0;
  503. }
  504. // Spawn next `alpha` futures
  505. spawn_futures(&mut nodes_to_visit, &mut futures).await;
  506. }
  507. info!(target: "dht::lookup", "[DHT] [LOOKUP] Lookup for {} completed", H::key_to_string(&key));
  508. let nodes: Vec<_> = result.into_iter().take(k).collect();
  509. (nodes, values)
  510. }
  511. /// Find `k` nodes closest to a key
  512. pub async fn lookup_nodes(&self, key: &blake3::Hash) -> Vec<H::Node> {
  513. info!(target: "dht::lookup_nodes", "[DHT] [LOOKUP] Starting node lookup for key {}", H::key_to_string(key));
  514. self.event_publisher.notify(DhtEvent::NodesLookupStarted { key: *key }).await;
  515. let (nodes, _) = self.lookup(*key, DhtLookupType::Nodes).await;
  516. self.event_publisher
  517. .notify(DhtEvent::NodesLookupCompleted { key: *key, nodes: nodes.clone() })
  518. .await;
  519. nodes
  520. }
  521. /// Find value for `key`
  522. pub async fn lookup_value(&self, key: &blake3::Hash) -> (Vec<H::Node>, Vec<H::Value>) {
  523. info!(target: "dht::lookup_value", "[DHT] [LOOKUP] Starting value lookup for key {}", H::key_to_string(key));
  524. self.event_publisher.notify(DhtEvent::ValueLookupStarted { key: *key }).await;
  525. let (nodes, values) = self.lookup(*key, DhtLookupType::Value).await;
  526. self.event_publisher
  527. .notify(DhtEvent::ValueLookupCompleted {
  528. key: *key,
  529. nodes: nodes.clone(),
  530. values: values.clone(),
  531. })
  532. .await;
  533. (nodes, values)
  534. }
  535. /// Update a channel's `last_used` field in the channel cache.
  536. pub async fn update_channel(&self, channel_id: u32) {
  537. let channel_cache_lock = self.channel_cache.clone();
  538. let mut channel_cache = channel_cache_lock.write().await;
  539. if let Some(cached) = channel_cache.get_mut(&channel_id) {
  540. cached.last_used = Timestamp::current_time();
  541. }
  542. }
  543. /// Get a channel (existing or create a new one) to `node`.
  544. /// Don't forget to call `cleanup_channel()` once you are done with it.
  545. pub async fn get_channel(&self, node: &H::Node) -> Result<(ChannelPtr, H::Node)> {
  546. let node_id = node.id();
  547. // Look in the channel cache for a channel connected to this node.
  548. // We skip direct session channels, for those we will call
  549. // `create_channel()` which increments the sessions's usage counter.
  550. let channel_cache = self.channel_cache.read().await.clone();
  551. if let Some((channel_id, cached)) = channel_cache
  552. .clone()
  553. .iter()
  554. .find(|(_, cached)| cached.node.clone().is_some_and(|n| n.id() == node_id))
  555. {
  556. if let Some(channel) = self.p2p.get_channel(*channel_id) {
  557. if channel.session_type_id() & SESSION_DIRECT == 0 {
  558. if channel.is_stopped() {
  559. self.cleanup_channel(channel).await;
  560. } else {
  561. return Ok((channel, cached.node.clone().unwrap()))
  562. }
  563. }
  564. }
  565. }
  566. self.create_channel_to_node(node).await
  567. }
  568. /// Create a channel in the direct session, ping the peer, add the
  569. /// DHT node to our buckets and the channel to our channel cache.
  570. pub async fn create_channel(&self, addr: &Url) -> Result<(ChannelPtr, H::Node)> {
  571. let external_addrs = self.p2p.hosts().external_addrs().await;
  572. if external_addrs.contains(addr) {
  573. return Err(Error::Custom(
  574. "Can't create a channel to our own external address".to_string(),
  575. ))
  576. }
  577. let channel = self.p2p.session_direct().get_channel(addr).await?;
  578. let channel_cache = self.channel_cache.read().await;
  579. if let Some(cached) = channel_cache.get(&channel.info.id) {
  580. if let Some(node) = &cached.node {
  581. return Ok((channel, node.clone()))
  582. }
  583. }
  584. drop(channel_cache);
  585. let node = self.ping(channel.clone()).await;
  586. // If ping failed, cleanup the channel and abort
  587. if let Err(e) = node {
  588. self.cleanup_channel(channel).await;
  589. return Err(e);
  590. }
  591. let node = node.unwrap();
  592. self.add_channel_to_cache(channel.info.id, &node).await;
  593. Ok((channel, node))
  594. }
  595. pub async fn create_channel_to_node(&self, node: &H::Node) -> Result<(ChannelPtr, H::Node)> {
  596. let net_settings = self.p2p.settings().read_arc().await;
  597. let active_profiles = net_settings.active_profiles.clone();
  598. drop(net_settings);
  599. // Create a channel
  600. let mut addrs = node.addresses().clone();
  601. addrs.retain(|addr| active_profiles.contains(&addr.scheme().to_string()));
  602. for addr in addrs {
  603. let res = self.create_channel(&addr).await;
  604. if res.is_err() {
  605. continue;
  606. }
  607. let (channel, node) = res.unwrap();
  608. return Ok((channel, node));
  609. }
  610. Err(Error::Custom("Could not create channel".to_string()))
  611. }
  612. /// Insert a channel to the DHT's channel cache. If the channel is already
  613. /// in the cache, `last_used` is updated.
  614. pub async fn add_channel_to_cache(&self, channel_id: u32, node: &H::Node) {
  615. let mut channel_cache = self.channel_cache.write().await;
  616. channel_cache
  617. .entry(channel_id)
  618. .and_modify(|c| c.last_used = Timestamp::current_time())
  619. .or_insert(ChannelCacheItem {
  620. node: Some(node.clone()),
  621. last_used: Timestamp::current_time(),
  622. ping_received: false,
  623. ping_sent: false,
  624. });
  625. }
  626. /// Wait until we received a DHT ping and sent a DHT ping on a channel.
  627. pub async fn wait_fully_pinged(&self, channel_id: u32) -> Result<()> {
  628. loop {
  629. let channel_cache = self.channel_cache.read().await;
  630. let cached = channel_cache
  631. .get(&channel_id)
  632. .ok_or(Error::Custom("Missing channel".to_string()))?;
  633. if cached.ping_received && cached.ping_sent {
  634. return Ok(())
  635. }
  636. drop(channel_cache);
  637. msleep(100).await;
  638. }
  639. }
  640. /// Call [`crate::net::session::DirectSession::cleanup_channel()`] and cleanup the DHT caches.
  641. pub async fn cleanup_channel(&self, channel: ChannelPtr) {
  642. let channel_cache_lock = self.channel_cache.clone();
  643. let mut channel_cache = channel_cache_lock.write().await;
  644. let mut ping_locks = self.ping_locks.lock().await;
  645. if self.p2p.session_direct().cleanup_channel(channel.clone()).await {
  646. channel_cache.remove(&channel.info.id);
  647. ping_locks.remove(&channel.info.id);
  648. }
  649. }
  650. }