Fix the mismatch number of peer connections

CodeChain-io · Sep 9, 2020 · 7200cc1 · 7200cc1
1 parent 49c72fa
commit 7200cc1
Showing 1 changed file with 25 additions and 7 deletions.
diff --git a/network/src/p2p/handler.rs b/network/src/p2p/handler.rs
@@ -506,6 +506,7 @@ impl IoHandler<Message> for Handler {
                 is_inbound: true,
             } => {
                 let mut inbound_connections = self.inbound_connections.write();
+                let outbound_connections = self.outbound_connections.write();
                 let target = connection.peer_addr();
                 self.peer_db.insert(*target);
                 if let Some(token) = self.inbound_tokens.lock().gen() {
@@ -526,10 +527,17 @@ impl IoHandler<Message> for Handler {
                             token
                         );
                     }
-
-                    let t = inbound_connections.insert(token, connection);
-                    assert!(t.is_none());
-                    io.register_stream(token);
+                    let can_insert = inbound_connections
+                        .values()
+                        .all(|in_connection| in_connection.peer_addr() != connection.peer_addr());
+                    let is_not_out = outbound_connections
+                        .values()
+                        .all(|out_connection| out_connection.peer_addr() != connection.peer_addr());
+                    if can_insert && is_not_out {
+                        let t = inbound_connections.insert(token, connection);
+                        assert!(t.is_none());
+                        io.register_stream(token);
+                    }
                 } else {
                     cwarn!(NETWORK, "Cannot establish an inbound connection");
                 }
@@ -539,6 +547,7 @@ impl IoHandler<Message> for Handler {
                 is_inbound: false,
             } => {
                 let mut outbound_connections = self.outbound_connections.write();
+                let inbound_connections = self.inbound_connections.write();
                 if let Some(token) = self.outbound_tokens.lock().gen() {
                     let peer_addr = *connection.peer_addr();
                     let remote_node_id = peer_addr.into();
@@ -570,9 +579,18 @@ impl IoHandler<Message> for Handler {
                             network_message_size,
                         );
                     }
-                    let t = outbound_connections.insert(token, connection);
-                    assert!(t.is_none());
-                    io.register_stream(token);
+                    let can_insert = outbound_connections
+                        .values()
+                        .all(|out_connection| out_connection.peer_addr() == connection.peer_addr());
+                    let is_not_in = inbound_connections
+                        .values()
+                        .all(|in_connection| in_connection.peer_addr() == connection.peer_addr());
+
+                    if can_insert && is_not_in {
+                        let t = outbound_connections.insert(token, connection);
+                        assert!(t.is_none());
+                        io.register_stream(token);
+                    }
                 } else {
                     cwarn!(NETWORK, "Cannot establish an outbound connection");
                 }