From abf659a01a654020018f0a94b958e444c08702b3 Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Mon, 1 Jun 2020 11:49:25 -0600
Subject: [PATCH] more consistent dialing proto selecting logic

---
 libp2p/multistream.nim | 23 ++++++++++++-----------
 libp2p/switch.nim      | 27 ++++++++++-----------------
 tests/testinterop.nim  |  4 ----
 3 files changed, 22 insertions(+), 32 deletions(-)

diff --git a/libp2p/multistream.nim b/libp2p/multistream.nim
index 15fbe73..f9b0e0c 100644
--- a/libp2p/multistream.nim
+++ b/libp2p/multistream.nim
@@ -71,18 +71,19 @@ proc select*(m: MultistreamSelect,
   trace "reading first requested proto"
   result.removeSuffix("\n")
   if result == proto[0]:
-    trace "succesfully selected ", proto = proto
+    trace "successfully selected ", proto = proto
     return
 
-  if not result.len > 0:
-    trace "selecting one of several protos"
-    for p in proto[1..<proto.len()]:
-      await conn.writeLp((p & "\n")) # select proto
-      result = string.fromBytes(await conn.readLp(1024)) # read the first proto
-      result.removeSuffix("\n")
-      if result == p:
-        trace "selected protocol", protocol = result
-        break
+  let protos = proto[1..<proto.len()]
+  trace "selecting one of several protos", protos = protos
+  for p in protos:
+    trace "selecting proto", proto = p
+    await conn.writeLp((p & "\n")) # select proto
+    result = string.fromBytes(await conn.readLp(1024)) # read the first proto
+    result.removeSuffix("\n")
+    if result == p:
+      trace "selected protocol", protocol = result
+      break
 
 proc select*(m: MultistreamSelect,
              conn: Connection,
@@ -157,7 +158,7 @@ proc addHandler*[T: LPProtocol](m: MultistreamSelect,
                                 matcher: Matcher = nil) =
   ## register a protocol
   # TODO: This is a bug in chronicles,
-  # it break if I uncoment this line.
+  # it break if I uncomment this line.
   # Which is almost the same as the
   # one on the next override of addHandler
   #
diff --git a/libp2p/switch.nim b/libp2p/switch.nim
index 6858259..6192c6d 100644
--- a/libp2p/switch.nim
+++ b/libp2p/switch.nim
@@ -44,7 +44,7 @@ type
       ms*: MultistreamSelect
       identity*: Identify
       streamHandler*: StreamHandler
-      secureManagers*: OrderedTable[string, Secure]
+      secureManagers*: seq[Secure]
       pubSub*: Option[PubSub]
       dialedPubSubPeers: HashSet[string]
 
@@ -52,17 +52,14 @@ proc newNoPubSubException(): ref CatchableError {.inline.} =
   result = newException(NoPubSubException, "no pubsub provided!")
 
 proc secure(s: Switch, conn: Connection): Future[Connection] {.async, gcsafe.} =
-  ## secure the incoming connection
-
-  let managers = toSeq(s.secureManagers.keys)
-  if managers.len == 0:
+  if s.secureManagers.len <= 0:
     raise newException(CatchableError, "No secure managers registered!")
 
-  let manager = await s.ms.select(conn, toSeq(s.secureManagers.values).mapIt(it.codec))
+  let manager = await s.ms.select(conn, s.secureManagers.mapIt(it.codec))
   if manager.len == 0:
     raise newException(CatchableError, "Unable to negotiate a secure channel!")
 
-  result = await s.secureManagers[manager].secure(conn, true)
+  result = await s.secureManagers.filterIt(it.codec == manager)[0].secure(conn, true)
 
 proc identify(s: Switch, conn: Connection): Future[PeerInfo] {.async, gcsafe.} =
   ## identify the connection
@@ -194,7 +191,7 @@ proc upgradeIncoming(s: Switch, conn: Connection) {.async, gcsafe.} =
                        {.async, gcsafe, closure.} =
     try:
       trace "Securing connection"
-      let secure = s.secureManagers[proto]
+      let secure = s.secureManagers.filterIt(it.codec == proto)[0]
       let sconn = await secure.secure(conn, false)
       if sconn.isNil:
         return
@@ -218,8 +215,8 @@ proc upgradeIncoming(s: Switch, conn: Connection) {.async, gcsafe.} =
     try:
       if (await ms.select(conn)): # just handshake
         # add the secure handlers
-        for k in s.secureManagers.keys:
-          ms.addHandler(k, securedHandler)
+        for k in s.secureManagers:
+          ms.addHandler(k.codec, securedHandler)
 
         # handle secured connections
       await ms.handle(conn)
@@ -428,7 +425,7 @@ proc newSwitch*(peerInfo: PeerInfo,
   result.muxed = initTable[string, Muxer]()
   result.identity = identity
   result.muxers = muxers
-  result.secureManagers = initOrderedTable[string, Secure]()
+  result.secureManagers = @secureManagers
   result.dialedPubSubPeers = initHashSet[string]()
 
   let s = result # can't capture result
@@ -467,14 +464,10 @@ proc newSwitch*(peerInfo: PeerInfo,
         if not(isNil(stream)):
           await stream.close()
 
-  for proto in secureManagers:
-    trace "adding secure manager ", codec = proto.codec
-    result.secureManagers[proto.codec] = proto
-
-  if result.secureManagers.len == 0:
+  if result.secureManagers.len <= 0:
     # use plain text if no secure managers are provided
     warn "no secure managers, falling back to plain text", codec = PlainTextCodec
-    result.secureManagers[PlainTextCodec] = Secure(newPlainText())
+    result.secureManagers &= Secure(newPlainText())
 
   if pubSub.isSome:
     result.pubSub = pubSub
diff --git a/tests/testinterop.nim b/tests/testinterop.nim
index 61ffefe..c33d268 100644
--- a/tests/testinterop.nim
+++ b/tests/testinterop.nim
@@ -83,7 +83,6 @@ proc testPubSubDaemonPublish(gossip: bool = false,
     let smsg = cast[string](data)
     check smsg == pubsubData
     times.inc()
-    echo "TIMES ", times
     if times >= count and not finished:
       finished = true
 
@@ -108,7 +107,6 @@ proc testPubSubDaemonPublish(gossip: bool = false,
 
   await wait(publisher(), 5.minutes) # should be plenty of time
 
-  echo "HEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEE"
   result = true
   await nativeNode.stop()
   await allFutures(awaiters)
@@ -144,7 +142,6 @@ proc testPubSubNodePublish(gossip: bool = false,
     let smsg = cast[string](message.data)
     check smsg == pubsubData
     times.inc()
-    echo "TIMES ", times
     if times >= count and not finished:
       finished = true
     result = true # don't cancel subscription
@@ -356,7 +353,6 @@ suite "Interop":
           check line == test
           await conn.writeLp(cast[seq[byte]](test))
           count.inc()
-          echo "COUNT ", count
 
         testFuture.complete(count)
         await conn.close()