refactor: processSphinxPacket and process replies and return in conn (#78)

Author: richard-ramos

examples/poc_resp_ping.nim

@@ -116,7 +116,7 @@ proc mixnetSimulation() {.async: (raises: [Exception]).} =
 
   let response = await pingProto[senderIndex].ping(conn)
 
-  await sleepAsync(1.seconds)
+  info "PING RESPONSE", response
 
   deleteNodeInfoFolder()
   deletePubInfoFolder()

mix/entry_connection.nim

@@ -1,5 +1,6 @@
 import hashes, chronos, stew/byteutils, results, chronicles
 import libp2p/stream/connection
+import libp2p/varint
 import ./mix_protocol
 import ./config
 from fragmentation import dataSize
@@ -44,27 +45,134 @@ type MixEntryConnection* = ref object of Connection
   mixDialer: MixDialer
   params: Opt[MixParameters]
 
+  incoming: AsyncQueue[seq[byte]]
+  incomingFut: Future[void]
+  replyReceivedFut: Future[void]
+  cached: seq[byte]
+
+method readOnce*(
+    s: MixEntryConnection, pbytes: pointer, nbytes: int
+): Future[int] {.async: (raises: [CancelledError, LPStreamError]), public.} =
+  if s.isEof:
+    raise newLPStreamEOFError()
+
+  try:
+    await s.replyReceivedFut
+    s.isEof = true
+    if s.cached.len == 0:
+      raise newLPStreamEOFError()
+  except CancelledError as exc:
+    raise exc
+  except LPStreamEOFError as exc:
+    raise exc
+  except CatchableError as exc:
+    raise (ref LPStreamError)(msg: "error in readOnce: " & exc.msg, parent: exc)
+
+  let toRead = min(nbytes, s.cached.len)
+  copyMem(pbytes, addr s.cached[0], toRead)
+  s.cached = s.cached[toRead ..^ 1]
+  return toRead
+
 method readExactly*(
-    self: MixEntryConnection, pbytes: pointer, nbytes: int
+    s: 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")
+  ## Waits for `nbytes` to be available, then read
+  ## them and return them
+  if s.atEof:
+    var ch: char
+    discard await s.readOnce(addr ch, 1)
+    raise newLPStreamEOFError()
+
+  if nbytes == 0:
+    return
+
+  logScope:
+    s
+    nbytes = nbytes
+    objName = s.objName
+
+  var pbuffer = cast[ptr UncheckedArray[byte]](pbytes)
+  var read = 0
+  while read < nbytes and not (s.atEof()):
+    read += await s.readOnce(addr pbuffer[read], nbytes - read)
+
+  if read == 0:
+    doAssert s.atEof()
+    trace "couldn't read all bytes, stream EOF", s, nbytes, read
+    # Re-readOnce to raise a more specific error than EOF
+    # Raise EOF if it doesn't raise anything(shouldn't happen)
+    discard await s.readOnce(addr pbuffer[read], nbytes - read)
+
+    raise newLPStreamEOFError()
+
+  if read < nbytes:
+    trace "couldn't read all bytes, incomplete data", s, nbytes, read
+    raise newLPStreamIncompleteError()
 
 method readLine*(
-    self: MixEntryConnection, limit = 0, sep = "\r\n"
+    s: MixEntryConnection, limit = 0, sep = "\r\n"
 ): Future[string] {.async: (raises: [CancelledError, LPStreamError]), public.} =
-  raise newException(LPStreamError, "readLine not implemented for MixEntryConnection")
+  ## Reads up to `limit` bytes are read, or a `sep` is found
+  # TODO replace with something that exploits buffering better
+  var lim = if limit <= 0: -1 else: limit
+  var state = 0
+
+  while true:
+    var ch: char
+    await readExactly(s, addr ch, 1)
+
+    if sep[state] == ch:
+      inc(state)
+      if state == len(sep):
+        break
+    else:
+      state = 0
+      if limit > 0:
+        let missing = min(state, lim - len(result) - 1)
+        result.add(sep[0 ..< missing])
+      else:
+        result.add(sep[0 ..< state])
+
+      result.add(ch)
+      if len(result) == lim:
+        break
 
 method readVarint*(
-    self: MixEntryConnection
+    conn: MixEntryConnection
 ): Future[uint64] {.async: (raises: [CancelledError, LPStreamError]), public.} =
-  raise newException(LPStreamError, "readVarint not implemented for MixEntryConnection")
+  var buffer: array[10, byte]
+
+  for i in 0 ..< len(buffer):
+    await conn.readExactly(addr buffer[i], 1)
+
+    var
+      varint: uint64
+      length: int
+    let res = PB.getUVarint(buffer.toOpenArray(0, i), length, varint)
+    if res.isOk():
+      return varint
+    if res.error() != VarintError.Incomplete:
+      break
+  if true: # can't end with a raise apparently
+    raise (ref InvalidVarintError)(msg: "Cannot parse varint")
 
 method readLp*(
-    self: MixEntryConnection, maxSize: int
+    s: MixEntryConnection, maxSize: int
 ): Future[seq[byte]] {.async: (raises: [CancelledError, LPStreamError]), public.} =
-  raise newException(LPStreamError, "readLp not implemented for MixEntryConnection")
+  ## read length prefixed msg, with the length encoded as a varint
+  let
+    length = await s.readVarint()
+    maxLen = uint64(if maxSize < 0: int.high else: maxSize)
+
+  if length > maxLen:
+    raise (ref MaxSizeError)(msg: "Message exceeds maximum length")
+
+  if length == 0:
+    return
+
+  var res = newSeqUninitialized[byte](length)
+  await s.readExactly(addr res[0], res.len)
+  res
 
 method write*(
     self: MixEntryConnection, msg: seq[byte]
@@ -112,6 +220,7 @@ proc shortLog*(self: MixEntryConnection): string {.raises: [].} =
 method closeImpl*(
     self: MixEntryConnection
 ): Future[void] {.async: (raises: [], raw: true).} =
+  self.incomingFut.cancelSoon()
   let fut = newFuture[void]()
   fut.complete()
   return fut
@@ -123,22 +232,6 @@ when defined(libp2p_agents_metrics):
   proc setShortAgent*(self: MixEntryConnection, shortAgent: string) =
     discard
 
-proc new*(
-    T: typedesc[MixEntryConnection],
-    srcMix: MixProtocol,
-    destination: Destination,
-    codec: string,
-    mixDialer: MixDialer,
-    params: Opt[MixParameters],
-): T =
-  let instance =
-    T(destination: destination, codec: codec, mixDialer: mixDialer, params: params)
-
-  when defined(libp2p_agents_metrics):
-    instance.shortAgent = connection.shortAgent
-
-  instance
-
 proc new*(
     T: typedesc[MixEntryConnection],
     srcMix: MixProtocol,
@@ -154,7 +247,20 @@ proc new*(
     else:
       0
 
-  var sendDialerFunc = proc(
+  var instance = T()
+  instance.destination = destination
+  instance.codec = codec
+  instance.params = Opt.some(params)
+
+  if expectReply:
+    instance.incoming = newAsyncQueue[seq[byte]]()
+    instance.replyReceivedFut = newFuture[void]()
+    let checkForIncoming = proc(): Future[void] {.async: (raises: [CancelledError]).} =
+      instance.cached = await instance.incoming.get()
+      instance.replyReceivedFut.complete()
+    instance.incomingFut = checkForIncoming()
+
+  instance.mixDialer = proc(
       msg: seq[byte], codec: string, dest: Destination
   ): Future[void] {.async: (raises: [CancelledError, LPStreamError]).} =
     try:
@@ -164,12 +270,17 @@ proc new*(
         else:
           (Opt.none(PeerId), Opt.some(MixDestination.init(dest.peerId, dest.address)))
 
-      await srcMix.anonymizeLocalProtocolSend(msg, codec, peerId, destination, surbs)
+      await srcMix.anonymizeLocalProtocolSend(
+        instance.incoming, 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, Opt.some(params))
+  when defined(libp2p_agents_metrics):
+    instance.shortAgent = connection.shortAgent
+
+  instance
 
 proc toConnection*(
     srcMix: MixProtocol,

mix/exit_layer.nim

@@ -79,7 +79,7 @@ proc reply(
     if not replyConn.isNil:
       await replyConn.close()
   try:
-    await replyConn.writeLp(response)
+    await replyConn.write(response)
   except LPStreamError as exc:
     error "could not reply", description = exc.msg
     mix_messages_error.inc(labelValues = ["ExitLayer", "REPLY_FAILED"])