remove message fragments

mixnet/config.py

@@ -49,19 +49,11 @@ class MixMembership:
 
     nodes: List[NodeInfo]
 
-    def generate_route(self, num_hops: int, last_mix: NodeInfo) -> list[NodeInfo]:
+    def generate_route(self, length: int) -> list[NodeInfo]:
         """
         Generate a mix route for a Sphinx packet.
-        The pre-selected mix_destination is used as a last mix node in the route,
-        so that associated packets can be merged together into a original message.
         """
-        return [*(self.choose() for _ in range(num_hops - 1)), last_mix]
-
-    def choose(self) -> NodeInfo:
-        """
-        Choose a mix node as a mix destination that will reconstruct a message from Sphinx packets.
-        """
-        return random.choice(self.nodes)
+        return [random.choice(self.nodes) for _ in range(length)]
 
 
 @dataclass

mixnet/node.py

@@ -4,15 +4,14 @@
 from typing import TypeAlias
 
 from pysphinx.sphinx import (
-    Payload,
     ProcessedFinalHopPacket,
     ProcessedForwardHopPacket,
     SphinxPacket,
 )
 
 from mixnet.config import GlobalConfig, NodeConfig
 from mixnet.nomssip import Nomssip
-from mixnet.packet import Fragment, MessageFlag, MessageReconstructor, PacketBuilder
+from mixnet.sphinx import SphinxPacketBuilder
 
 BroadcastChannel: TypeAlias = asyncio.Queue[bytes]
 
@@ -28,7 +27,6 @@ class Node:
     config: NodeConfig
     global_config: GlobalConfig
     nomssip: Nomssip
-    reconstructor: MessageReconstructor
     broadcast_channel: BroadcastChannel
 
     def __init__(self, config: NodeConfig, global_config: GlobalConfig):
@@ -42,18 +40,17 @@ def __init__(self, config: NodeConfig, global_config: GlobalConfig):
             ),
             self.__process_msg,
         )
-        self.reconstructor = MessageReconstructor()
         self.broadcast_channel = asyncio.Queue()
 
     @staticmethod
     def __calculate_message_size(global_config: GlobalConfig) -> int:
         """
         Calculate the actual message size to be gossiped, which depends on the maximum length of mix path.
         """
-        sample_packet, _ = PacketBuilder.build_real_packets(
+        sample_sphinx_packet, _ = SphinxPacketBuilder.build(
             bytes(1), global_config.membership, global_config.max_mix_path_length
-        )[0]
-        return len(sample_packet.bytes())
+        )
+        return len(sample_sphinx_packet.bytes())
 
     async def __process_msg(self, msg: bytes) -> None:
         """
@@ -83,24 +80,11 @@ async def __process_sphinx_packet(
                 case ProcessedForwardHopPacket():
                     return processed.next_packet
                 case ProcessedFinalHopPacket():
-                    return await self.__process_sphinx_payload(processed.payload)
+                    return processed.payload.recover_plain_playload()
         except ValueError:
             # Return nothing, if it cannot be unwrapped by the private key of this node.
             return None
 
-    async def __process_sphinx_payload(self, payload: Payload) -> bytes | None:
-        """
-        Process the Sphinx payload if possible
-        """
-        msg_with_flag = self.reconstructor.add(
-            Fragment.from_bytes(payload.recover_plain_playload())
-        )
-        if msg_with_flag is not None:
-            flag, msg = PacketBuilder.parse_msg_and_flag(msg_with_flag)
-            if flag == MessageFlag.MESSAGE_FLAG_REAL:
-                return msg
-        return None
-
     def connect(self, peer: Node):
         """
         Establish a duplex connection with a peer node.
@@ -117,9 +101,9 @@ async def send_message(self, msg: bytes):
         """
         # Here, we handle the case in which a msg is split into multiple Sphinx packets.
         # But, in practice, we expect a message to be small enough to fit in a single Sphinx packet.
-        for packet, _ in PacketBuilder.build_real_packets(
+        sphinx_packet, _ = SphinxPacketBuilder.build(
             msg,
             self.global_config.membership,
             self.config.mix_path_length,
-        ):
-            await self.nomssip.gossip(packet.bytes())
+        )
+        await self.nomssip.gossip(sphinx_packet.bytes())

mixnet/sphinx.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from typing import List, Tuple
+
+from pysphinx.payload import Payload
+from pysphinx.sphinx import SphinxPacket
+
+from mixnet.config import MixMembership, NodeInfo
+
+
+class SphinxPacketBuilder:
+    @staticmethod
+    def build(
+        message: bytes, membership: MixMembership, path_len: int
+    ) -> Tuple[SphinxPacket, List[NodeInfo]]:
+        if path_len <= 0:
+            raise ValueError("path_len must be greater than 0")
+        if len(message) > Payload.max_plain_payload_size():
+            raise ValueError("message is too long")
+
+        route = membership.generate_route(path_len)
+        # We don't need the destination (defined in the Loopix Sphinx spec)
+        # because the last mix will broadcast the fully unwrapped message.
+        # Later, we will optimize the Sphinx according to our requirements.
+        dummy_destination = route[-1]
+
+        packet = SphinxPacket.build(
+            message,
+            route=[mixnode.sphinx_node() for mixnode in route],
+            destination=dummy_destination.sphinx_node(),
+        )
+        return (packet, route)