diff --git a/p2p/auth.py b/p2p/auth.py index 555bbc1520..f04b1ec1e4 100644 --- a/p2p/auth.py +++ b/p2p/auth.py @@ -57,8 +57,20 @@ async def handshake( use_eip8 = False initiator = HandshakeInitiator(remote, privkey, use_eip8, token) reader, writer = await initiator.connect() - aes_secret, mac_secret, egress_mac, ingress_mac = await _handshake( - initiator, reader, writer, token) + try: + aes_secret, mac_secret, egress_mac, ingress_mac = await _handshake( + initiator, reader, writer, token) + except Exception: + # Note: This is one of two places where we manually handle closing the + # reader/writer connection pair in the event of an error during the + # peer connection and handshake process. + # See `p2p.peer.handshake` for the other. + if not reader.at_eof(): + reader.feed_eof() + writer.close() + await asyncio.sleep(0) + raise + return aes_secret, mac_secret, egress_mac, ingress_mac, reader, writer diff --git a/p2p/peer.py b/p2p/peer.py index 1cd90a825d..e7d5f180e3 100644 --- a/p2p/peer.py +++ b/p2p/peer.py @@ -127,8 +127,21 @@ async def handshake(remote: Node, factory: 'BasePeerFactory') -> 'BasePeer': connection=connection, inbound=False, ) - await peer.do_p2p_handshake() - await peer.do_sub_proto_handshake() + + try: + await peer.do_p2p_handshake() + await peer.do_sub_proto_handshake() + except Exception: + # Note: This is one of two places where we manually handle closing the + # reader/writer connection pair in the event of an error during the + # peer connection and handshake process. + # See `p2p.auth.handshake` for the other. + if not reader.at_eof(): + reader.feed_eof() + writer.close() + await asyncio.sleep(0) + raise + return peer @@ -342,9 +355,8 @@ def close(self) -> None: If the streams have already been closed, do nothing. """ - if self.reader.at_eof(): - return - self.reader.feed_eof() + if not self.reader.at_eof(): + self.reader.feed_eof() self.writer.close() @property