refactor: exit node (#62)

* refactor: exit node * refactor: entry/exit will serialize embedded dest as 0 in not present
2026-01-09 23:08:09 -05:00 · 2025-07-25 15:10:10 -04:00
parent afe4237bf2
commit cead65d489
5 changed files with 168 additions and 134 deletions
--- a/mix/mix_protocol.nim
+++ b/mix/mix_protocol.nim
@@ -52,8 +52,20 @@ proc cryptoRandomInt(max: int): Result[int, string] =
  let value = cast[uint64](bytes)
  return ok(int(value mod uint64(max)))

+proc exitNodeIsDestination(
+    mixProto: MixProtocol, msg: MixMessage
+) {.async: (raises: [CancelledError]).} =
+  let exitConn = MixExitConnection.new(msg.message)
+  trace "Received: ", receiver = multiAddr, message = message
+  await mixProto.pHandler(exitConn, msg.codec)
+  try:
+    await exitConn.close()
+  except CatchableError as e:
+    error "Failed to close exit connection: ", err = e.msg
+  return
+
 proc handleMixNodeConnection(
-    mixProto: MixProtocol, conn: Connection, codec: string
+    mixProto: MixProtocol, conn: Connection
 ) {.async: (raises: [CancelledError]).} =
  var receivedBytes: seq[byte]
  try:
@@ -74,14 +86,17 @@ proc handleMixNodeConnection(
  # Process the packet
  let (multiAddr, _, mixPrivKey, _, _) = getMixNodeInfo(mixProto.mixNodeInfo)

-  let processedPktRes = processSphinxPacket(
-    receivedBytes, mixPrivKey, mixProto.tagManager, not codec.destIsExit
-  )
-  if processedPktRes.isErr:
-    error "Failed to process Sphinx packet", err = processedPktRes.error
+  let sphinxPacket = SphinxPacket.deserialize(receivedBytes).valueOr:
+    error "Sphinx packet deserialization error", err = error
+    mix_messages_error.inc(labelValues = ["Intermediate/Exit", "INVALID_SPHINX"])
+    return
+
+  let (nextHop, delay, processedPkt, status) = processSphinxPacket(
+    sphinxPacket, mixPrivKey, mixProto.tagManager
+  ).valueOr:
+    error "Failed to process Sphinx packet", err = error
    mix_messages_error.inc(labelValues = ["Intermediate/Exit", "INVALID_SPHINX"])
    return
-  let (nextHop, delay, processedPkt, status) = processedPktRes.get()

  case status
  of Exit:
@@ -105,18 +120,8 @@ proc handleMixNodeConnection(
    trace "Exit node - Received mix message: ",
      receiver = multiAddr, message = deserialized.message

-    if destIsExit(deserialized.codec):
-      let exitConn = MixExitConnection.new(deserialized.message)
-      trace "Received: ", receiver = multiAddr, message = message
-      await mixProto.pHandler(exitConn, deserialized.codec)
-      if exitConn != nil:
-        try:
-          await exitConn.close()
-        except CatchableError as e:
-          error "Failed to close exit connection: ", err = e.msg
-        return
-    elif nextHop != Hop() or delay != @[]:
-      error "Next hop and delay must be empty"
+    if nextHop == Hop() and delay == @[]:
+      await mixProto.exitNodeIsDestination(deserialized)
      return

    # Add delay
@@ -334,15 +339,17 @@ proc anonymizeLocalProtocolSend*(
    mix_messages_error.inc(labelValues = ["Entry", "NON_RECOVERABLE"])
    return

-  var destHop = Opt.none(Hop)
-  if not exitNodeIsDestination:
-    #Encode destination
-    let dest = $destMultiAddr & "/p2p/" & $destPeerId
-    let destAddrBytes = multiAddrToBytes(dest).valueOr:
-      error "Failed to convert multiaddress to bytes", err = error
-      mix_messages_error.inc(labelValues = ["Entry", "INVALID_DEST"])
-      return
-    destHop = Opt.some(Hop.init(destAddrBytes))
+  let destHop =
+    if not exitNodeIsDestination:
+      #Encode destination
+      let dest = $destMultiAddr & "/p2p/" & $destPeerId
+      let destAddrBytes = multiAddrToBytes(dest).valueOr:
+        error "Failed to convert multiaddress to bytes", err = error
+        mix_messages_error.inc(labelValues = ["Entry", "INVALID_DEST"])
+        return
+      Hop.init(destAddrBytes)
+    else:
+      Hop()

  # Wrap in Sphinx packet
  let sphinxPacket = wrapInSphinxPacket(
@@ -387,7 +394,7 @@ proc anonymizeLocalProtocolSend*(

 method init*(mixProtocol: MixProtocol) {.gcsafe, raises: [].} =
  proc handle(conn: Connection, proto: string) {.async: (raises: [CancelledError]).} =
-    await mixProtocol.handleMixNodeConnection(conn, proto)
+    await mixProtocol.handleMixNodeConnection(conn)

  mixProtocol.codecs = @[MixProtocolID]
  mixProtocol.handler = handle
--- a/mix/protocol.nim
+++ b/mix/protocol.nim
@@ -1,19 +1,10 @@
-import chronos, std/enumerate
-import
-  libp2p/[builders, protocols/ping, protocols/pubsub/gossipsub/types, stream/connection]
+import chronos, chronicles, std/enumerate
+import libp2p/[builders, stream/connection]

 type ProtocolHandler* = proc(conn: Connection, codec: string): Future[void] {.
  async: (raises: [CancelledError])
 .}

-# TODO: this is temporary while I attempt to extract protocol specific logic from mix
-func destIsExit*(proto: string): bool =
-  return
-    not (
-      proto == GossipSubCodec_12 or proto == GossipSubCodec_11 or
-      proto == GossipSubCodec_10
-    )
-
 method callHandler*(
    switch: Switch, conn: Connection, codec: string
 ): Future[void] {.base, async.} =
--- a/mix/serialization.nim
+++ b/mix/serialization.nim
@@ -50,20 +50,36 @@ type Hop* = object
  MultiAddress: seq[byte]

 proc init*(T: typedesc[Hop], multiAddress: seq[byte]): T =
-  T(MultiAddress: multiAddress)
+  T(
+    MultiAddress:
+      if multiAddress == newSeq[byte](addrSize):
+        @[]
+      else:
+        multiAddress
+  )

 proc getHop*(hop: Hop): seq[byte] =
  return hop.MultiAddress

 proc serialize*(hop: Hop): Result[seq[byte], string] =
-  if len(hop.MultiAddress) != addrSize:
+  if hop.MultiAddress.len == 0:
+    return ok(newSeq[byte](addrSize))
+  elif len(hop.MultiAddress) != addrSize:
    return err("MultiAddress must be exactly " & $addrSize & " bytes")
  return ok(hop.MultiAddress)

 proc deserialize*(T: typedesc[Hop], data: openArray[byte]): Result[T, string] =
  if len(data) != addrSize:
    return err("MultiAddress must be exactly " & $addrSize & " bytes")
-  return ok(T(MultiAddress: @data))
+  ok(
+    T(
+      MultiAddress:
+        if data == newSeq[byte](addrSize):
+          @[]
+        else:
+          @data
+    )
+  )

 type RoutingInfo* = object
  Addr: Hop
--- a/mix/sphinx.nim
+++ b/mix/sphinx.nim
@@ -111,13 +111,15 @@ proc generateRandomDelay(): seq[byte] =
  return toseq(delayBytes)
 ]#

+const paddingLength = (((t + 1) * (r - L)) + 2) * k
+
 # Function to compute betas, gammas, and deltas
 proc computeBetaGammaDelta(
    s: seq[seq[byte]],
    hop: openArray[Hop],
    msg: Message,
    delay: openArray[seq[byte]],
-    destHop: Opt[Hop],
+    destHop: Hop,
 ): Result[(seq[byte], seq[byte], seq[byte]), string] = # TODO: name tuples
  let sLen = s.len
  var
@@ -141,15 +143,8 @@ proc computeBetaGammaDelta(

    # Compute Beta and Gamma
    if i == sLen - 1:
-      var padding: seq[byte]
-      if destHop.isSome:
-        let destBytes = destHop.get().serialize().valueOr:
-            return err("Error in destination address serialization: " & error)
-        let paddingLength = (((t + 1) * (r - L)) + 2) * k
-        padding = destBytes & delay[i] & newSeq[byte](paddingLength)
-      else:
-        let paddingLength = (((t + 1) * (r - L)) + t + 2) * k
-        padding = newSeq[byte](paddingLength)
+      let destBytes = ?destHop.serialize()
+      let padding = destBytes & delay[i] & newSeq[byte](paddingLength)

      let aesRes = aes_ctr(beta_aes_key, beta_iv, padding).valueOr:
        return err("Error in aes: " & error)
@@ -183,7 +178,7 @@ proc wrapInSphinxPacket*(
    publicKeys: openArray[FieldElement],
    delay: openArray[seq[byte]],
    hop: openArray[Hop],
-    destHop: Opt[Hop],
+    destHop: Hop,
 ): Result[seq[byte], string] =
  # Compute alphas and shared secrets
  let (alpha_0, s) = computeAlpha(publicKeys).valueOr:
@@ -202,15 +197,8 @@ proc wrapInSphinxPacket*(
  return ok(serialized)

 proc processSphinxPacket*(
-    serSphinxPacket: seq[byte],
-    privateKey: FieldElement,
-    tm: var TagManager,
-    isDestEmbedded: bool,
+    sphinxPacket: SphinxPacket, privateKey: FieldElement, tm: var TagManager
 ): Result[(Hop, seq[byte], seq[byte], ProcessingStatus), string] = # TODO: named touple
-  # Deserialize the Sphinx packet
-  let sphinxPacket = SphinxPacket.deserialize(serSphinxPacket).valueOr:
-    return err("Sphinx packet deserialization error: " & error)
-
  let
    (header, payload) = sphinxPacket.getSphinxPacket()
    (alpha, beta, gamma) = getHeader(header)
@@ -260,33 +248,15 @@ proc processSphinxPacket*(
    return err("Error in aes: " & error)

  # Check if B has the required prefix for the original message
-  paddingLength =
-    if isDestEmbedded:
-      (((t + 1) * (r - L)) + 2) * k
-    else:
-      (((t + 1) * (r - L)) + t + 2) * k
-
  zeroPadding = newSeq[byte](paddingLength)

-  let bOffset =
-    if isDestEmbedded:
-      (t * k)
-    else:
-      0
-
-  if B[bOffset .. bOffset + paddingLength - 1] == zeroPadding:
+  if B[(t * k) .. (t * k) + paddingLength - 1] == zeroPadding:
    let msg = Message.deserialize(delta_prime).valueOr:
      return err("Message deserialization error: " & error)
-
    let content = msg.getContent()
-
-    if isDestEmbedded:
-      let hop = Hop.deserialize(B[0 .. addrSize - 1]).valueOr:
-        return err(error)
-      return
-        ok((hop, B[addrSize .. ((t * k) - 1)], content[0 .. messageSize - 1], Exit))
-    else:
-      return ok((Hop(), @[], content[0 .. messageSize - 1], Exit))
+    let hop = Hop.deserialize(B[0 .. addrSize - 1]).valueOr:
+      return err(error)
+    return ok((hop, B[addrSize .. ((t * k) - 1)], content[0 .. messageSize - 1], Exit))
  else:
    # Extract routing information from B
    let routingInfo = RoutingInfo.deserialize(B).valueOr:
--- a/tests/test_sphinx.nim
+++ b/tests/test_sphinx.nim
@@ -64,47 +64,66 @@ suite "Sphinx Tests":
  test "sphinx_wrap_and_process":
    let (message, privateKeys, publicKeys, delay, hops) = createDummyData()

-    let packetRes = wrapInSphinxPacket(message, publicKeys, delay, hops, Opt.none(Hop))
-    if packetRes.isErr:
-      error "Sphinx wrap error", err = packetRes.error
-    let packet = packetRes.get()
+    let packetBytesRes = wrapInSphinxPacket(message, publicKeys, delay, hops, Hop())
+    if packetBytesRes.isErr:
+      error "Sphinx wrap error", err = packetBytesRes.error
+      fail()
+    let packetBytes = packetBytesRes.get()

-    if packet.len != packetSize:
+    if packetBytes.len != packetSize:
      error "Packet length is not valid",
-        pkt_len = $(packet.len), expected_len = $packetSize
+        pkt_len = $(packetBytes.len), expected_len = $packetSize
      fail()

-    let res1 = processSphinxPacket(packet, privateKeys[0], tm, false)
+    let packetRes = SphinxPacket.deserialize(packetBytes)
+    if packetRes.isErr:
+      error "Sphinx wrap error", err = packetBytesRes.error
+      fail()
+    let packet = packetRes.get()
+
+    let res1 = processSphinxPacket(packet, privateKeys[0], tm)
    if res1.isErr:
      error "Error in Sphinx processing", err = res1.error
      fail()
-    let (address1, delay1, processedPacket1, status1) = res1.get()
+    let (address1, delay1, processedPacket1Bytes, status1) = res1.get()

    if status1 != Intermediate:
      error "Processing status should be Intermediate"
      fail()

-    if processedPacket1.len != packetSize:
+    if processedPacket1Bytes.len != packetSize:
      error "Packet length is not valid",
-        pkt_len = $(processedPacket1.len), expected_len = $packetSize
+        pkt_len = $(processedPacket1Bytes.len), expected_len = $packetSize
      fail()

-    let res2 = processSphinxPacket(processedPacket1, privateKeys[1], tm, false)
+    let processedPacket1Res = SphinxPacket.deserialize(processedPacket1Bytes)
+    if processedPacket1Res.isErr:
+      error "Sphinx wrap error", err = processedPacket1Res.error
+      fail()
+    let processedPacket1 = processedPacket1Res.get()
+
+    let res2 = processSphinxPacket(processedPacket1, privateKeys[1], tm)
    if res2.isErr:
      error "Error in Sphinx processing", err = res2.error
      fail()
-    let (address2, delay2, processedPacket2, status2) = res2.get()
+    let (address2, delay2, processedPacket2Bytes, status2) = res2.get()

    if status2 != Intermediate:
      error "Processing status should be Intermediate"
      fail()

-    if processedPacket2.len != packetSize:
+    if processedPacket2Bytes.len != packetSize:
      error "Packet length is not valid",
-        pkt_len = $(processedPacket2.len), expected_len = $packetSize
+        pkt_len = $(processedPacket2Bytes.len), expected_len = $packetSize
      fail()

-    let res3 = processSphinxPacket(processedPacket2, privateKeys[2], tm, false)
+    let processedPacket2Res = SphinxPacket.deserialize(processedPacket2Bytes)
+    if processedPacket2Res.isErr:
+      error "Sphinx wrap error", err = processedPacket2Res.error
+      fail()
+    let processedPacket2 = processedPacket2Res.get()
+
+    let res3 = processSphinxPacket(processedPacket2, privateKeys[2], tm)
    if res3.isErr:
      error "Error in Sphinx processing", err = res3.error
      fail()
@@ -122,7 +141,7 @@ suite "Sphinx Tests":
  test "sphinx_wrap_empty_public_keys":
    let (message, _, _, delay, _) = createDummyData()

-    let packetRes = wrapInSphinxPacket(message, @[], delay, @[], Opt.none(Hop))
+    let packetRes = wrapInSphinxPacket(message, @[], delay, @[], Hop())
    if packetRes.isOk:
      error "Expected Sphinx wrap error when public keys are empty, but got success"
      fail()
@@ -130,21 +149,27 @@ suite "Sphinx Tests":
  test "sphinx_process_invalid_mac":
    let (message, privateKeys, publicKeys, delay, hops) = createDummyData()

-    let packetRes = wrapInSphinxPacket(message, publicKeys, delay, hops, Opt.none(Hop))
-    if packetRes.isErr:
-      error "Sphinx wrap error", err = packetRes.error
-    let packet = packetRes.get()
+    let packetBytesRes = wrapInSphinxPacket(message, publicKeys, delay, hops, Hop())
+    if packetBytesRes.isErr:
+      error "Sphinx wrap error", err = packetBytesRes.error
+    let packetBytes = packetBytesRes.get()

-    if packet.len != packetSize:
+    if packetBytes.len != packetSize:
      error "Packet length is not valid",
-        pkt_len = $(packet.len), expected_len = $packetSize
+        pkt_len = $(packetBytes.len), expected_len = $packetSize
      fail()

    # Corrupt the MAC for testing
-    var tamperedPacket = packet
-    tamperedPacket[0] = packet[0] xor 0x01
+    var tamperedPacketBytes = packetBytes
+    tamperedPacketBytes[0] = packetBytes[0] xor 0x01

-    let res = processSphinxPacket(tamperedPacket, privateKeys[0], tm, false)
+    let tamperedPacketRes = SphinxPacket.deserialize(tamperedPacketBytes)
+    if tamperedPacketRes.isErr:
+      error "Sphinx wrap error", err = tamperedPacketRes.error
+      fail()
+    let tamperedPacket = tamperedPacketRes.get()
+
+    let res = processSphinxPacket(tamperedPacket, privateKeys[0], tm)
    if res.isErr:
      error "Error in Sphinx processing", err = res.error
      fail()
@@ -157,18 +182,24 @@ suite "Sphinx Tests":
  test "sphinx_process_duplicate_tag":
    let (message, privateKeys, publicKeys, delay, hops) = createDummyData()

-    let packetRes = wrapInSphinxPacket(message, publicKeys, delay, hops, Opt.none(Hop))
-    if packetRes.isErr:
-      error "Sphinx wrap error", err = packetRes.error
-    let packet = packetRes.get()
+    let packetBytesRes = wrapInSphinxPacket(message, publicKeys, delay, hops, Hop())
+    if packetBytesRes.isErr:
+      error "Sphinx wrap error", err = packetBytesRes.error
+    let packetBytes = packetBytesRes.get()

-    if packet.len != packetSize:
+    if packetBytes.len != packetSize:
      error "Packet length is not valid",
-        pkt_len = $(packet.len), expected_len = $packetSize
+        pkt_len = $(packetBytes.len), expected_len = $packetSize
      fail()

+    let packetRes = SphinxPacket.deserialize(packetBytes)
+    if packetRes.isErr:
+      error "Sphinx wrap error", err = packetRes.error
+      fail()
+    let packet = packetRes.get()
+
    # Process the packet twice to test duplicate tag handling
-    let res1 = processSphinxPacket(packet, privateKeys[0], tm, false)
+    let res1 = processSphinxPacket(packet, privateKeys[0], tm)
    if res1.isErr:
      error "Error in Sphinx processing", err = res1.error
      fail()
@@ -178,7 +209,7 @@ suite "Sphinx Tests":
      error "Processing status should be Intermediate"
      fail()

-    let res2 = processSphinxPacket(packet, privateKeys[0], tm, false)
+    let res2 = processSphinxPacket(packet, privateKeys[0], tm)
    if res2.isErr:
      error "Error in Sphinx processing", err = res2.error
      fail()
@@ -198,49 +229,68 @@ suite "Sphinx Tests":
        message[i] = byte(rand(256))
      let paddedMessage = padMessage(message, messageSize)

-      let packetRes = wrapInSphinxPacket(
-        Message.init(paddedMessage), publicKeys, delay, hops, Opt.none(Hop)
-      )
-      if packetRes.isErr:
-        error "Sphinx wrap error", err = packetRes.error
-      let packet = packetRes.get()
+      let packetResBytes =
+        wrapInSphinxPacket(Message.init(paddedMessage), publicKeys, delay, hops, Hop())
+      if packetResBytes.isErr:
+        error "Sphinx wrap error", err = packetResBytes.error
+      let packetBytes = packetResBytes.get()

-      if packet.len != packetSize:
+      if packetBytes.len != packetSize:
        error "Packet length is not valid",
-          pkt_len = $(packet.len), expected_len = $packetSize, msg_len = $messageSize
+          pkt_len = $(packetBytes.len),
+          expected_len = $packetSize,
+          msg_len = $messageSize
        fail()

-      let res1 = processSphinxPacket(packet, privateKeys[0], tm, false)
+      let packetRes = SphinxPacket.deserialize(packetBytes)
+      if packetRes.isErr:
+        error "Sphinx wrap error", err = packetRes.error
+        fail()
+      let packet = packetRes.get()
+
+      let res1 = processSphinxPacket(packet, privateKeys[0], tm)
      if res1.isErr:
        error "Error in Sphinx processing", err = res1.error
        fail()
-      let (address1, delay1, processedPacket1, status1) = res1.get()
+      let (address1, delay1, processedPacket1Bytes, status1) = res1.get()

      if status1 != Intermediate:
        error "Processing status should be Intermediate"
        fail()

-      if processedPacket1.len != packetSize:
+      if processedPacket1Bytes.len != packetSize:
        error "Packet length is not valid",
-          pkt_len = $(processedPacket1.len), expected_len = $packetSize
+          pkt_len = $(processedPacket1Bytes.len), expected_len = $packetSize
        fail()

-      let res2 = processSphinxPacket(processedPacket1, privateKeys[1], tm, false)
+      let processedPacket1Res = SphinxPacket.deserialize(processedPacket1Bytes)
+      if processedPacket1Res.isErr:
+        error "Sphinx wrap error", err = processedPacket1Res.error
+        fail()
+      let processedPacket1 = processedPacket1Res.get()
+
+      let res2 = processSphinxPacket(processedPacket1, privateKeys[1], tm)
      if res2.isErr:
        error "Error in Sphinx processing", err = res2.error
        fail()
-      let (address2, delay2, processedPacket2, status2) = res2.get()
+      let (address2, delay2, processedPacket2Bytes, status2) = res2.get()

      if status2 != Intermediate:
        error "Processing status should be Intermediate"
        fail()

-      if processedPacket2.len != packetSize:
+      if processedPacket2Bytes.len != packetSize:
        error "Packet length is not valid",
-          pkt_len = $(processedPacket2.len), expected_len = $packetSize
+          pkt_len = $(processedPacket2Bytes.len), expected_len = $packetSize
        fail()

-      let res3 = processSphinxPacket(processedPacket2, privateKeys[2], tm, false)
+      let processedPacket2Res = SphinxPacket.deserialize(processedPacket2Bytes)
+      if processedPacket2Res.isErr:
+        error "Sphinx wrap error", err = processedPacket2Res.error
+        fail()
+      let processedPacket2 = processedPacket2Res.get()
+
+      let res3 = processSphinxPacket(processedPacket2, privateKeys[2], tm)
      if res3.isErr:
        error "Error in Sphinx processing", err = res3.error
        fail()