litep2p/protocol/libp2p/kademlia/
bucket.rs

1// Copyright 2018-2019 Parity Technologies (UK) Ltd.
2// Copyright 2023 litep2p developers
3//
4// Permission is hereby granted, free of charge, to any person obtaining a
5// copy of this software and associated documentation files (the "Software"),
6// to deal in the Software without restriction, including without limitation
7// the rights to use, copy, modify, merge, publish, distribute, sublicense,
8// and/or sell copies of the Software, and to permit persons to whom the
9// Software is furnished to do so, subject to the following conditions:
10//
11// The above copyright notice and this permission notice shall be included in
12// all copies or substantial portions of the Software.
13//
14// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
15// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
16// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
17// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
18// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
19// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
20// DEALINGS IN THE SOFTWARE.
21
22//! Kademlia k-bucket implementation.
23
24use crate::{
25    protocol::libp2p::kademlia::types::{ConnectionType, KademliaPeer, Key},
26    PeerId,
27};
28
29/// K-bucket entry.
30#[derive(Debug, PartialEq, Eq)]
31pub enum KBucketEntry<'a> {
32    /// Entry points to local node.
33    LocalNode,
34
35    /// Occupied entry to a connected node.
36    Occupied(&'a mut KademliaPeer),
37
38    /// Vacant entry.
39    Vacant(&'a mut KademliaPeer),
40
41    /// Entry not found and any present entry cannot be replaced.
42    NoSlot,
43}
44
45impl<'a> KBucketEntry<'a> {
46    /// Insert new entry into the entry if possible.
47    pub fn insert(&'a mut self, new: KademliaPeer) {
48        if let KBucketEntry::Vacant(old) = self {
49            old.peer = new.peer;
50            old.key = Key::from(new.peer);
51            old.addresses = new.addresses;
52            old.connection = new.connection;
53        }
54    }
55}
56
57/// Kademlia k-bucket.
58pub struct KBucket {
59    // TODO: store peers in a btreemap with increasing distance from local key?
60    nodes: Vec<KademliaPeer>,
61}
62
63impl KBucket {
64    /// Create new [`KBucket`].
65    pub fn new() -> Self {
66        Self {
67            nodes: Vec::with_capacity(20),
68        }
69    }
70
71    /// Get entry into the bucket.
72    // TODO: this is horrible code
73    pub fn entry<K: Clone>(&mut self, key: Key<K>) -> KBucketEntry<'_> {
74        for i in 0..self.nodes.len() {
75            if self.nodes[i].key == key {
76                return KBucketEntry::Occupied(&mut self.nodes[i]);
77            }
78        }
79
80        if self.nodes.len() < 20 {
81            self.nodes.push(KademliaPeer::new(
82                PeerId::random(),
83                vec![],
84                ConnectionType::NotConnected,
85            ));
86            let len = self.nodes.len() - 1;
87            return KBucketEntry::Vacant(&mut self.nodes[len]);
88        }
89
90        for i in 0..self.nodes.len() {
91            match self.nodes[i].connection {
92                ConnectionType::NotConnected | ConnectionType::CannotConnect => {
93                    return KBucketEntry::Vacant(&mut self.nodes[i]);
94                }
95                _ => continue,
96            }
97        }
98
99        KBucketEntry::NoSlot
100    }
101
102    /// Get iterator over the k-bucket, sorting the k-bucket entries in increasing order
103    /// by distance.
104    pub fn closest_iter<K: Clone>(&self, target: &Key<K>) -> impl Iterator<Item = KademliaPeer> {
105        let mut nodes = self.nodes.clone();
106        nodes.sort_by(|a, b| target.distance(&a.key).cmp(&target.distance(&b.key)));
107        nodes.into_iter().filter(|peer| !peer.addresses.is_empty())
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114
115    #[test]
116    fn closest_iter() {
117        let mut bucket = KBucket::new();
118
119        // add some random nodes to the bucket
120        let _ = (0..10)
121            .map(|_| {
122                let peer = PeerId::random();
123                bucket.nodes.push(KademliaPeer::new(peer, vec![], ConnectionType::Connected));
124
125                peer
126            })
127            .collect::<Vec<_>>();
128
129        let target = Key::from(PeerId::random());
130        let mut iter = bucket.closest_iter(&target);
131        let mut prev = None;
132
133        while let Some(node) = iter.next() {
134            if let Some(distance) = prev {
135                assert!(distance < target.distance(&node.key));
136            }
137
138            prev = Some(target.distance(&node.key));
139        }
140    }
141
142    #[test]
143    fn ignore_peers_with_no_addresses() {
144        let mut bucket = KBucket::new();
145
146        // add peers with no addresses to the bucket
147        let _ = (0..10)
148            .map(|_| {
149                let peer = PeerId::random();
150                bucket.nodes.push(KademliaPeer::new(
151                    peer,
152                    vec![],
153                    ConnectionType::NotConnected,
154                ));
155
156                peer
157            })
158            .collect::<Vec<_>>();
159
160        // add three peers with an address
161        let _ = (0..3)
162            .map(|_| {
163                let peer = PeerId::random();
164                bucket.nodes.push(KademliaPeer::new(
165                    peer,
166                    vec!["/ip6/::/tcp/0".parse().unwrap()],
167                    ConnectionType::Connected,
168                ));
169
170                peer
171            })
172            .collect::<Vec<_>>();
173
174        let target = Key::from(PeerId::random());
175        let mut iter = bucket.closest_iter(&target);
176        let mut prev = None;
177        let mut num_peers = 0usize;
178
179        while let Some(node) = iter.next() {
180            if let Some(distance) = prev {
181                assert!(distance < target.distance(&node.key));
182            }
183
184            num_peers += 1;
185            prev = Some(target.distance(&node.key));
186        }
187
188        assert_eq!(num_peers, 3usize);
189    }
190}
litep2p/protocol/libp2p/kademlia/bucket.rs

litep2p/protocol/libp2p/kademlia/
bucket.rs
