Fix duplicate connection bugs

CodeChain-io · Sep 7, 2020 · 49c72fa · 49c72fa
1 parent bc771a4
commit 49c72fa
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 23 deletions.
diff --git a/network/src/p2p/handler.rs b/network/src/p2p/handler.rs
@@ -517,13 +517,15 @@ impl IoHandler<Message> for Handler {
                         remote_node_id,
                         token
                     );
-                    assert_eq!(
-                        None,
-                        self.remote_node_ids_reverse.write().insert(remote_node_id, token),
-                        "{}:{} is already registered",
-                        remote_node_id,
-                        token
-                    );
+                    if !self.remote_node_ids_reverse.write().contains_key(&remote_node_id) {
+                        assert_eq!(
+                            None,
+                            self.remote_node_ids_reverse.write().insert(remote_node_id, token),
+                            "{}:{} is already registered",
+                            remote_node_id,
+                            token
+                        );
+                    }
 
                     let t = inbound_connections.insert(token, connection);
                     assert!(t.is_none());
@@ -547,13 +549,15 @@ impl IoHandler<Message> for Handler {
                         remote_node_id,
                         token
                     );
-                    assert_eq!(
-                        None,
-                        self.remote_node_ids_reverse.write().insert(remote_node_id, token),
-                        "{}:{} is already registered",
-                        remote_node_id,
-                        token
-                    );
+                    if !self.remote_node_ids_reverse.write().contains_key(&remote_node_id) {
+                        assert_eq!(
+                            None,
+                            self.remote_node_ids_reverse.write().insert(remote_node_id, token),
+                            "{}:{} is already registered",
+                            remote_node_id,
+                            token
+                        );
+                    }
 
                     let mut network_message_size = 0;
                     for (name, versions) in self.client.extension_versions() {

diff --git a/sync/src/block/extension.rs b/sync/src/block/extension.rs
@@ -476,22 +476,27 @@ impl NetworkExtension<Event> for Extension {
         cinfo!(SYNC, "New peer detected #{}", id);
         self.send_status(id);
 
-        let t = self.connected_nodes.insert(*id);
-        debug_assert!(t, "{} is already added to peer list", id);
+        if !self.connected_nodes.contains(id) {
+            let t = self.connected_nodes.insert(*id);
+            debug_assert!(t, "{} is already added to peer list", id);
+        }
 
         let token = self.token_generator.gen().expect("Token generator is full");
         let token_info = TokenInfo {
             node_id: *id,
             request_id: None,
         };
 
-        let t = self.requests.insert(*id, Vec::new());
-        debug_assert_eq!(None, t);
-        let t = self.tokens_info.insert(token, token_info);
-        debug_assert_eq!(None, t);
-        let t = self.tokens.insert(*id, token);
-        debug_assert_eq!(None, t);
-        debug_assert!(t.is_none());
+        if !self.requests.contains_key(id) {
+            let t = self.requests.insert(*id, Vec::new());
+            debug_assert_eq!(None, t);
+        }
+        self.tokens_info.entry(token).or_insert(token_info);
+        if !self.tokens.contains_key(id) {
+            let t = self.tokens.insert(*id, token);
+            debug_assert_eq!(None, t);
+            debug_assert!(t.is_none());
+        }
     }
 
     fn on_node_removed(&mut self, id: &NodeId) {