From 619ed94e6421f58219c9adcea66a5de293d873c7 Mon Sep 17 00:00:00 2001
From: folex <0xdxdy@gmail.com>
Date: Thu, 19 Mar 2020 12:11:04 +0300
Subject: [PATCH] Add public key to buckets WIP

---
 protocols/kad/src/addresses.rs |  2 +-
 protocols/kad/src/behaviour.rs | 31 +++++++++++++++++--------------
 protocols/kad/src/contact.rs   |  2 +-
 3 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/protocols/kad/src/addresses.rs b/protocols/kad/src/addresses.rs
index 1198a927..02a14c4c 100644
--- a/protocols/kad/src/addresses.rs
+++ b/protocols/kad/src/addresses.rs
@@ -23,7 +23,7 @@ use smallvec::SmallVec;
 use std::fmt;
 
 /// A non-empty list of (unique) addresses of a peer in the routing table.
-#[derive(Clone)]
+#[derive(Clone, PartialEq, Eq)]
 pub struct Addresses {
     addrs: SmallVec<[Multiaddr; 6]>,
 }
diff --git a/protocols/kad/src/behaviour.rs b/protocols/kad/src/behaviour.rs
index 4cbd75b9..6ee35ed6 100644
--- a/protocols/kad/src/behaviour.rs
+++ b/protocols/kad/src/behaviour.rs
@@ -620,8 +620,9 @@ where
         let key = kbucket::Key::new(peer.clone());
         match self.kbuckets.entry(&key) {
             kbucket::Entry::Present(mut entry, old_status) => {
-                if let Some(address) = address {
-                    if entry.value().insert(address) {
+                if let Some(contact) = contact {
+                    if *entry.value() != contact { // TODO: what about public key change?
+                        *entry.value() = contact; // TODO: is there a better way to do that?
                         self.queued_events.push_back(NetworkBehaviourAction::GenerateEvent(
                             KademliaEvent::RoutingUpdated {
                                 peer,
@@ -637,8 +638,8 @@ where
             },
 
             kbucket::Entry::Pending(mut entry, old_status) => {
-                if let Some(address) = address {
-                    entry.value().insert(address);
+                if let Some(contact) = contact {
+                    *entry.value() = contact;
                 }
                 if old_status != new_status {
                     entry.update(new_status);
@@ -649,11 +650,11 @@ where
                 // Only connected nodes with a known address are newly inserted.
                 if new_status == NodeStatus::Connected {
                     if let Some(contact) = contact {
-                        match entry.insert(contact, new_status) {
+                        match entry.insert(contact.clone(), new_status) {
                             kbucket::InsertResult::Inserted => {
                                 let event = KademliaEvent::RoutingUpdated {
                                     peer: peer.clone(),
-                                    addresses: contact.addresses.clone(),
+                                    addresses: contact.addresses,
                                     old_peer: None,
                                 };
                                 self.queued_events.push_back(
@@ -1103,15 +1104,17 @@ where
             ConnectedPoint::Listener { .. } => None,
         };
 
-        let contact = self.queries.iter().find_map(|q| q.inner.contacts.get(&peer));
-        let contact = match (address, public_key) {
-            (Some(addr), Some(pk)) => Contact::new(Addresses::new(addr), pk.clone()),
-            _ => None
-        };
+        let contact = self.queries
+            .iter()
+            .find_map(|q| q.inner.contacts.get(&peer)) // Option<&Contact>
+            .as_mut() // &mut Option<&Contact>
+            .and_then(|c|
+                new_address.map(|addr| { c.insert(addr); c }) // insert(&mut self)
+            )
+            .cloned(); // Option<&Contact>
 
-
-
-        self.connection_updated(peer.clone(), address, NodeStatus::Connected);
+        // Need Option<Contact> here
+        self.connection_updated(peer.clone(), contact, NodeStatus::Connected);
         self.connected_peers.insert(peer);
     }
 
diff --git a/protocols/kad/src/contact.rs b/protocols/kad/src/contact.rs
index 1248b00c..c1424877 100644
--- a/protocols/kad/src/contact.rs
+++ b/protocols/kad/src/contact.rs
@@ -23,7 +23,7 @@ use libp2p_core::identity::ed25519::PublicKey;
 use crate::protocol::KadPeer;
 use smallvec::SmallVec;
 
-#[derive(Clone)]
+#[derive(Clone, PartialEq, Eq)]
 pub struct Contact {
     pub addresses: Addresses,
     pub public_key: PublicKey