From 34e330353fe5234817ca1e613ed98bc5233b2ed5 Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Mon, 18 Jan 2021 16:27:29 -0600
Subject: [PATCH] better `upgraded` lifetime handling (avoid NPE) (#506)

* avoid npe on connection upgrade

* add `onUpgraded` event
---
 libp2p/connmanager.nim       |  4 +---
 libp2p/stream/connection.nim | 15 +++++++++++++++
 libp2p/switch.nim            | 14 +++++---------
 3 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/libp2p/connmanager.nim b/libp2p/connmanager.nim
index 970f93597..b9d2855ed 100644
--- a/libp2p/connmanager.nim
+++ b/libp2p/connmanager.nim
@@ -240,9 +240,7 @@ proc cleanupConn(c: ConnManager, conn: Connection) {.async.} =
 proc onConnUpgraded(c: ConnManager, conn: Connection) {.async.} =
   try:
     trace "Triggering connect events", conn
-    doAssert(not isNil(conn.upgraded),
-      "The `upgraded` event hasn't been properly initialized!")
-    conn.upgraded.complete()
+    conn.upgrade()
 
     let peerId = conn.peerInfo.peerId
     await c.triggerPeerEvents(
diff --git a/libp2p/stream/connection.nim b/libp2p/stream/connection.nim
index eae67c645..a523673f1 100644
--- a/libp2p/stream/connection.nim
+++ b/libp2p/stream/connection.nim
@@ -36,6 +36,21 @@ type
 
 proc timeoutMonitor(s: Connection) {.async, gcsafe.}
 
+proc isUpgraded*(s: Connection): bool =
+  if not isNil(s.upgraded):
+    return s.upgraded.finished
+
+proc upgrade*(s: Connection, failed: ref Exception = nil) =
+  if not isNil(failed):
+    s.upgraded.fail(failed)
+    return
+
+  s.upgraded.complete()
+
+proc onUpgrade*(s: Connection) {.async.} =
+  if not isNil(s.upgraded):
+    await s.upgraded
+
 func shortLog*(conn: Connection): string =
   if conn.isNil: "Connection(nil)"
   elif conn.peerInfo.isNil: $conn.oid
diff --git a/libp2p/switch.nim b/libp2p/switch.nim
index 7e8914993..e513c864d 100644
--- a/libp2p/switch.nim
+++ b/libp2p/switch.nim
@@ -244,10 +244,8 @@ proc upgradeIncoming(s: Switch, incomingConn: Connection) {.async, gcsafe.} = #
       await ms.handle(cconn)
     except CatchableError as exc:
       debug "Exception in secure handler during incoming upgrade", msg = exc.msg, conn
-      if  not isNil(cconn) and
-          not isNil(cconn.upgraded) and
-          not(cconn.upgraded.finished):
-        cconn.upgraded.fail(exc)
+      if not cconn.isUpgraded:
+        cconn.upgrade(exc)
     finally:
       if not isNil(cconn):
         await cconn.close()
@@ -265,10 +263,8 @@ proc upgradeIncoming(s: Switch, incomingConn: Connection) {.async, gcsafe.} = #
     await ms.handle(incomingConn, active = true)
   except CatchableError as exc:
     debug "Exception upgrading incoming", exc = exc.msg
-    if  not isNil(incomingConn) and
-        not isNil(incomingConn.upgraded) and
-        not(incomingConn.upgraded.finished):
-      incomingConn.upgraded.fail(exc)
+    if not incomingConn.isUpgraded:
+      incomingConn.upgrade(exc)
   finally:
     if not isNil(incomingConn):
       await incomingConn.close()
@@ -441,7 +437,7 @@ proc upgradeMonitor(conn: Connection, upgrades: AsyncSemaphore) {.async.} =
     # upgrade, this timeout guarantees that a
     # "hanged" remote doesn't hold the upgrade
     # forever
-    await conn.upgraded.wait(30.seconds) # wait for connection to be upgraded
+    await conn.onUpgrade.wait(30.seconds) # wait for connection to be upgraded
     trace "Connection upgrade succeeded"
   except CatchableError as exc:
     libp2p_failed_upgrades_incoming.inc()