feat: create surbs and send response back (#73)

Author: richard-ramos

.github/workflows/examples.yml

@@ -14,4 +14,5 @@ jobs:
        nimble c ./examples/poc_gossipsub.nim
        nimble c ./examples/poc_gossipsub_repeated_runs.nim
        nimble c ./examples/poc_noresp_ping.nim
+       nimble c ./examples/poc_resp_ping.nim
       nim-versions: '["version-2-0", "version-2-2", "devel"]'

examples/poc_gossipsub.nim

@@ -114,7 +114,7 @@ proc makeMixConnCb(mixProto: MixProtocol): CustomConnCreationProc =
       destAddr: Option[MultiAddress], destPeerId: PeerId, codec: string
   ): Connection {.gcsafe, raises: [].} =
     try:
-      return mixProto.toConnection(destPeerId, codec)
+      return mixProto.toConnection(destPeerId, codec).get()
     except CatchableError as e:
       error "Error during execution of MixEntryConnection callback: ", err = e.msg
       return nil

examples/poc_gossipsub_repeated_runs.nim

@@ -116,7 +116,7 @@ proc makeMixConnCb(mixProto: MixProtocol): CustomConnCreationProc =
       destAddr: Option[MultiAddress], destPeerId: PeerId, codec: string
   ): Connection {.gcsafe, raises: [].} =
     try:
-      return mixProto.toConnection(destPeerId, codec)
+      return mixProto.toConnection(destPeerId, codec).get()
     except CatchableError as e:
       error "Error during execution of MixEntryConnection callback: ", err = e.msg
       return nil

examples/poc_noresp_ping.nim

@@ -85,7 +85,7 @@ proc mixnetSimulation() {.async: (raises: [Exception]).} =
 
     let proto = MixProtocol.new(index, numberOfNodes, nodes[index]).valueOr:
       error "Mix protocol initialization failed", err = error
-      return
+      return # We'll fwd requests, so let's register how should the exit node behave
 
     mixProto.add(proto)
 
@@ -109,7 +109,9 @@ proc mixnetSimulation() {.async: (raises: [Exception]).} =
       nodes[receiverIndex].peerInfo.peerId, nodes[receiverIndex].peerInfo.addrs[0]
     ),
     NoRespPingCodec,
-  )
+  ).valueOr:
+    error "could not obtain connection", err = error
+    return
 
   discard await noRespPingProto[senderIndex].noRespPing(conn)
   await sleepAsync(1.seconds)

examples/poc_resp_ping.nim

@@ -0,0 +1,125 @@
+import chronicles, chronos, results, strutils
+import std/[enumerate, sysrand]
+import libp2p
+import libp2p/[crypto/secp, protocols/ping]
+import ../mix
+
+proc cryptoRandomInt(max: int): Result[int, string] =
+  if max == 0:
+    return err("Max cannot be zero.")
+  var bytes: array[8, byte]
+  discard urandom(bytes)
+  let value = cast[uint64](bytes)
+  return ok(int(value mod uint64(max)))
+
+proc createSwitch(libp2pPrivKey: SkPrivateKey, multiAddr: MultiAddress): Switch =
+  result = SwitchBuilder
+    .new()
+    .withPrivateKey(PrivateKey(scheme: Secp256k1, skkey: libp2pPrivKey))
+    .withAddress(multiAddr)
+    .withRng(crypto.newRng())
+    .withYamux()
+    .withTcpTransport()
+    .withNoise()
+    .build()
+
+# Set up nodes
+proc setUpNodes(numNodes: int): seq[Switch] =
+  # This is not actually GC-safe
+  {.gcsafe.}:
+    # Initialize mix nodes
+    discard initializeMixNodes(numNodes)
+
+    var nodes: seq[Switch] = @[]
+
+    for index, node in enumerate(mixNodes):
+      # Write public info of all mix nodes
+      let nodePubInfoRes = getMixPubInfoByIndex(index)
+      if nodePubInfoRes.isErr:
+        error "Get mix pub info by index error", err = nodePubInfoRes.error
+        continue
+      let nodePubInfo = nodePubInfoRes.get()
+
+      let writePubRes = writeMixPubInfoToFile(nodePubInfo, index)
+      if writePubRes.isErr:
+        error "Failed to write pub info to file", nodeIndex = index
+        continue
+
+      # Write info of all mix nodes
+      let writeNodeRes = writeMixNodeInfoToFile(node, index)
+      if writeNodeRes.isErr:
+        error "Failed to write mix info to file", nodeIndex = index
+        continue
+
+      # Extract private key and multiaddress
+      let (multiAddrStr, _, _, _, libp2pPrivKey) = getMixNodeInfo(node)
+
+      let multiAddr = MultiAddress.init(multiAddrStr.split("/p2p/")[0]).valueOr:
+        error "Failed to initialize MultiAddress", err = error
+        return
+
+      # Create switch
+      let switch = createSwitch(libp2pPrivKey, multiAddr)
+      if not switch.isNil:
+        nodes.add(switch)
+      else:
+        warn "Failed to set up node", nodeIndex = index
+
+    return nodes
+
+proc mixnetSimulation() {.async: (raises: [Exception]).} =
+  let
+    numberOfNodes = 10
+    nodes = setUpNodes(numberOfNodes)
+
+  var
+    mixProto: seq[MixProtocol] = @[]
+    pingProto: seq[Ping] = @[]
+
+  # Start nodes
+  for index, _ in enumerate(nodes):
+    pingProto.add(Ping.new())
+
+    let proto = MixProtocol.new(index, numberOfNodes, nodes[index]).valueOr:
+      error "Mix protocol initialization failed", err = error
+      return
+
+    # We'll fwd requests, so let's register how should the exit node will read responses
+    proto.registerFwdReadBehavior(PingCodec, readExactly(32))
+    
+    mixProto.add(proto)
+
+    nodes[index].mount(pingProto[index])
+    nodes[index].mount(mixProto[index])
+
+    await nodes[index].start()
+  await sleepAsync(1.seconds)
+
+  let cryptoRandomIntResult = cryptoRandomInt(numberOfNodes)
+  if cryptoRandomIntResult.isErr:
+    error "Failed to generate random number", err = cryptoRandomIntResult.error
+    return
+  let senderIndex = cryptoRandomIntResult.value
+  var receiverIndex = 0
+  if senderIndex < numberOfNodes - 1:
+    receiverIndex = senderIndex + 1
+
+  let conn = mixProto[senderIndex].toConnection(
+    Destination.forwardToAddr(
+      nodes[receiverIndex].peerInfo.peerId, nodes[receiverIndex].peerInfo.addrs[0]
+    ),
+    PingCodec,
+    Opt.some(MixParameters(expectReply: Opt.some(true), numSurbs: Opt.some(byte(1)))),
+  ).valueOr:
+    error "Could not obtain connection", err = error
+    return
+
+  let response = await pingProto[senderIndex].ping(conn)
+
+  await sleepAsync(1.seconds)
+
+  deleteNodeInfoFolder()
+  deletePubInfoFolder()
+
+when isMainModule:
+  waitFor(mixnetSimulation())

mix.nim

@@ -1,5 +1,7 @@
 import results
-import ./mix/[mix_protocol, mix_node, entry_connection]
+import chronos
+import libp2p
+import ./mix/[mix_protocol, mix_node, entry_connection, exit_layer]
 
 export results
 
@@ -21,3 +23,20 @@ export Destination
 export DestinationType
 export forwardToAddr
 export mixNode
+export MixParameters
+export fwdReadBehaviorCb
+export registerFwdReadBehavior
+
+proc readLp*(maxSize: int): fwdReadBehaviorCb =
+  return proc(
+      conn: Connection
+  ): Future[seq[byte]] {.async: (raises: [CancelledError, LPStreamError]).} =
+    await conn.readLp(maxSize)
+
+proc readExactly*(nBytes: int): fwdReadBehaviorCb =
+  return proc(
+      conn: Connection
+  ): Future[seq[byte]] {.async: (raises: [CancelledError, LPStreamError]).} =
+    let buf = newSeqUninitialized[byte](nBytes)
+    await conn.readExactly(addr buf[0], nBytes)
+    return buf

mix/config.nim

@@ -12,3 +12,8 @@ const
   packetSize* = 4608 # Total packet size (from spec)
   messageSize* = packetSize - headerSize - k # Size of the message itself
   payloadSize* = messageSize + k # Total payload size
+  surbSize* = headerSize + k + addrSize
+    # Size of a surb packet inside the message payload
+  surbLenSize* = 1 # Size of the field storing the number of surbs
+  surbIdLen* = 16 # Size of the identifier used when sending a message with surb
+  defaultSurbs* = uint8(4) # Default number of SURBs to send

mix/entry_connection.nim

@@ -1,6 +1,7 @@
 import hashes, chronos, stew/byteutils, results, chronicles
 import libp2p/stream/connection
 import ./mix_protocol
+import ./config
 from fragmentation import dataSize
 
 type
@@ -33,14 +34,20 @@ type MixDialer* = proc(
   msg: seq[byte], codec: string, destination: Destination
 ): Future[void] {.async: (raises: [CancelledError, LPStreamError], raw: true).}
 
+type MixParameters* = object
+  expectReply*: Opt[bool]
+  numSurbs*: Opt[uint8]
+
 type MixEntryConnection* = ref object of Connection
   destination: Destination
   codec: string
   mixDialer: MixDialer
+  params: Opt[MixParameters]
 
 method readExactly*(
     self: MixEntryConnection, pbytes: pointer, nbytes: int
 ): Future[void] {.async: (raises: [CancelledError, LPStreamError]), public.} =
+  await sleepAsync(10.minutes) # TODO: implement readExactly
   raise
     newException(LPStreamError, "readExactly not implemented for MixEntryConnection")
 
@@ -102,9 +109,6 @@ method writeLp*(
 proc shortLog*(self: MixEntryConnection): string {.raises: [].} =
   "[MixEntryConnection] Destination: " & $self.destination
 
-method initStream*(self: MixEntryConnection) =
-  discard
-
 method closeImpl*(
     self: MixEntryConnection
 ): Future[void] {.async: (raises: [], raw: true).} =
@@ -125,8 +129,10 @@ proc new*(
     destination: Destination,
     codec: string,
     mixDialer: MixDialer,
+    params: Opt[MixParameters],
 ): T =
-  let instance = T(destination: destination, codec: codec, mixDialer: mixDialer)
+  let instance =
+    T(destination: destination, codec: codec, mixDialer: mixDialer, params: params)
 
   when defined(libp2p_agents_metrics):
     instance.shortAgent = connection.shortAgent
@@ -138,7 +144,16 @@ proc new*(
     srcMix: MixProtocol,
     destination: Destination,
     codec: string,
+    params: Opt[MixParameters],
 ): T {.raises: [].} =
+  let params = params.get(MixParameters())
+  let expectReply = params.expectReply.get(false)
+  let surbs =
+    if expectReply:
+      params.numSurbs.get(defaultSurbs)
+    else:
+      0
+
   var sendDialerFunc = proc(
       msg: seq[byte], codec: string, dest: Destination
   ): Future[void] {.async: (raises: [CancelledError, LPStreamError]).} =
@@ -149,19 +164,28 @@ proc new*(
         else:
           (Opt.none(PeerId), Opt.some(MixDestination.init(dest.peerId, dest.address)))
 
-      await srcMix.anonymizeLocalProtocolSend(msg, codec, peerId, destination)
+      await srcMix.anonymizeLocalProtocolSend(msg, codec, peerId, destination, surbs)
     except CatchableError as e:
       error "Error during execution of anonymizeLocalProtocolSend: ", err = e.msg
     return
 
-  T.new(srcMix, destination, codec, sendDialerFunc)
+  T.new(srcMix, destination, codec, sendDialerFunc, Opt.some(params))
 
 proc toConnection*(
-    srcMix: MixProtocol, destination: Destination | PeerId, codec: string
-): Connection {.gcsafe, raises: [].} =
+    srcMix: MixProtocol,
+    destination: Destination | PeerId,
+    codec: string,
+    params: Opt[MixParameters] = Opt.none(MixParameters),
+): Result[Connection, string] {.gcsafe, raises: [].} =
   let dest =
     when destination is PeerId:
       Destination.mixNode(destination)
     else:
       destination
-  MixEntryConnection.new(srcMix, dest, codec)
+
+  if dest.kind == DestinationType.ForwardAddr and
+      params.get(MixParameters()).expectReply.get(false) and
+      not srcMix.hasFwdBehavior(codec):
+    return err("no forward behavior for codec")
+
+  ok(MixEntryConnection.new(srcMix, dest, codec, params))

mix/exit_connection.nim

@@ -4,6 +4,7 @@ from fragmentation import dataSize
 
 type MixExitConnection* = ref object of Connection
   message: seq[byte]
+  response: seq[byte]
 
 method join*(
     self: MixExitConnection
@@ -117,8 +118,7 @@ method readLp*(
 method write*(
     self: MixExitConnection, msg: seq[byte]
 ): Future[void] {.async: (raises: [CancelledError, LPStreamError], raw: true), public.} =
-  # TODO: dial back
-  discard
+  self.response.add(msg)
 
 proc write*(
     self: MixExitConnection, msg: string
@@ -148,7 +148,7 @@ method writeLp*(
   buf[0 ..< vbytes.len] = vbytes.toOpenArray(0, vbytes.len - 1)
   buf[vbytes.len ..< buf.len] = msg
 
-  # TODO: dial back
+  self.write(buf)
 
 method writeLp*(
     self: MixExitConnection, msg: string
@@ -171,6 +171,11 @@ method closeImpl*(
 func hash*(self: MixExitConnection): Hash =
   discard
 
+proc getResponse*(self: MixExitConnection): seq[byte] =
+  let r = self.response
+  self.response = @[]
+  return r
+
 proc new*(T: typedesc[MixExitConnection], message: seq[byte]): T =
   let instance = T(message: message)