Fix the mismatch number of peer connections

CodeChain-io · Sep 7, 2020 · 0db816d · 0db816d
1 parent 49c72fa
commit 0db816d
Showing 1 changed file with 37 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,23 @@ impl IoHandler<Message> for Handler {
                             token
                         );
                     }
-
-                    let t = inbound_connections.insert(token, connection);
-                    assert!(t.is_none());
-                    io.register_stream(token);
+                    let mut can_insert: bool = true;
+                    let mut is_not_out: bool = true;
+                    for (_, in_connection) in inbound_connections.iter() {
+                        if in_connection.peer_addr() == connection.peer_addr() {
+                            can_insert = false;
+                        }
+                    }
+                    for (_, out_connection) in outbound_connections.iter() {
+                        if out_connection.peer_addr() == connection.peer_addr() {
+                            is_not_out = false;
+                        }
+                    }
+                    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 +553,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 +585,24 @@ impl IoHandler<Message> for Handler {
                             network_message_size,
                         );
                     }
-                    let t = outbound_connections.insert(token, connection);
-                    assert!(t.is_none());
-                    io.register_stream(token);
+                    let mut can_insert: bool = true;
+                    let mut is_not_in: bool = true;
+                    for (_, out_connection) in outbound_connections.iter() {
+                        if out_connection.peer_addr() == connection.peer_addr() {
+                            can_insert = false;
+                        }
+                    }
+                    for (_, in_connection) in inbound_connections.iter() {
+                        if in_connection.peer_addr() == connection.peer_addr() {
+                            is_not_in = false;
+                        }
+                    }
+
+                    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");
                 }