From 4ace70d53b0b0e3b58cd3bead70b967d34bd03f3 Mon Sep 17 00:00:00 2001
From: diegomrsantos <diego@status.im>
Date: Wed, 25 Jan 2023 11:19:03 +0100
Subject: [PATCH] Connect is able to force a new connection (#849)

---
 libp2p/dial.nim      |  3 ++-
 libp2p/dialer.nim    | 47 +++++++++++++++++++++++++-------------------
 libp2p/switch.nim    |  5 +++--
 tests/testdialer.nim | 38 +++++++++++++++++++++++++++++++++++
 4 files changed, 70 insertions(+), 23 deletions(-)
 create mode 100644 tests/testdialer.nim

diff --git a/libp2p/dial.nim b/libp2p/dial.nim
index ff0ea0833..b4c205af7 100644
--- a/libp2p/dial.nim
+++ b/libp2p/dial.nim
@@ -27,7 +27,8 @@ method connect*(
   self: Dial,
   peerId: PeerId,
   addrs: seq[MultiAddress],
-  forceDial = false) {.async, base.} =
+  forceDial = false,
+  reuseConnection = true) {.async, base.} =
   ## connect remote peer without negotiating
   ## a protocol
   ##
diff --git a/libp2p/dialer.nim b/libp2p/dialer.nim
index c9adb12d6..d2ba4eeeb 100644
--- a/libp2p/dialer.nim
+++ b/libp2p/dialer.nim
@@ -147,11 +147,28 @@ proc dialAndUpgrade(
         if not isNil(result):
           return result
 
+proc tryReusingConnection(self: Dialer, peerId: PeerId): Future[Opt[Connection]] {.async.} =
+  var conn = self.connManager.selectConn(peerId)
+  if conn == nil:
+    return Opt.none(Connection)
+
+  if conn.atEof or conn.closed:
+    # This connection should already have been removed from the connection
+    # manager - it's essentially a bug that we end up here - we'll fail
+    # for now, hoping that this will clean themselves up later...
+    warn "dead connection in connection manager", conn
+    await conn.close()
+    raise newException(DialFailedError, "Zombie connection encountered")
+
+  trace "Reusing existing connection", conn, direction = $conn.dir
+  return Opt.some(conn)
+
 proc internalConnect(
   self: Dialer,
   peerId: Opt[PeerId],
   addrs: seq[MultiAddress],
-  forceDial: bool):
+  forceDial: bool,
+  reuseConnection = true):
   Future[Connection] {.async.} =
   if Opt.some(self.localPeerId) == peerId:
     raise newException(CatchableError, "can't dial self!")
@@ -161,24 +178,13 @@ proc internalConnect(
   try:
     await lock.acquire()
 
-    # Check if we have a connection already and try to reuse it
-    var conn =
-      if peerId.isSome: self.connManager.selectConn(peerId.get())
-      else: nil
-    if conn != nil:
-      if conn.atEof or conn.closed:
-        # This connection should already have been removed from the connection
-        # manager - it's essentially a bug that we end up here - we'll fail
-        # for now, hoping that this will clean themselves up later...
-        warn "dead connection in connection manager", conn
-        await conn.close()
-        raise newException(DialFailedError, "Zombie connection encountered")
-
-      trace "Reusing existing connection", conn, direction = $conn.dir
-      return conn
+    if peerId.isSome and reuseConnection:
+      let connOpt = await self.tryReusingConnection(peerId.get())
+      if connOpt.isSome:
+        return connOpt.get()
 
     let slot = self.connManager.getOutgoingSlot(forceDial)
-    conn =
+    let conn =
       try:
         await self.dialAndUpgrade(peerId, addrs)
       except CatchableError as exc:
@@ -207,15 +213,16 @@ method connect*(
   self: Dialer,
   peerId: PeerId,
   addrs: seq[MultiAddress],
-  forceDial = false) {.async.} =
+  forceDial = false,
+  reuseConnection = true) {.async.} =
   ## connect remote peer without negotiating
   ## a protocol
   ##
 
-  if self.connManager.connCount(peerId) > 0:
+  if self.connManager.connCount(peerId) > 0 and reuseConnection:
     return
 
-  discard await self.internalConnect(Opt.some(peerId), addrs, forceDial)
+  discard await self.internalConnect(Opt.some(peerId), addrs, forceDial, reuseConnection)
 
 method connect*(
   self: Dialer,
diff --git a/libp2p/switch.nim b/libp2p/switch.nim
index 851a84594..d28ae1b1a 100644
--- a/libp2p/switch.nim
+++ b/libp2p/switch.nim
@@ -148,10 +148,11 @@ method connect*(
   s: Switch,
   peerId: PeerId,
   addrs: seq[MultiAddress],
-  forceDial = false): Future[void] {.public.} =
+  forceDial = false,
+  reuseConnection = true): Future[void] {.public.} =
   ## Connects to a peer without opening a stream to it
 
-  s.dialer.connect(peerId, addrs, forceDial)
+  s.dialer.connect(peerId, addrs, forceDial, reuseConnection)
 
 method connect*(
   s: Switch,
diff --git a/tests/testdialer.nim b/tests/testdialer.nim
new file mode 100644
index 000000000..1008915d5
--- /dev/null
+++ b/tests/testdialer.nim
@@ -0,0 +1,38 @@
+# Nim-LibP2P
+# Copyright (c) 2023 Status Research & Development GmbH
+# Licensed under either of
+#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
+#  * MIT license ([LICENSE-MIT](LICENSE-MIT))
+# at your option.
+# This file may not be copied, modified, or distributed except according to
+# those terms.
+
+import std/options
+import chronos
+import unittest2
+import ../libp2p/[builders,
+                  switch]
+import ./helpers
+
+suite "Dialer":
+  teardown:
+    checkTrackers()
+
+  asyncTest "Connect forces a new connection":
+
+    let
+      src = newStandardSwitch()
+      dst = newStandardSwitch()
+
+    await dst.start()
+
+    await src.connect(dst.peerInfo.peerId, dst.peerInfo.addrs)
+    check src.connManager.connCount(dst.peerInfo.peerId) == 1
+
+    await src.connect(dst.peerInfo.peerId, dst.peerInfo.addrs)
+    check src.connManager.connCount(dst.peerInfo.peerId) == 1
+
+    await src.connect(dst.peerInfo.peerId, dst.peerInfo.addrs, true, false)
+    check src.connManager.connCount(dst.peerInfo.peerId) == 2
+
+    await allFutures(src.stop(), dst.stop())