From b1dd564c32ed81d23d7cae5e1fa521cc9313dde1 Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Thu, 13 Feb 2020 10:19:02 -0500
Subject: [PATCH] make sure we don't exit the handler prematurely

---
 libp2p/muxers/muxer.nim | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/libp2p/muxers/muxer.nim b/libp2p/muxers/muxer.nim
index 6a60169..bebd19f 100644
--- a/libp2p/muxers/muxer.nim
+++ b/libp2p/muxers/muxer.nim
@@ -46,11 +46,14 @@ proc newMuxerProvider*(creator: MuxerConstructor, codec: string): MuxerProvider
 method init(c: MuxerProvider) =
   proc handler(conn: Connection, proto: string) {.async, gcsafe, closure.} =
     let muxer = c.newMuxer(conn)
+    var handlerFut: Future[void] = newFuture[void]()
     if not isNil(c.muxerHandler):
-      asyncCheck c.muxerHandler(muxer)
+      handlerFut = c.muxerHandler(muxer)
+    else:
+      handlerFut.complete()
 
     if not isNil(c.streamHandler):
       muxer.streamHandler = c.streamHandler
 
-    await muxer.handle()
+    await allFutures(muxer.handle(), handlerFut)
   c.handler = handler