From 74f8b5b5f1331c36228284dfa0c2388c3f431fb8 Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Sat, 23 May 2020 10:51:54 -0600
Subject: [PATCH] resource cleanup

---
 libp2p/protocols/secure/noise.nim  |  3 +-
 libp2p/protocols/secure/secio.nim  |  7 +---
 libp2p/protocols/secure/secure.nim | 55 +++++++++++++++++-------------
 3 files changed, 35 insertions(+), 30 deletions(-)

diff --git a/libp2p/protocols/secure/noise.nim b/libp2p/protocols/secure/noise.nim
index f83bac42b..0a13abd14 100644
--- a/libp2p/protocols/secure/noise.nim
+++ b/libp2p/protocols/secure/noise.nim
@@ -500,7 +500,8 @@ method handshake*(p: Noise, conn: Connection, initiator: bool = false): Future[S
       raise newException(NoiseHandshakeError, "Noise handshake, peer infos don't match! " & $pid & " != " & $conn.peerInfo.peerId)
 
   var secure = new NoiseConnection
-  inc getConnectionTracker().opened
+  secure.initStream()
+
   secure.stream = conn
   secure.closeEvent = newAsyncEvent()
   secure.peerInfo = PeerInfo.init(remotePubKey)
diff --git a/libp2p/protocols/secure/secio.nim b/libp2p/protocols/secure/secio.nim
index 599ff8295..d2edd4f5f 100644
--- a/libp2p/protocols/secure/secio.nim
+++ b/libp2p/protocols/secure/secio.nim
@@ -246,9 +246,8 @@ proc newSecioConn(conn: Connection,
   ## cipher algorithm ``cipher``, stretched keys ``secrets`` and order
   ## ``order``.
   new result
-
+  result.initStream()
   result.stream = conn
-  result.closeEvent = newAsyncEvent()
 
   let i0 = if order < 0: 1 else: 0
   let i1 = if order < 0: 0 else: 1
@@ -267,10 +266,6 @@ proc newSecioConn(conn: Connection,
                           secrets.ivOpenArray(i1))
 
   result.peerInfo = PeerInfo.init(remotePubKey)
-  when chronicles.enabledLogLevel == LogLevel.TRACE:
-    result.oid = genOid()
-
-  inc  getConnectionTracker().opened
 
 proc transactMessage(conn: Connection,
                      msg: seq[byte]): Future[seq[byte]] {.async.} =
diff --git a/libp2p/protocols/secure/secure.nim b/libp2p/protocols/secure/secure.nim
index 39f3e6910..fd1de092a 100644
--- a/libp2p/protocols/secure/secure.nim
+++ b/libp2p/protocols/secure/secure.nim
@@ -44,9 +44,8 @@ method init*(s: Secure) {.gcsafe.} =
       discard await s.handleConn(conn, false)
       trace "connection secured"
     except CatchableError as exc:
-      if not conn.closed():
-        warn "securing connection failed", msg = exc.msg
-        await conn.close()
+      warn "securing connection failed", msg = exc.msg
+      await conn.close()
 
   s.handler = handle
 
@@ -63,32 +62,42 @@ method readExactly*(s: SecureConn,
                     pbytes: pointer,
                     nbytes: int):
                     Future[void] {.async, gcsafe.} =
-  if nbytes == 0:
-    return
+  try:
+    if nbytes == 0:
+      return
 
-  while s.buf.data().len < nbytes:
-    # TODO write decrypted content straight into buf using `prepare`
-    let buf = await s.readMessage()
-    if buf.len == 0:
-      raise newLPStreamIncompleteError()
-    s.buf.add(buf)
+    while s.buf.data().len < nbytes:
+      # TODO write decrypted content straight into buf using `prepare`
+      let buf = await s.readMessage()
+      if buf.len == 0:
+        raise newLPStreamIncompleteError()
+      s.buf.add(buf)
 
-  var p = cast[ptr UncheckedArray[byte]](pbytes)
-  let consumed = s.buf.consumeTo(toOpenArray(p, 0, nbytes - 1))
-  doAssert consumed == nbytes, "checked above"
+    var p = cast[ptr UncheckedArray[byte]](pbytes)
+    let consumed = s.buf.consumeTo(toOpenArray(p, 0, nbytes - 1))
+    doAssert consumed == nbytes, "checked above"
+  except CatchableError as exc:
+    trace "exception reading from secure connection", exc = exc.msg
+    await s.close() # make sure to close the wrapped connection
+    raise exc
 
 method readOnce*(s: SecureConn,
                  pbytes: pointer,
                  nbytes: int):
                  Future[int] {.async, gcsafe.} =
-  if nbytes == 0:
-    return 0
+  try:
+    if nbytes == 0:
+      return 0
 
-  if s.buf.data().len() == 0:
-    let buf = await s.readMessage()
-    if buf.len == 0:
-      raise newLPStreamIncompleteError()
-    s.buf.add(buf)
+    if s.buf.data().len() == 0:
+      let buf = await s.readMessage()
+      if buf.len == 0:
+        raise newLPStreamIncompleteError()
+      s.buf.add(buf)
 
-  var p = cast[ptr UncheckedArray[byte]](pbytes)
-  return s.buf.consumeTo(toOpenArray(p, 0, nbytes - 1))
+    var p = cast[ptr UncheckedArray[byte]](pbytes)
+    return s.buf.consumeTo(toOpenArray(p, 0, nbytes - 1))
+  except CatchableError as exc:
+    trace "exception reading from secure connection", exc = exc.msg
+    await s.close() # make sure to close the wrapped connection
+    raise exc