From bd41e0461e485c32525d6cb68cd7e1606157305b Mon Sep 17 00:00:00 2001
From: Nazar Mokrynskyi <nazar@mokrynskyi.com>
Date: Mon, 7 Feb 2022 23:56:30 +0200
Subject: [PATCH] protocols/kad: Require owned key in get_record (#2477)

Co-authored-by: Max Inden <mail@max-inden.de>
---
 examples/distributed-key-value-store.rs |  2 +-
 protocols/kad/CHANGELOG.md              |  4 ++++
 protocols/kad/src/behaviour.rs          | 10 +++++-----
 protocols/kad/src/behaviour/test.rs     | 10 ++++++----
 protocols/kad/src/kbucket.rs            |  2 +-
 5 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/examples/distributed-key-value-store.rs b/examples/distributed-key-value-store.rs
index 3f92ebbe..6bf28bf0 100644
--- a/examples/distributed-key-value-store.rs
+++ b/examples/distributed-key-value-store.rs
@@ -187,7 +187,7 @@ fn handle_input_line(kademlia: &mut Kademlia<MemoryStore>, line: String) {
                     }
                 }
             };
-            kademlia.get_record(&key, Quorum::One);
+            kademlia.get_record(key, Quorum::One);
         }
         Some("GET_PROVIDERS") => {
             let key = {
diff --git a/protocols/kad/CHANGELOG.md b/protocols/kad/CHANGELOG.md
index d8d2c3c2..37cde354 100644
--- a/protocols/kad/CHANGELOG.md
+++ b/protocols/kad/CHANGELOG.md
@@ -4,6 +4,10 @@
 
 - Update to `libp2p-swarm` `v0.34.0`.
 
+- Require owned key in `get_record()` method (see [PR 2477]).
+
+[PR 2477]: https://github.com/libp2p/rust-libp2p/pull/2477
+
 # 0.34.0 [2022-01-27]
 
 - Update dependencies.
diff --git a/protocols/kad/src/behaviour.rs b/protocols/kad/src/behaviour.rs
index b213f964..196a86f5 100644
--- a/protocols/kad/src/behaviour.rs
+++ b/protocols/kad/src/behaviour.rs
@@ -679,13 +679,13 @@ where
     ///
     /// The result of this operation is delivered in a
     /// [`KademliaEvent::OutboundQueryCompleted{QueryResult::GetRecord}`].
-    pub fn get_record(&mut self, key: &record::Key, quorum: Quorum) -> QueryId {
+    pub fn get_record(&mut self, key: record::Key, quorum: Quorum) -> QueryId {
         let quorum = quorum.eval(self.queries.config().replication_factor);
         let mut records = Vec::with_capacity(quorum.get());
 
-        if let Some(record) = self.store.get(key) {
+        if let Some(record) = self.store.get(&key) {
             if record.is_expired(Instant::now()) {
-                self.store.remove(key)
+                self.store.remove(&key)
             } else {
                 records.push(PeerRecord {
                     peer: None,
@@ -697,7 +697,7 @@ where
         let done = records.len() >= quorum.get();
         let target = kbucket::Key::new(key.clone());
         let info = QueryInfo::GetRecord {
-            key: key.clone(),
+            key,
             records,
             quorum,
             cache_candidates: BTreeMap::new(),
@@ -1231,7 +1231,7 @@ where
 
                 if let Some(target) = remaining.next() {
                     let info = QueryInfo::Bootstrap {
-                        peer: target.clone().into_preimage(),
+                        peer: *target.preimage(),
                         remaining: Some(remaining),
                     };
                     let peers = self.kbuckets.closest_keys(&target);
diff --git a/protocols/kad/src/behaviour/test.rs b/protocols/kad/src/behaviour/test.rs
index 21eb76c1..aa462202 100644
--- a/protocols/kad/src/behaviour/test.rs
+++ b/protocols/kad/src/behaviour/test.rs
@@ -448,7 +448,7 @@ fn get_record_not_found() {
     let target_key = record::Key::from(random_multihash());
     let qid = swarms[0]
         .behaviour_mut()
-        .get_record(&target_key, Quorum::One);
+        .get_record(target_key.clone(), Quorum::One);
 
     block_on(poll_fn(move |ctx| {
         for swarm in &mut swarms {
@@ -761,7 +761,7 @@ fn get_record() {
     swarms[2].behaviour_mut().store.put(record.clone()).unwrap();
     let qid = swarms[0]
         .behaviour_mut()
-        .get_record(&record.key, Quorum::One);
+        .get_record(record.key.clone(), Quorum::One);
 
     block_on(poll_fn(move |ctx| {
         for swarm in &mut swarms {
@@ -817,7 +817,9 @@ fn get_record_many() {
     }
 
     let quorum = Quorum::N(NonZeroUsize::new(num_results).unwrap());
-    let qid = swarms[0].behaviour_mut().get_record(&record.key, quorum);
+    let qid = swarms[0]
+        .behaviour_mut()
+        .get_record(record.key.clone(), quorum);
 
     block_on(poll_fn(move |ctx| {
         for swarm in &mut swarms {
@@ -1116,7 +1118,7 @@ fn disjoint_query_does_not_finish_before_all_paths_did() {
     let (mut alice, mut bob, mut trudy) = (alice.1, bob.1, trudy.1);
 
     // Have `alice` query the Dht for `key` with a quorum of 1.
-    alice.behaviour_mut().get_record(&key, Quorum::One);
+    alice.behaviour_mut().get_record(key, Quorum::One);
 
     // The default peer timeout is 10 seconds. Choosing 1 seconds here should
     // give enough head room to prevent connections to `bob` to time out.
diff --git a/protocols/kad/src/kbucket.rs b/protocols/kad/src/kbucket.rs
index 11140778..8b85a65b 100644
--- a/protocols/kad/src/kbucket.rs
+++ b/protocols/kad/src/kbucket.rs
@@ -240,7 +240,7 @@ where
     /// increasing distance.
     pub fn closest_keys<'a, T>(&'a mut self, target: &'a T) -> impl Iterator<Item = TKey> + 'a
     where
-        T: Clone + AsRef<KeyBytes>,
+        T: AsRef<KeyBytes>,
     {
         let distance = self.local_key.as_ref().distance(target);
         ClosestIter {