diff --git a/p2p/peer.py b/p2p/peer.py index f86e77cc1b..6e85157c6e 100644 --- a/p2p/peer.py +++ b/p2p/peer.py @@ -599,14 +599,11 @@ def is_full(self) -> bool: return len(self) >= self.max_peers def is_valid_connection_candidate(self, candidate: Node) -> bool: - max_matching_ip = 2 - matching_ip = 0 - for peer in self.connected_nodes.values(): - if candidate.address.ip == peer.remote.address.ip: - matching_ip += 1 - if (matching_ip > max_matching_ip): - return False - return True + # connect to no more then 2 nodes with the same IP + nodes_by_ip = groupby( + operator.attrgetter('remote.address.ip'), self.connected_nodes.values()) + matching_ip_nodes = nodes_by_ip.get(candidate.address.ip, []) + return len(matching_ip_nodes) <= 2 def get_nodes_to_connect(self) -> Generator[Node, None, None]: for node in self.discovery.get_random_nodes(self.max_peers): diff --git a/p2p/server.py b/p2p/server.py index ef02269c60..f6f0edd24a 100644 --- a/p2p/server.py +++ b/p2p/server.py @@ -355,25 +355,21 @@ async def _receive_handshake( inbound=True, ) - total_peers_inbound = 0.0 - total_peers = 0.0 - - for current_peers in self.peer_pool.connected_nodes.values(): - total_peers += 1 - if current_peers.inbound: - total_peers_inbound += 1 - # make sure to have atleast 1/4 outbound connections - if self.peer_pool.is_full: peer.disconnect(DisconnectReason.too_many_peers) - elif total_peers > 1 and total_peers_inbound // total_peers > 0.75: - peer.disconnect(DisconnectReason.useless_peer) - elif not self.peer_pool.is_valid_connection_candidate(peer.remote): - peer.disconnect(DisconnectReason.useless_peer) else: - # We use self.wait() here as a workaround for - # https://github.com/ethereum/py-evm/issues/670. - await self.wait(self.do_handshake(peer)) + total_peers = len(self.peer_pool.connected_nodes) + inbound_peers = len( + [peer for peer in self.peer_pool.connected_nodes.values() if peer.inbound]) + if total_peers > 1 and inbound_peers / total_peers > 0.75: + # make sure to have at least 1/4 outbound connections + peer.disconnect(DisconnectReason.useless_peer) + elif not self.peer_pool.is_valid_connection_candidate(peer.remote): + peer.disconnect(DisconnectReason.useless_peer) + else: + # We use self.wait() here as a workaround for + # https://github.com/ethereum/py-evm/issues/670. + await self.wait(self.do_handshake(peer)) async def do_handshake(self, peer: BasePeer) -> None: await peer.do_p2p_handshake(), diff --git a/tests/p2p/test_server.py b/tests/p2p/test_server.py index e5b00ef652..b178261bd1 100644 --- a/tests/p2p/test_server.py +++ b/tests/p2p/test_server.py @@ -51,6 +51,9 @@ class MockPeerPool(): def is_valid_connection_candidate(self, node): return True + def __len__(self): + return len(self.connected_nodes) + def get_server(privkey, address, peer_class): base_db = MemoryDB()