p2p/network/kad/
p2p_network_kad_internals.rs

1use std::{
2    fmt::Debug,
3    ops::{Add, Shr, Sub},
4};
5
6use crypto_bigint::{ArrayEncoding, Encoding, U256};
7use derive_more::From;
8use libp2p_identity::DecodingError;
9use malloc_size_of_derive::MallocSizeOf;
10use multiaddr::Multiaddr;
11use openmina_core::bug_condition;
12use serde::{Deserialize, Serialize};
13use sha2::{Digest, Sha256};
14
15use crate::{
16    ConnectionType, P2pNetworkKademliaMultiaddrError, P2pNetworkKademliaPeerIdError, PeerId,
17};
18
19use super::CID;
20
21mod u256_serde {
22    use std::array::TryFromSliceError;
23
24    use crypto_bigint::{Encoding, U256};
25    use serde::{Deserialize, Serialize};
26
27    pub fn serialize<S>(value: &U256, serializer: S) -> Result<S::Ok, S::Error>
28    where
29        S: serde::Serializer,
30    {
31        if serializer.is_human_readable() {
32            let hex = hex::encode(value.to_be_bytes());
33            serializer.serialize_str(&hex)
34        } else {
35            value.serialize(serializer)
36        }
37    }
38
39    fn decode_error<E: serde::de::Error>(e: hex::FromHexError) -> E {
40        E::custom(format!("error converting from hex string: {e}"))
41    }
42
43    fn bytes_error<E: serde::de::Error>(e: TryFromSliceError) -> E {
44        E::custom(format!("error converting from slice to array: {e}"))
45    }
46
47    pub fn deserialize<'de, D>(deserializer: D) -> Result<U256, D::Error>
48    where
49        D: serde::Deserializer<'de>,
50    {
51        if deserializer.is_human_readable() {
52            let s = String::deserialize(deserializer)?;
53            let bytes = hex::decode(s)
54                .map_err(decode_error)?
55                .as_slice()
56                .try_into()
57                .map_err(bytes_error)?;
58            let u256 = U256::from_be_bytes(bytes);
59            Ok(u256)
60        } else {
61            U256::deserialize(deserializer)
62        }
63    }
64}
65
66/// Kademlia key, sha256 of the node's peer id.
67#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, MallocSizeOf)]
68pub struct P2pNetworkKadKey(
69    #[serde(with = "u256_serde")]
70    #[ignore_malloc_size_of = "doesn't allocate"]
71    U256,
72);
73
74impl P2pNetworkKadKey {
75    pub fn distance(self, rhs: &Self) -> P2pNetworkKadDist {
76        P2pNetworkKadDist(self.0 ^ rhs.0)
77    }
78}
79
80#[derive(Clone, Debug, Serialize, PartialEq, Deserialize, thiserror::Error, MallocSizeOf)]
81pub enum P2pNetworkKadKeyError {
82    #[error("decoding error")]
83    DecodingError,
84}
85
86impl TryFrom<&PeerId> for P2pNetworkKadKey {
87    type Error = P2pNetworkKadKeyError;
88
89    fn try_from(value: &PeerId) -> Result<Self, Self::Error> {
90        P2pNetworkKadKey::try_from(*value)
91    }
92}
93
94impl TryFrom<PeerId> for P2pNetworkKadKey {
95    type Error = P2pNetworkKadKeyError;
96
97    fn try_from(value: PeerId) -> Result<Self, Self::Error> {
98        Ok(P2pNetworkKadKey::from(
99            CID::try_from(value).map_err(|_| P2pNetworkKadKeyError::DecodingError)?,
100        ))
101    }
102}
103
104impl From<CID> for P2pNetworkKadKey {
105    fn from(value: CID) -> Self {
106        P2pNetworkKadKey(U256::from_be_byte_array(Sha256::digest(value.0)))
107    }
108}
109
110impl Debug for P2pNetworkKadKey {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        if f.alternate() {
113            let bytes = self.0.to_be_bytes();
114            f.write_str(&hex::encode(bytes))
115        } else {
116            f.debug_tuple("P2pNetworkKadKey").field(&self.0).finish()
117        }
118    }
119}
120
121impl Add<&P2pNetworkKadDist> for &P2pNetworkKadKey {
122    type Output = P2pNetworkKadKey;
123
124    #[allow(clippy::suspicious_arithmetic_impl, clippy::arithmetic_side_effects)]
125    fn add(self, rhs: &P2pNetworkKadDist) -> Self::Output {
126        P2pNetworkKadKey(self.0 ^ rhs.0)
127    }
128}
129
130impl Sub for P2pNetworkKadKey {
131    type Output = P2pNetworkKadDist;
132
133    #[allow(clippy::suspicious_arithmetic_impl, clippy::arithmetic_side_effects)]
134    fn sub(self, rhs: Self) -> Self::Output {
135        self.distance(&rhs)
136    }
137}
138
139impl Sub<&P2pNetworkKadKey> for &P2pNetworkKadKey {
140    type Output = P2pNetworkKadDist;
141
142    #[allow(clippy::suspicious_arithmetic_impl, clippy::arithmetic_side_effects)]
143    fn sub(self, rhs: &P2pNetworkKadKey) -> Self::Output {
144        self.distance(rhs)
145    }
146}
147
148impl Sub<P2pNetworkKadKey> for &P2pNetworkKadKey {
149    type Output = P2pNetworkKadDist;
150
151    #[allow(clippy::suspicious_arithmetic_impl, clippy::arithmetic_side_effects)]
152    fn sub(self, rhs: P2pNetworkKadKey) -> Self::Output {
153        self.distance(&rhs)
154    }
155}
156
157impl Sub<&P2pNetworkKadKey> for P2pNetworkKadKey {
158    type Output = P2pNetworkKadDist;
159
160    #[allow(clippy::suspicious_arithmetic_impl, clippy::arithmetic_side_effects)]
161    fn sub(self, rhs: &P2pNetworkKadKey) -> Self::Output {
162        self.distance(rhs)
163    }
164}
165
166impl Shr<usize> for &P2pNetworkKadKey {
167    type Output = P2pNetworkKadKey;
168
169    fn shr(self, rhs: usize) -> Self::Output {
170        P2pNetworkKadKey(self.0 >> rhs)
171    }
172}
173
174/// Kademlia distance between two nodes, calculated as `XOR` of their keys.
175#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
176pub struct P2pNetworkKadDist(#[serde(with = "u256_serde")] U256);
177
178impl Debug for P2pNetworkKadDist {
179    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
180        if f.alternate() {
181            let bytes = self.0.to_be_bytes();
182            f.write_str(&hex::encode(bytes))
183        } else {
184            f.debug_tuple("P2pNetworkKadDist").field(&self.0).finish()
185        }
186    }
187}
188
189impl P2pNetworkKadDist {
190    /// Returns (maximal possible) index of the bucket for this distance. In
191    /// other words, this is the length of the common prefix of two nodes
192    /// withing this distance from each other.
193    fn to_index(&self) -> usize {
194        256 - self.0.bits_vartime()
195    }
196}
197
198/// Converts a K-bucket index to distance as `2^n - 1`. For `i`, any node with its keys
199/// having at least `i` highest bits in common are within this distance.
200impl From<usize> for P2pNetworkKadDist {
201    fn from(value: usize) -> Self {
202        P2pNetworkKadDist(U256::MAX >> value)
203    }
204}
205
206/// Kademlia routing table, with `K` parameter, the maximum number of records
207/// for each bucket. Usually it is set to `20`.
208#[derive(Clone, Serialize, Deserialize)]
209pub struct P2pNetworkKadRoutingTable<const K: usize = 20> {
210    /// SHA256 of the current node's id.
211    pub this_key: P2pNetworkKadKey,
212    /// Kademlia K-buckets. Under index `i` located elements within distance
213    /// `2^(256-i)` from the current node at most. If there is also `i+1` (i.e.
214    /// `i` is not the last index), then distance from the current node to the
215    /// elements under `i` are greater than `2^(256-i-1)`.
216    pub buckets: Vec<P2pNetworkKadBucket<K>>,
217}
218
219impl<const K: usize> Default for P2pNetworkKadRoutingTable<K> {
220    fn default() -> Self {
221        Self {
222            this_key: P2pNetworkKadKey(U256::ZERO),
223            buckets: Default::default(),
224        }
225    }
226}
227
228impl<const K: usize> Debug for P2pNetworkKadRoutingTable<K> {
229    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
230        if f.alternate() {
231            if f.sign_plus() {
232                writeln!(f, "this key: {:#?}", self.this_key)?;
233            }
234            for (i, bucket) in self.buckets.iter().enumerate() {
235                writeln!(f, "{i}: {}", bucket.0.len())?;
236                if f.sign_plus() {
237                    for entry in &bucket.0 {
238                        writeln!(f, "    {:#.*?}", i, &entry.key)?;
239                    }
240                }
241            }
242            Ok(())
243        } else {
244            f.debug_struct("P2pNetworkKadRoutingTable")
245                .field("this_key", &self.this_key)
246                .field("buckets", &self.buckets)
247                .finish()
248        }
249    }
250}
251
252#[derive(Debug, thiserror::Error)]
253#[error("bucket capacity exceeded")]
254pub struct P2pNetworkKadRoutingTableInsertError;
255
256impl<const K: usize> P2pNetworkKadRoutingTable<K> {
257    pub fn new(this_entry: P2pNetworkKadEntry) -> Self {
258        let this_key = this_entry.key;
259        let global_bucket = P2pNetworkKadBucket(vec![this_entry]);
260        let buckets = vec![global_bucket];
261        P2pNetworkKadRoutingTable { this_key, buckets }
262    }
263
264    /// Inserts the `entry` into the routing table.
265    ///
266    /// Returns `Ok(true)` it is added as a new entry, and `Ok(false)` if an
267    /// existing one has been updated.
268    ///
269    /// If there is no space for adding the new entry (i.e. corresponding
270    /// K-bucket is full and cannot be split), then `Err(_)` value is returned.
271    ///
272    /// TODO: should it be just another variant in `Ok(_)`?
273    ///
274    /// Filters `entry.addrs` for supported addresses if non of addresses are supported returns `Err(_)`
275    pub fn insert(
276        &mut self,
277        entry: P2pNetworkKadEntry,
278    ) -> Result<bool, P2pNetworkKadRoutingTableInsertError> {
279        if entry.addrs.is_empty() {
280            return Err(P2pNetworkKadRoutingTableInsertError);
281        }
282        // distance to this node
283        let dist = self.this_key - entry.key;
284
285        // index of the closest k-bucket that can contain this node.
286        let index = dist.to_index();
287
288        let mut max_index = self.buckets.len() - 1;
289        loop {
290            if index < max_index {
291                // bucket cannot be split
292                if self.buckets[index].can_insert(&entry) {
293                    break Ok(self.buckets[index].insert(entry));
294                } else {
295                    break Err(P2pNetworkKadRoutingTableInsertError);
296                }
297            } else if self.buckets[max_index].can_insert(&entry) {
298                break Ok(self.buckets[max_index].insert(entry));
299            } else {
300                max_index += 1;
301                let split_dist = max_index.into();
302                let Some((bucket1, bucket2)) = self
303                    .buckets
304                    .pop()
305                    .map(|b| b.split(|e| (self.this_key - e.key) >= split_dist))
306                else {
307                    bug_condition!("should be unreachable");
308                    return Err(P2pNetworkKadRoutingTableInsertError);
309                };
310                self.buckets.extend([bucket1, bucket2]);
311            }
312        }
313    }
314
315    /// Looks up a Kademlia entry with the specified `key`.
316    pub fn look_up(&self, key: &P2pNetworkKadKey) -> Option<&P2pNetworkKadEntry> {
317        // distance to this node
318        let dist = self.this_key - key;
319
320        // index of the closest k-bucket that can contain this node.
321        let index = dist.to_index().min(self.buckets.len() - 1);
322
323        self.buckets[index].iter().find(|e| &e.key == key)
324    }
325
326    /// FIND_NODE backend. Returns iterator of nodes closest to the specified
327    /// `key`, excluding nodes that correspond to the `key` itself and
328    /// `self.this_key`.
329    ///
330    /// The number of entries is limited to 20.
331    pub fn find_node<'a>(
332        &'a self,
333        key: &'a P2pNetworkKadKey,
334    ) -> impl Iterator<Item = &'a P2pNetworkKadEntry> {
335        self.closest_peers(key).take(20)
336    }
337
338    /// Returns iterator of nodes closest to the current one.
339    /// TODO: use reverse order over bucket, from the deepest bucket.
340    pub fn closest_peers<'a>(&'a self, key: &'a P2pNetworkKadKey) -> ClosestPeers<'a, K> {
341        ClosestPeers::new(self, key)
342    }
343
344    #[cfg(test)]
345    fn assert_k_buckets(&self) {
346        let mut prev_dist = None;
347        for (i, bucket) in self.buckets.iter().enumerate().rev() {
348            assert!(bucket.0.len() <= K, "{self:+#?}");
349            let dist = P2pNetworkKadDist::from(i);
350            for entry in &bucket.0 {
351                assert!(
352                    self.this_key - entry.key <= dist,
353                    "for {:#?} at {i} distance {:#?} is too big, expecting at most {dist:#?}\nrouting table:\n{self:+#?}",
354                    entry.key,
355                    self.this_key - entry.key,
356                );
357                if let Some(prev_dist) = &prev_dist {
358                    assert!(
359                        &(self.this_key - entry.key) > prev_dist,
360                        "distance too small: {:#?}\nrouting table:\n{:+#?}\ndist: {:#?}\nprev_dist: {:#?}",
361                        entry.key,
362                        self,
363                        self.this_key - entry.key,
364                        prev_dist,
365                    );
366                }
367            }
368            prev_dist = Some(dist);
369        }
370    }
371}
372
373impl Extend<P2pNetworkKadEntry> for P2pNetworkKadRoutingTable {
374    fn extend<T: IntoIterator<Item = P2pNetworkKadEntry>>(&mut self, iter: T) {
375        for entry in iter {
376            // TODO(akoptelov): log addition status?
377            let _ = self.insert(entry);
378        }
379    }
380}
381
382#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, MallocSizeOf)]
383pub struct P2pNetworkKadEntry {
384    pub key: P2pNetworkKadKey,
385    pub peer_id: PeerId,
386    #[with_malloc_size_of_func = "measurement::multiaddr_vec"]
387    addrs: Vec<Multiaddr>,
388    #[ignore_malloc_size_of = "doesn't allocate"]
389    pub connection: ConnectionType,
390}
391
392impl P2pNetworkKadEntry {
393    pub const MAX_ADDRS: usize = 16;
394
395    pub fn new(peer_id: PeerId, addrs: Vec<Multiaddr>) -> Result<Self, P2pNetworkKadKeyError> {
396        if addrs.len() > Self::MAX_ADDRS {
397            openmina_core::log::info!(
398                openmina_core::log::system_time();
399                kind = "P2pNetworkKadEntry new",
400                summary = format!("truncating {addrs:?} to {} elements", Self::MAX_ADDRS),
401            );
402        }
403
404        Ok(P2pNetworkKadEntry {
405            key: peer_id.try_into()?,
406            peer_id,
407            addrs: addrs.into_iter().take(Self::MAX_ADDRS).collect(),
408            connection: ConnectionType::NotConnected,
409        })
410    }
411
412    pub fn dist(&self, other: &P2pNetworkKadEntry) -> P2pNetworkKadDist {
413        self.key - other.key
414    }
415
416    pub fn addresses(&self) -> &Vec<Multiaddr> {
417        &self.addrs
418    }
419}
420
421#[derive(Clone, Debug, Serialize, PartialEq, Deserialize, thiserror::Error, MallocSizeOf)]
422pub enum P2pNetworkKadEntryTryFromError {
423    #[error(transparent)]
424    PeerId(#[from] P2pNetworkKademliaPeerIdError),
425    #[error(transparent)]
426    Key(#[from] P2pNetworkKadKeyError),
427    #[error(transparent)]
428    Multiaddr(#[from] P2pNetworkKademliaMultiaddrError),
429}
430
431impl TryFrom<super::mod_Message::Peer<'_>> for P2pNetworkKadEntry {
432    type Error = P2pNetworkKadEntryTryFromError;
433
434    fn try_from(value: super::mod_Message::Peer) -> Result<Self, Self::Error> {
435        let peer_id = super::peer_id_try_from_bytes(value.id)?;
436        let key = peer_id.try_into()?;
437        let addrs = value
438            .addrs
439            .into_iter()
440            .map(super::multiaddr_try_from_bytes)
441            .collect::<Result<_, _>>()?;
442        let connection = value.connection.into();
443        Ok(P2pNetworkKadEntry {
444            peer_id,
445            key,
446            addrs,
447            connection,
448        })
449    }
450}
451
452impl<'a> TryFrom<&'a P2pNetworkKadEntry> for super::mod_Message::Peer<'a> {
453    type Error = DecodingError;
454
455    fn try_from(value: &'a P2pNetworkKadEntry) -> Result<Self, Self::Error> {
456        Ok(super::mod_Message::Peer {
457            id: (&value.peer_id).try_into()?,
458            addrs: value
459                .addrs
460                .iter()
461                .map(|addr| addr.as_ref().into())
462                .collect(),
463            connection: value.connection.into(),
464        })
465    }
466}
467
468pub struct ClosestPeers<'a, const K: usize> {
469    table: &'a P2pNetworkKadRoutingTable<K>,
470    index_iter: std::vec::IntoIter<usize>,
471    bucket_index: usize,
472    bucket_iterator: std::vec::IntoIter<&'a P2pNetworkKadEntry>,
473    key: &'a P2pNetworkKadKey,
474}
475
476impl<'a, const K: usize> ClosestPeers<'a, K> {
477    fn new(table: &'a P2pNetworkKadRoutingTable<K>, key: &'a P2pNetworkKadKey) -> Self {
478        let dist = table.this_key - key;
479        let mut index_iter = Self::bucket_index_iterator(dist, table.buckets.len());
480        let bucket_index = index_iter
481            .next()
482            .expect("implementation should ensure there is at least one bucket");
483        let bucket_iterator =
484            Self::get_bucket_iter(&table.buckets[bucket_index], key, &table.this_key);
485        // println!(">>> first bucket {}", bucket_index);
486        ClosestPeers {
487            table,
488            index_iter,
489            bucket_index,
490            bucket_iterator,
491            key,
492        }
493    }
494
495    fn bucket_index_iterator(
496        dist: P2pNetworkKadDist,
497        buckets_len: usize,
498    ) -> std::vec::IntoIter<usize> {
499        let (mut ones, zeroes) =
500            (0..buckets_len).partition::<Vec<_>, _>(|index| dist.0.bit_vartime(255 - *index));
501        ones.extend(zeroes.into_iter().rev());
502        let it: std::vec::IntoIter<usize> = ones.into_iter();
503        it
504    }
505
506    fn get_bucket_iter(
507        bucket: &'a P2pNetworkKadBucket<K>,
508        key: &P2pNetworkKadKey,
509        this_key: &P2pNetworkKadKey,
510    ) -> std::vec::IntoIter<&'a P2pNetworkKadEntry> {
511        let mut vec = Vec::from_iter(
512            bucket
513                .into_iter()
514                .filter(|e| &e.key != key && &e.key != this_key),
515        );
516        vec.sort_by_cached_key(|entry| key - entry.key);
517        vec.into_iter()
518    }
519}
520
521impl<'a, const K: usize> Iterator for ClosestPeers<'a, K> {
522    type Item = &'a P2pNetworkKadEntry;
523
524    fn next(&mut self) -> Option<Self::Item> {
525        Some(loop {
526            if let Some(item) = self.bucket_iterator.next() {
527                break item;
528            }
529            self.bucket_index = self.index_iter.next()?;
530            self.bucket_iterator = Self::get_bucket_iter(
531                &self.table.buckets[self.bucket_index],
532                self.key,
533                &self.table.this_key,
534            );
535        })
536    }
537}
538
539#[derive(Clone, Debug, Default, Serialize, Deserialize, From, MallocSizeOf)]
540pub struct P2pNetworkKadBucket<const K: usize>(Vec<P2pNetworkKadEntry>);
541
542impl<const K: usize> P2pNetworkKadBucket<K> {
543    pub fn len(&self) -> usize {
544        self.0.len()
545    }
546
547    pub fn is_empty(&self) -> bool {
548        self.0.is_empty()
549    }
550
551    pub fn iter(&self) -> std::slice::Iter<'_, P2pNetworkKadEntry> {
552        self.0.iter()
553    }
554
555    /// Checks if the `entry` can be inserted into this bucket, that is, either
556    /// there is a free slot, or the entry with this peer_id already there and
557    /// only needs to be updated.
558    fn can_insert(&self, entry: &P2pNetworkKadEntry) -> bool {
559        self.len() < K || self.0.iter().any(|e| e.key == entry.key)
560    }
561
562    /// Inserts an entry into the bucket. Returns true if a new entry is added,
563    /// false if an existing one is updated.
564    fn insert(&mut self, entry: P2pNetworkKadEntry) -> bool {
565        if let Some(pos) = self.0.iter().position(|e| e.key == entry.key) {
566            let e = &mut self.0[pos];
567
568            if e.peer_id != entry.peer_id {
569                bug_condition!(
570                    "Kad entry peer_id mismatch {:?} != {:?}",
571                    e.peer_id,
572                    entry.peer_id
573                );
574            }
575
576            let addrs = entry
577                .addrs
578                .into_iter()
579                .filter(|addr| !e.addrs.contains(addr))
580                .collect::<Vec<_>>();
581
582            for addr in addrs {
583                if e.addrs.len() >= P2pNetworkKadEntry::MAX_ADDRS {
584                    openmina_core::warn!(
585                        openmina_core::log::system_time();
586                        kind = "P2pNetworkKadBucket insert",
587                        peer_id = e.peer_id.to_string(),
588                        summary = format!("Skipping updates to Kad entry multiaddress list"),
589                    );
590                    break;
591                }
592
593                if !e.addrs.contains(&addr) {
594                    e.addrs.push(addr);
595                }
596            }
597            false
598        } else {
599            if self.len() >= K {
600                bug_condition!("Kad bucket len {:?} >= K ({:?})", self.len(), K);
601            }
602            self.0.push(entry);
603            true
604        }
605    }
606
607    /// Splits this bucket into two, keeping entries that are not closer to the
608    /// current node than the `dist`.
609    fn split<F: Fn(&P2pNetworkKadEntry) -> bool>(self, f: F) -> (Self, Self) {
610        self.into_iter().partition(f)
611    }
612}
613
614impl<const K: usize> FromIterator<P2pNetworkKadEntry> for P2pNetworkKadBucket<K> {
615    fn from_iter<T: IntoIterator<Item = P2pNetworkKadEntry>>(iter: T) -> Self {
616        P2pNetworkKadBucket(Vec::from_iter(iter))
617    }
618}
619
620impl<const K: usize> IntoIterator for P2pNetworkKadBucket<K> {
621    type Item = P2pNetworkKadEntry;
622
623    type IntoIter = <Vec<P2pNetworkKadEntry> as IntoIterator>::IntoIter;
624
625    fn into_iter(self) -> Self::IntoIter {
626        self.0.into_iter()
627    }
628}
629
630impl<const K: usize> Extend<P2pNetworkKadEntry> for P2pNetworkKadBucket<K> {
631    fn extend<T: IntoIterator<Item = P2pNetworkKadEntry>>(&mut self, iter: T) {
632        self.0.extend(iter)
633    }
634}
635
636impl<'a, const K: usize> IntoIterator for &'a P2pNetworkKadBucket<K> {
637    type Item = &'a P2pNetworkKadEntry;
638
639    type IntoIter = std::slice::Iter<'a, P2pNetworkKadEntry>;
640
641    fn into_iter(self) -> Self::IntoIter {
642        self.0.as_slice().iter()
643    }
644}
645
646#[cfg(test)]
647mod tests {
648    use std::collections::BTreeSet;
649
650    use crypto_bigint::{Random, U256};
651    use multiaddr::multiaddr;
652
653    use crate::{identity::SecretKey, PeerId, CID};
654
655    use super::{P2pNetworkKadEntry, P2pNetworkKadKey, P2pNetworkKadRoutingTable};
656
657    const THIS_KEY: P2pNetworkKadKey = P2pNetworkKadKey(U256::ZERO);
658
659    fn this_key() -> P2pNetworkKadKey {
660        THIS_KEY
661    }
662
663    fn key_pow_2(pow: usize) -> P2pNetworkKadKey {
664        P2pNetworkKadKey(U256::ONE.shl_vartime(pow))
665    }
666
667    fn key_rand() -> P2pNetworkKadKey {
668        P2pNetworkKadKey(U256::random(&mut rand::thread_rng()))
669    }
670
671    fn peer_id_rand() -> PeerId {
672        crate::identity::SecretKey::rand().public_key().peer_id()
673    }
674
675    fn entry(key: P2pNetworkKadKey) -> P2pNetworkKadEntry {
676        let peer_id = peer_id_rand();
677        P2pNetworkKadEntry {
678            key,
679            peer_id,
680            addrs: vec![multiaddr!(Ip4([0; 4]), Tcp(1000_u16))],
681            connection: super::ConnectionType::Connected,
682        }
683    }
684
685    fn entry_with_peer_id(peer_id: PeerId) -> P2pNetworkKadEntry {
686        let key = peer_id.try_into().expect("Error converting PeerId");
687        P2pNetworkKadEntry {
688            key,
689            peer_id,
690            addrs: vec![multiaddr!(Ip4([0; 4]), Tcp(1000_u16))],
691            connection: super::ConnectionType::Connected,
692        }
693    }
694
695    #[test]
696    fn test_key_generation() {
697        let random_peer_id = SecretKey::rand().public_key().peer_id();
698        let libp2p_peer_id =
699            libp2p_identity::PeerId::try_from(random_peer_id).expect("Conversion failed");
700        let cid = CID::from(libp2p_peer_id);
701
702        let key0 = P2pNetworkKadKey::try_from(&random_peer_id).expect("Conversion failed");
703        let key1 = P2pNetworkKadKey::try_from(random_peer_id).expect("Conversion failed");
704        let key2 = P2pNetworkKadKey::from(cid);
705
706        assert_eq!(key0, key1);
707        assert_eq!(key1, key2);
708    }
709
710    #[test]
711    fn test_256_keys() {
712        let mut rt: P2pNetworkKadRoutingTable = P2pNetworkKadRoutingTable::new(entry(this_key()));
713
714        for i in 0..255 {
715            println!("=== adding {i}...");
716            let key = key_pow_2(255 - i);
717            let entry = entry(key);
718            let _ = rt.insert(entry);
719            rt.assert_k_buckets();
720        }
721        println!("routing table: {rt:+#?}");
722    }
723
724    #[test]
725    fn test_256_keys_rev() {
726        let mut rt: P2pNetworkKadRoutingTable = P2pNetworkKadRoutingTable::new(entry(this_key()));
727
728        for i in 0..255 {
729            println!("=== adding {i}...");
730            let key = key_pow_2(i);
731            let entry = entry(key);
732            let _ = rt.insert(entry);
733            rt.assert_k_buckets();
734        }
735        println!("routing table: {rt:+#?}");
736    }
737
738    #[test]
739    fn test_rand_keys() {
740        let mut rt: P2pNetworkKadRoutingTable = P2pNetworkKadRoutingTable::new(entry(this_key()));
741        for _ in 0..(256 * 256) {
742            let key = key_rand();
743            let entry = entry(key);
744            let _ = rt.insert(entry);
745            rt.assert_k_buckets();
746        }
747        println!("routing table: {rt:+#?}");
748    }
749
750    #[test]
751    fn test_rand_peers_rand_this() {
752        let mut rt: P2pNetworkKadRoutingTable =
753            P2pNetworkKadRoutingTable::new(entry_with_peer_id(peer_id_rand()));
754        for _ in 0..(256 * 256) {
755            let peer_id = peer_id_rand();
756            let entry = entry_with_peer_id(peer_id);
757            let _ = rt.insert(entry);
758            rt.assert_k_buckets();
759        }
760        println!("routing table: {rt:+#?}");
761    }
762
763    #[test]
764    fn test_rand_peers() {
765        let mut rt: P2pNetworkKadRoutingTable = P2pNetworkKadRoutingTable::new(entry(this_key()));
766        for _ in 0..(256 * 256) {
767            let peer_id = peer_id_rand();
768            let entry = entry_with_peer_id(peer_id);
769            let _ = rt.insert(entry);
770            rt.assert_k_buckets();
771        }
772        println!("routing table: {rt:+#?}");
773    }
774
775    #[test]
776    fn test_find_node_zero() {
777        let this_entry = entry_with_peer_id(peer_id_rand());
778        let mut rt: P2pNetworkKadRoutingTable = P2pNetworkKadRoutingTable::new(this_entry.clone());
779        for _ in 0..(256 * 32) {
780            let peer_id = peer_id_rand();
781            let entry = entry_with_peer_id(peer_id);
782            let _ = rt.insert(entry);
783            rt.assert_k_buckets();
784        }
785
786        let entry = entry(this_key());
787        let iter = rt.find_node(&entry.key);
788
789        // let dist = entry.dist(&this_entry);
790        // let index = dist.to_index().min(rt.buckets.len());
791
792        // let iter = rt.buckets[index..]
793        //     .iter()
794        //     .flat_map(|bucket| bucket.iter())
795        //     .filter(|e| &e.key != &entry.key)
796        //     .take(20);
797        let closest = BTreeSet::from_iter(iter);
798        println!("{}", closest.len());
799
800        let max_closest_dist = closest
801            .iter()
802            .max_by_key(|e| entry.dist(e))
803            .expect("Failed to find entry");
804        let min_non_closest_dist = rt
805            .buckets
806            .iter()
807            .flat_map(|e| e.iter())
808            .filter(|e| !closest.contains(*e))
809            .min_by_key(|e| entry.dist(e))
810            .expect("Failed to find entry");
811
812        let max = entry.dist(max_closest_dist);
813        let min = entry.dist(min_non_closest_dist);
814        println!(
815            "farthest {:#?} is closer than the closest {:#?}",
816            max_closest_dist.key, min_non_closest_dist.key
817        );
818        assert!(min > max);
819    }
820
821    #[test]
822    fn test_find_node_rand() {
823        let mut rt: P2pNetworkKadRoutingTable = P2pNetworkKadRoutingTable::new(entry(this_key()));
824        for _ in 0..(256 * 32) {
825            let peer_id = peer_id_rand();
826            let entry = entry_with_peer_id(peer_id);
827            let _ = rt.insert(entry);
828            rt.assert_k_buckets();
829        }
830
831        for _ in 0..(1024 * 16) {
832            let peer_id = peer_id_rand();
833            let entry = entry_with_peer_id(peer_id);
834
835            let find_node = rt.find_node(&entry.key);
836            let closest = BTreeSet::from_iter(find_node);
837
838            let max_closest_dist = closest
839                .iter()
840                .max_by_key(|e| entry.dist(e))
841                .expect("Error finding entry");
842            let min_non_closest_dist = rt
843                .buckets
844                .iter()
845                .flat_map(|e| e.iter())
846                .filter(|e| e.key != entry.key && e.key != rt.this_key)
847                .filter(|e| !closest.contains(*e))
848                .min_by_key(|e| entry.dist(e))
849                .expect("Error finding entry");
850
851            let max = entry.dist(max_closest_dist);
852            let min = entry.dist(min_non_closest_dist);
853            if max > min {
854                println!(
855                    "farthest {:#?} should be closer than the closest {:#?}",
856                    max_closest_dist.key, min_non_closest_dist.key
857                );
858                panic!("min is {min:#?}\nmax is {max:#?}");
859            }
860        }
861    }
862
863    /// Tests that `find_node` returns entries in order of increasing distance.
864    #[test]
865    fn test_closest_peers_rand() {
866        let mut rt: P2pNetworkKadRoutingTable = P2pNetworkKadRoutingTable::new(entry(this_key()));
867        for _ in 0..(256 * 32) {
868            let peer_id = peer_id_rand();
869            let entry = entry_with_peer_id(peer_id);
870            let _ = rt.insert(entry);
871            rt.assert_k_buckets();
872        }
873
874        for _ in 0..(16 * 1024) {
875            let peer_id = peer_id_rand();
876            let entry = entry_with_peer_id(peer_id);
877
878            let mut prev = None;
879            for e in rt.find_node(&entry.key) {
880                let dist = entry.dist(e);
881                if let Some((prev, prev_dist)) = prev {
882                    if dist <= prev_dist {
883                        panic!(
884                            "incorrect order\n{:#?}\nshould go before\n{:#?}",
885                            e.key, prev
886                        );
887                    }
888                }
889                prev = Some((e.clone(), dist));
890            }
891        }
892    }
893}
894mod measurement {
895    use std::mem;
896
897    use malloc_size_of::{MallocSizeOf, MallocSizeOfOps};
898
899    use super::{Multiaddr, P2pNetworkKadBucket, P2pNetworkKadRoutingTable};
900
901    pub fn multiaddr_vec(v: &Vec<Multiaddr>, _ops: &mut MallocSizeOfOps) -> usize {
902        v.capacity() * mem::size_of::<Multiaddr>() + v.iter().map(Multiaddr::len).sum::<usize>()
903    }
904
905    impl<const K: usize> MallocSizeOf for P2pNetworkKadRoutingTable<K> {
906        fn size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
907            self.this_key.size_of(ops)
908                + self.buckets.capacity() * mem::size_of::<P2pNetworkKadBucket<K>>()
909                + self.buckets.iter().map(|b| b.size_of(ops)).sum::<usize>()
910        }
911    }
912}
p2p/network/kad/p2p_network_kad_internals.rs

p2p/network/kad/
p2p_network_kad_internals.rs
