Add NetworkBehaviour::inject_replaced (#914)

* Add NetworkBehaviour::inject_replaced

* Address style

* Forgot to call set_disconnected

* Also add incoming addresses to kbuckets

protocols/kad/src/behaviour.rs

@@ -291,7 +291,7 @@ where
             .unwrap_or_else(Vec::new)
     }
 
-    fn inject_connected(&mut self, id: PeerId, _: ConnectedPoint) {
+    fn inject_connected(&mut self, id: PeerId, endpoint: ConnectedPoint) {
         if let Some(pos) = self.pending_rpcs.iter().position(|(p, _)| p == &id) {
             let (_, rpc) = self.pending_rpcs.remove(pos);
             self.queued_events.push(NetworkBehaviourAction::SendEvent {
@@ -309,6 +309,14 @@ where
             _ => ()
         }
 
+        if let ConnectedPoint::Dialer { address } = endpoint {
+            if let Some(list) = self.kbuckets.entry_mut(&id) {
+                if list.iter().all(|a| *a != address) {
+                    list.push(address);
+                }
+            }
+        }
+
         self.connected_peers.insert(id);
     }
 
@@ -319,6 +327,28 @@ where
         for (query, _, _) in self.active_queries.values_mut() {
             query.inject_rpc_error(id);
         }
+
+        self.kbuckets.set_disconnected(&id);
+    }
+
+    fn inject_replaced(&mut self, peer_id: PeerId, _: ConnectedPoint, new_endpoint: ConnectedPoint) {
+        // We need to re-send the active queries.
+        for (query_id, (query, _, _)) in self.active_queries.iter() {
+            if query.is_waiting(&peer_id) {
+                self.queued_events.push(NetworkBehaviourAction::SendEvent {
+                    peer_id: peer_id.clone(),
+                    event: query.target().to_rpc_request(*query_id),
+                });
+            }
+        }
+
+        if let ConnectedPoint::Dialer { address } = new_endpoint {
+            if let Some(list) = self.kbuckets.entry_mut(&peer_id) {
+                if list.iter().all(|a| *a != address) {
+                    list.push(address);
+                }
+            }
+        }
     }
 
     fn inject_node_event(&mut self, source: PeerId, event: KademliaHandlerEvent<QueryId>) {

protocols/kad/src/query.rs

@@ -216,6 +216,33 @@ impl QueryState {
         }
     }
 
+    /// Returns true if we are waiting for a query answer from that peer.
+    ///
+    /// After `poll()` returned `SendRpc`, this function will return `true`.
+    pub fn is_waiting(&self, id: &PeerId) -> bool {
+        let state = self
+            .closest_peers
+            .iter()
+            .filter_map(
+                |(peer_id, state)| {
+                    if peer_id == id {
+                        Some(state)
+                    } else {
+                        None
+                    }
+                },
+            )
+            .next();
+
+        match state {
+            Some(&QueryPeerState::InProgress(_)) => true,
+            Some(&QueryPeerState::NotContacted) => false,
+            Some(&QueryPeerState::Succeeded) => false,
+            Some(&QueryPeerState::Failed) => false,
+            None => false,
+        }
+    }
+
     /// After `poll()` returned `SendRpc`, this function should be called if we were unable to
     /// reach the peer, or if an error of some sort happened.
     ///