From c7606ce2eb7bf3356fff2e87e9c2cf292bdb3ce2 Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Wed, 25 Sep 2019 18:11:20 -0600
Subject: [PATCH] readMsg should never return an empty buffer

---
 libp2p/muxers/mplex/mplex.nim | 10 +-----
 tests/testidentify.nim        | 66 ++++++++++++++++-------------------
 2 files changed, 32 insertions(+), 44 deletions(-)

diff --git a/libp2p/muxers/mplex/mplex.nim b/libp2p/muxers/mplex/mplex.nim
index 63596e53f..c7963cef2 100644
--- a/libp2p/muxers/mplex/mplex.nim
+++ b/libp2p/muxers/mplex/mplex.nim
@@ -55,15 +55,7 @@ method handle*(m: Mplex) {.async, gcsafe.} =
   trace "starting mplex main loop"
   try:
     while not m.connection.closed:
-      let msgRes = await m.connection.readMsg()
-      if msgRes.isNone:
-        # TODO: this should ideally be poll(timeout = 100)
-        # but chronos poll() doesnt take a timeout, so it
-        # might and does hang
-        await sleepAsync(100.millis) # yield to async loop
-        continue
-
-      let (id, msgType, data) = msgRes.get()
+      let (id, msgType, data) = (await m.connection.readMsg()).get()
       let initiator = bool(ord(msgType) and 1)
       var channel: LPChannel
       if MessageType(msgType) != MessageType.New:
diff --git a/tests/testidentify.nim b/tests/testidentify.nim
index 78b9dc4de..9b4b1baa0 100644
--- a/tests/testidentify.nim
+++ b/tests/testidentify.nim
@@ -14,52 +14,48 @@ import ../libp2p/protocols/identify,
 suite "Identify":
   test "handle identify message":
     proc testHandle(): Future[bool] {.async.} =
-      let ma: MultiAddress = Multiaddress.init("/ip4/127.0.0.1/tcp/53360")
+      let ma: MultiAddress = Multiaddress.init("/ip4/0.0.0.0/tcp/0")
 
       let remoteSeckey = PrivateKey.random(RSA)
       var remotePeerInfo: PeerInfo
       var serverFut: Future[void]
-      var transport: TcpTransport
-      proc receiver() {.async.} =
-        remotePeerInfo.peerId = some(PeerID.init(remoteSeckey))
-        remotePeerInfo.addrs.add(ma)
-        remotePeerInfo.protocols.add("/test/proto1/1.0.0")
-        remotePeerInfo.protocols.add("/test/proto2/1.0.0")
+      remotePeerInfo.peerId = some(PeerID.init(remoteSeckey))
+      remotePeerInfo.addrs.add(ma)
+      remotePeerInfo.protocols.add("/test/proto1/1.0.0")
+      remotePeerInfo.protocols.add("/test/proto2/1.0.0")
 
-        let identifyProto = newIdentify(remotePeerInfo)
-        let msListen = newMultistream()
+      let identifyProto1 = newIdentify(remotePeerInfo)
+      let msListen = newMultistream()
 
-        msListen.addHandler(IdentifyCodec, identifyProto)
-        proc connHandler(conn: Connection): Future[void] {.async, gcsafe.} =
-          await msListen.handle(conn)
+      msListen.addHandler(IdentifyCodec, identifyProto1)
+      proc connHandler(conn: Connection): Future[void] {.async, gcsafe.} =
+        await msListen.handle(conn)
 
-        transport = newTransport(TcpTransport)
-        serverFut = await transport.listen(ma, connHandler)
+      var transport1 = newTransport(TcpTransport)
+      serverFut = await transport1.listen(ma, connHandler)
 
-      proc sender() {.async.} =
-        let msDial = newMultistream()
-        let transport: TcpTransport = newTransport(TcpTransport)
-        let conn = await transport.dial(ma)
+      let msDial = newMultistream()
+      let transport2: TcpTransport = newTransport(TcpTransport)
+      let conn = await transport2.dial(transport1.ma)
 
-        let seckey = PrivateKey.random(RSA)
-        var peerInfo: PeerInfo
-        peerInfo.peerId = some(PeerID.init(seckey))
-        peerInfo.addrs.add(ma)
+      let seckey = PrivateKey.random(RSA)
+      var peerInfo: PeerInfo
+      peerInfo.peerId = some(PeerID.init(seckey))
+      peerInfo.addrs.add(ma)
 
-        let identifyProto = newIdentify(peerInfo)
-        let res = await msDial.select(conn, IdentifyCodec)
-        let id = await identifyProto.identify(conn, remotePeerInfo)
+      let identifyProto2 = newIdentify(peerInfo)
+      let res = await msDial.select(conn, IdentifyCodec)
+      let id = await identifyProto2.identify(conn, remotePeerInfo)
 
-        check id.pubKey.get() == remoteSeckey.getKey()
-        check id.addrs[0] == ma
-        check id.protoVersion.get() == ProtoVersion
-        # check id.agentVersion.get() == AgentVersion
-        check id.protos == @["/test/proto1/1.0.0", "/test/proto2/1.0.0"]
+      check id.pubKey.get() == remoteSeckey.getKey()
+      check id.addrs[0] == ma
+      check id.protoVersion.get() == ProtoVersion
+      # check id.agentVersion.get() == AgentVersion
+      check id.protos == @["/test/proto1/1.0.0", "/test/proto2/1.0.0"]
 
-        await conn.close()
+      await conn.close()
 
-      await allFutures(sender(), receiver())
-      await transport.close()
+      await transport1.close()
       await serverFut
       result = true
 
@@ -68,7 +64,7 @@ suite "Identify":
 
   test "handle failed identify":
     proc testHandleError() {.async.} =
-      let ma: MultiAddress = Multiaddress.init("/ip4/127.0.0.1/tcp/53361")
+      let ma: MultiAddress = Multiaddress.init("/ip4/0.0.0.0/tcp/0")
 
       let remoteSeckey = PrivateKey.random(RSA)
       var remotePeerInfo: PeerInfo
@@ -87,7 +83,7 @@ suite "Identify":
 
       let msDial = newMultistream()
       let transport2: TcpTransport = newTransport(TcpTransport)
-      let conn = await transport2.dial(ma)
+      let conn = await transport2.dial(transport1.ma)
 
       let seckey = PrivateKey.random(RSA)
       var localPeerInfo: PeerInfo