Use Into<kbucket::Key<K>>.

libp2p · Dec 7, 2020 · 69bd277 · 69bd277
1 parent d21bd64
commit 69bd277
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 12 deletions.
diff --git a/core/src/peer_id.rs b/core/src/peer_id.rs
@@ -153,6 +153,12 @@ impl From<PeerId> for Multihash {
     }
 }
 
+impl From<PeerId> for Vec<u8> {
+    fn from(peer_id: PeerId) -> Self {
+        peer_id.to_bytes()
+    }
+}
+
 #[derive(Debug, Error)]
 pub enum ParseError {
     #[error("base-58 decode error: {0}")]

diff --git a/protocols/kad/src/behaviour.rs b/protocols/kad/src/behaviour.rs
@@ -47,7 +47,7 @@ use libp2p_swarm::{
 };
 use log::{info, debug, warn};
 use smallvec::SmallVec;
-use std::{borrow::{Borrow, Cow}, error, iter, time::Duration};
+use std::{borrow::Cow, error, iter, time::Duration};
 use std::collections::{HashSet, VecDeque};
 use std::fmt;
 use std::num::NonZeroUsize;
@@ -551,9 +551,9 @@ where
     pub fn kbucket<K>(&mut self, key: K)
         -> Option<kbucket::KBucketRef<'_, kbucket::Key<PeerId>, Addresses>>
     where
-        K: Borrow<[u8]> + Clone
+        K: Into<kbucket::Key<K>> + Clone
     {
-        self.kbuckets.bucket(&kbucket::Key::new(key))
+        self.kbuckets.bucket(&key.into())
     }
 
     /// Initiates an iterative query for the closest peers to the given key.
@@ -562,10 +562,10 @@ where
     /// [`KademliaEvent::QueryResult{QueryResult::GetClosestPeers}`].
     pub fn get_closest_peers<K>(&mut self, key: K) -> QueryId
     where
-        K: Borrow<[u8]> + Clone
+        K: Into<kbucket::Key<K>> + Into<Vec<u8>> + Clone
     {
-        let info = QueryInfo::GetClosestPeers { key: key.borrow().to_vec() };
-        let target = kbucket::Key::new(key);
+        let info = QueryInfo::GetClosestPeers { key: key.clone().into() };
+        let target: kbucket::Key<K> = key.into();
         let peers = self.kbuckets.closest_keys(&target);
         let inner = QueryInner::new(info);
         self.queries.add_iter_closest(target.clone(), peers, inner)

diff --git a/protocols/kad/src/behaviour/test.rs b/protocols/kad/src/behaviour/test.rs
@@ -240,7 +240,7 @@ fn query_iter() {
         // propagate forwards through the list of peers.
         let search_target = PeerId::random();
         let search_target_key = kbucket::Key::from(search_target);
-        let qid = swarms[0].get_closest_peers(search_target.to_bytes());
+        let qid = swarms[0].get_closest_peers(search_target);
 
         match swarms[0].query(&qid) {
             Some(q) => match q.info() {
@@ -310,7 +310,7 @@ fn unresponsive_not_returned_direct() {
 
     // Ask first to search a random value.
     let search_target = PeerId::random();
-    swarms[0].get_closest_peers(search_target.to_bytes());
+    swarms[0].get_closest_peers(search_target);
 
     block_on(
         poll_fn(move |ctx| {
@@ -360,7 +360,7 @@ fn unresponsive_not_returned_indirect() {
 
     // Ask second to search a random value.
     let search_target = PeerId::random();
-    swarms[1].get_closest_peers(search_target.to_bytes());
+    swarms[1].get_closest_peers(search_target);
 
     block_on(
         poll_fn(move |ctx| {
@@ -885,7 +885,7 @@ fn exceed_jobs_max_queries() {
     let (_addr, mut swarm) = build_node();
     let num = JOBS_MAX_QUERIES + 1;
     for _ in 0 .. num {
-        swarm.get_closest_peers(PeerId::random().to_bytes());
+        swarm.get_closest_peers(PeerId::random());
     }
 
     assert_eq!(swarm.queries.size(), num);
@@ -1072,7 +1072,7 @@ fn manual_bucket_inserts() {
     // that none of them was inserted into a bucket.
     let mut routable = Vec::new();
     // Start an iterative query from the first peer.
-    swarms[0].1.get_closest_peers(PeerId::random().to_bytes());
+    swarms[0].1.get_closest_peers(PeerId::random());
     block_on(poll_fn(move |ctx| {
         for (_, swarm) in swarms.iter_mut() {
             loop {
@@ -1084,7 +1084,7 @@ fn manual_bucket_inserts() {
                         routable.push(peer);
                         if expected.is_empty() {
                             for peer in routable.iter() {
-                                let bucket = swarm.kbucket(peer.to_bytes()).unwrap();
+                                let bucket = swarm.kbucket(*peer).unwrap();
                                 assert!(bucket.iter().all(|e| e.node.key.preimage() != peer));
                             }
                             return Poll::Ready(())