diff --git a/src/common/routing_table.rs b/src/common/routing_table.rs index c2db290..dd7d6ae 100644 --- a/src/common/routing_table.rs +++ b/src/common/routing_table.rs @@ -3,7 +3,8 @@ use std::slice::Iter; use std::{collections::BTreeMap, rc::Rc}; -use crate::common::{Id, Node, MAX_DISTANCE}; +use crate::common::{Id, Node}; +use crate::rpc::ClosestNodes; /// K = the default maximum size of a k-bucket. pub const MAX_BUCKET_SIZE_K: usize = 20; @@ -78,46 +79,31 @@ impl RoutingTable { /// Return the closest nodes to the target while prioritizing secure nodes, /// as defined in [BEP_0042](https://www.bittorrent.org/beps/bep_0042.html) pub fn closest(&self, target: &Id) -> Vec> { - let mut result = Vec::with_capacity(MAX_BUCKET_SIZE_K); - let mut unsecure = Vec::with_capacity(MAX_BUCKET_SIZE_K); + let mut closest = ClosestNodes::new(*target); - let distance = self.id.distance(target); - - for i in - // First search in closest nodes - (distance..=MAX_DISTANCE) - // if we don't have enough close nodes, populate from other buckets - .chain((0..distance).rev()) - { - match &self.buckets.get(&i) { - Some(bucket) => { - for node in bucket.iter() { - if result.len() < MAX_BUCKET_SIZE_K { - if node.is_secure() { - result.push(node.clone()); - } else { - unsecure.push(node.clone()) - } - } else { - return result; - } - } - } - None => continue, - } + for node in self.to_vec() { + closest.add(node); } - if result.len() < MAX_BUCKET_SIZE_K { - for node in unsecure { - if result.len() < MAX_BUCKET_SIZE_K { - result.push(node); - } else { - break; - } - } + closest.nodes()[..MAX_BUCKET_SIZE_K].to_vec() + } + + /// Secure version of [Self::closest] that tries to circumvent sybil attacks. + pub fn closest_secure( + &self, + target: &Id, + dht_size_estimate: usize, + subnets: usize, + ) -> Vec> { + let mut closest = ClosestNodes::new(*target); + + for node in self.to_vec() { + closest.add(node); } - result + closest + .take_until_secure(dht_size_estimate, subnets) + .to_vec() } pub fn is_empty(&self) -> bool { diff --git a/src/rpc.rs b/src/rpc.rs index 2bdffac..d0cfadb 100644 --- a/src/rpc.rs +++ b/src/rpc.rs @@ -416,7 +416,11 @@ impl Rpc { // Seed the query either with the closest nodes from the routing table, or the // bootstrapping nodes if the closest nodes are not enough. - let routing_table_closest = self.routing_table.closest(&target); + let routing_table_closest = self.routing_table.closest_secure( + &target, + self.responders_based_dht_size_estimates_sum as usize / self.closest_nodes.len().max(1), + self.subnets_sum / self.closest_nodes.len().max(1), + ); // If we don't have enough or any closest nodes, call the bootstraping nodes. if routing_table_closest.is_empty() || routing_table_closest.len() < self.bootstrap.len() {