From 35ff99829e741c3fe0d50eb41800b8321cf6ea8a Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Mon, 8 Jun 2020 17:42:27 -0600
Subject: [PATCH] close streams (#208)

* close streams

* don't warn on missing proto
---
 libp2p/multistream.nim        |  2 +-
 libp2p/protocols/identify.nim | 13 +++++++++----
 libp2p/switch.nim             | 13 ++++++++-----
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/libp2p/multistream.nim b/libp2p/multistream.nim
index 1221c2ec1..ca24c8695 100644
--- a/libp2p/multistream.nim
+++ b/libp2p/multistream.nim
@@ -145,7 +145,7 @@ proc handle*(m: MultistreamSelect, conn: Connection) {.async, gcsafe.} =
               await conn.writeLp((h.proto & "\n"))
               await h.protocol.handler(conn, ms)
               return
-          warn "no handlers for ", protocol = ms
+          debug "no handlers for ", protocol = ms
           await conn.write(Na)
   except CatchableError as exc:
     trace "exception in multistream", exc = exc.msg
diff --git a/libp2p/protocols/identify.nim b/libp2p/protocols/identify.nim
index cdf051cf3..4d7b2985c 100644
--- a/libp2p/protocols/identify.nim
+++ b/libp2p/protocols/identify.nim
@@ -112,10 +112,15 @@ proc newIdentify*(peerInfo: PeerInfo): Identify =
 
 method init*(p: Identify) =
   proc handle(conn: Connection, proto: string) {.async, gcsafe, closure.} =
-    trace "handling identify request"
-    var pb = encodeMsg(p.peerInfo, await conn.getObservedAddrs())
-    await conn.writeLp(pb.buffer)
-    # await conn.close() #TODO: investigate why this breaks
+    try:
+      try:
+        trace "handling identify request"
+        var pb = encodeMsg(p.peerInfo, await conn.getObservedAddrs())
+        await conn.writeLp(pb.buffer)
+      finally:
+        await conn.close()
+    except CatchableError as exc:
+      trace "exception in identify handler", exc = exc.msg
 
   p.handler = handle
   p.codec = IdentifyCodec
diff --git a/libp2p/switch.nim b/libp2p/switch.nim
index 7f317a479..7cb5e2787 100644
--- a/libp2p/switch.nim
+++ b/libp2p/switch.nim
@@ -110,6 +110,8 @@ proc mux(s: Switch, conn: Connection): Future[void] {.async, gcsafe.} =
 
   # new stream for identify
   var stream = await muxer.newStream()
+  # call muxer handler, this should
+  # not end until muxer ends
   let handlerFut = muxer.handle()
 
   # add muxer handler cleanup proc
@@ -117,11 +119,12 @@ proc mux(s: Switch, conn: Connection): Future[void] {.async, gcsafe.} =
     trace "muxer handler completed for peer",
       peer = conn.peerInfo.id
 
-  # do identify first, so that we have a
-  # PeerInfo in case we didn't before
-  conn.peerInfo = await s.identify(stream)
-
-  await stream.close() # close identify stream
+  try:
+    # do identify first, so that we have a
+    # PeerInfo in case we didn't before
+    conn.peerInfo = await s.identify(stream)
+  finally:
+    await stream.close() # close identify stream
 
   trace "connection's peerInfo", peerInfo = $conn.peerInfo