From 3df92e0d675845611d57eef433edbc7fa1e1ac89 Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Sat, 31 Aug 2019 11:58:49 -0600
Subject: [PATCH] wip: switch

---
 libp2p/protocol.nim  |  2 +-
 libp2p/switch.nim    | 39 ++++++++++++++++++++++++--------
 tests/testswitch.nim | 54 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 84 insertions(+), 11 deletions(-)
 create mode 100644 tests/testswitch.nim

diff --git a/libp2p/protocol.nim b/libp2p/protocol.nim
index d5627df..613398c 100644
--- a/libp2p/protocol.nim
+++ b/libp2p/protocol.nim
@@ -27,4 +27,4 @@ proc newProtocol*(p: typedesc[LPProtocol],
   result.peerInfo = peerInfo
   result.init()
 
-method init*(p: LPProtocol) {.base.} = discard
+method init*(p: LPProtocol) {.base, gcsafe.} = discard
diff --git a/libp2p/switch.nim b/libp2p/switch.nim
index 6ef0db5..7e86d5e 100644
--- a/libp2p/switch.nim
+++ b/libp2p/switch.nim
@@ -7,10 +7,10 @@
 ## This file may not be copied, modified, or distributed except according to
 ## those terms.
 
-import tables
+import tables, sequtils
 import chronos
 import connection, transport, stream, 
-       multistreamselect, protocol,
+       multistream, protocol,
        peerinfo, multiaddress
 
 type
@@ -26,24 +26,43 @@ proc newSwitch*(peerInfo: PeerInfo, transports: seq[Transport]): Switch =
   result.peerInfo = peerInfo
   result.ms = newMultistream()
   result.transports = transports
-  result.protocols = newSeq[LPProtocol]()
   result.connections = newTable[string, Connection]()
+  result.protocols = newSeq[LPProtocol]()
 
-proc dial*(s: Switch, peer: PeerInfo, proto: string = ""): Future[Connection] {.async.} = discard
+proc secure(s: Switch, conn: Connection) {.async, gcsafe.} = 
+  ## secure the incoming connection
+  discard
 
-proc mount*(s: Switch, protocol: LPProtocol) = discard
+proc processConn(s: Switch, conn: Connection, proto: string = "") {.async, gcsafe.} =
+  defer: await s.ms.handle(conn) # fire up the handler 
+  await s.secure(conn) # secure the connection
+  if proto.len > 0:
+    echo "SELECTING"
+    if not await s.ms.select(conn, proto):
+      raise newException(CatchableError, "Unable to select protocol: " & proto)
+
+proc dial*(s: Switch, peer: PeerInfo, proto: string = ""): Future[Connection] {.async.} = 
+  for t in s.transports: # for each transport
+    for a in peer.addrs: # for each address
+      if t.handles(a): # check if it can dial it
+        result = await t.dial(a)
+        await s.processConn(result, proto)
+
+proc mount*[T: LPProtocol](s: Switch, proto: T) = 
+  if isNil(proto.handler):
+    raise newException(CatchableError, 
+    "Protocol has to define a handle method or proc")
+
+  s.ms.addHandler(proto.codec, proto)
 
 proc start*(s: Switch) {.async.} = 
-  # TODO: how bad is it that this is a closure?
   proc handle(conn: Connection): Future[void] {.closure, gcsafe.} = 
-    s.ms.handle(conn)
+    s.processConn(conn)
 
   for t in s.transports: # for each transport
     for a in s.peerInfo.addrs:
       if t.handles(a): # check if it handles the multiaddr
         await t.listen(a, handle) # listen for incoming connections
-        break
 
 proc stop*(s: Switch) {.async.} = 
-  for c in s.connections.values:
-    await c.close()
+  await allFutures(s.transports.mapIt(it.close()))
diff --git a/tests/testswitch.nim b/tests/testswitch.nim
new file mode 100644
index 0000000..abeee12
--- /dev/null
+++ b/tests/testswitch.nim
@@ -0,0 +1,54 @@
+import unittest
+import chronos
+import ../libp2p/switch, ../libp2p/multistream,
+       ../libp2p/identify, ../libp2p/connection,
+       ../libp2p/transport, ../libp2p/tcptransport,
+       ../libp2p/multiaddress, ../libp2p/peerinfo,
+       ../libp2p/crypto/crypto, ../libp2p/peer,
+       ../libp2p/protocol
+
+const TestCodec = "/test/proto/1.0.0"
+
+type
+  TestProto = ref object of LPProtocol
+
+method init(p: TestProto) {.gcsafe.} =
+  proc handle(conn: Connection, proto: string) {.async, gcsafe.} = 
+    let msg = cast[string](await conn.readLp())
+    check "Hello!" == msg
+    await conn.writeLp("Hello!")
+
+  p.handler = handle
+
+suite "Switch":
+  test "e2e use switch": 
+    proc testSwitch(): Future[bool] {.async, gcsafe.} =
+      let ma1: MultiAddress = Multiaddress.init("/ip4/127.0.0.1/tcp/53370")
+      let ma2: MultiAddress = Multiaddress.init("/ip4/127.0.0.1/tcp/53371")
+
+      var peerInfo1, peerInfo2: PeerInfo
+      var switch1, switch2: Switch
+      proc createSwitch(ma: MultiAddress): (Switch, PeerInfo) =
+        let seckey = PrivateKey.random(RSA)
+        var peerInfo: PeerInfo
+        peerInfo.peerId = PeerID.init(seckey)
+        peerInfo.addrs.add(ma)
+        let switch = newSwitch(peerInfo, @[Transport(newTransport(TcpTransport))])
+        result = (switch, peerInfo)
+
+      (switch1, peerInfo1) = createSwitch(ma1)
+      let testProto = newProtocol(TestProto, peerInfo1)
+      switch1.mount(testProto)
+      await switch1.start()
+
+      (switch2, peerInfo2) = createSwitch(ma2)
+      await switch2.start()
+      let conn = await switch2.dial(peerInfo1, TestCodec)
+      await conn.writeLp("Hello!")
+      let msg = cast[string](await conn.readLp())
+      check "Hello!" == msg
+
+      result = true
+
+    check:
+      waitFor(testSwitch()) == true