From 6fca8f723244c2e7f920a6aa18689ae61e257d03 Mon Sep 17 00:00:00 2001
From: Piper Merriam <piper@pipermerriam.com>
Date: Mon, 19 Nov 2018 13:39:35 -0700
Subject: [PATCH] Ensure network connection is closed when handshake with peer
 fails (#1476)

* Ensure network connection is closed when handshake with peer fails

* clean up remaining unclosed transport warnings

* Add comment
---
 p2p/auth.py | 16 ++++++++++++++--
 p2p/peer.py | 22 +++++++++++++++++-----
 2 files changed, 31 insertions(+), 7 deletions(-)

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