From adfed3a6223d6f21bedcf90557cfb34ba09a417e Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Wed, 10 Jul 2024 09:28:50 +0900
Subject: [PATCH] use mix_path_length to make Sphinx packets

---
 mixnet/config.py      |  2 +-
 mixnet/node.py        |  6 ++++--
 mixnet/packet.py      | 15 +++++++++------
 mixnet/test_packet.py |  4 ++--
 4 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/mixnet/config.py b/mixnet/config.py
index d1f44b4..4525041 100644
--- a/mixnet/config.py
+++ b/mixnet/config.py
@@ -22,7 +22,7 @@ class GlobalConfig:
 @dataclass
 class NodeConfig:
     private_key: X25519PrivateKey
-    mix_path_length: int  # TODO: use this when creating Sphinx packets
+    mix_path_length: int
     gossip: GossipConfig
 
 
diff --git a/mixnet/node.py b/mixnet/node.py
index 717483f..ed26d35 100644
--- a/mixnet/node.py
+++ b/mixnet/node.py
@@ -35,7 +35,7 @@ class Node:
         self.broadcast_channel = asyncio.Queue()
 
         sample_packet, _ = PacketBuilder.build_real_packets(
-            bytes(1), global_config.membership
+            bytes(1), global_config.membership, self.global_config.max_mix_path_length
         )[0]
         self.packet_size = len(sample_packet.bytes())
 
@@ -102,7 +102,9 @@ class Node:
 
     async def send_message(self, msg: bytes):
         for packet, _ in PacketBuilder.build_real_packets(
-            msg, self.global_config.membership
+            msg,
+            self.global_config.membership,
+            self.config.mix_path_length,
         ):
             await self.mixgossip_channel.gossip(
                 Node.__build_msg(MsgType.REAL, packet.bytes())
diff --git a/mixnet/packet.py b/mixnet/packet.py
index 1ac833b..513b9f0 100644
--- a/mixnet/packet.py
+++ b/mixnet/packet.py
@@ -23,24 +23,27 @@ class MessageFlag(Enum):
 class PacketBuilder:
     @staticmethod
     def build_real_packets(
-        message: bytes, membership: MixMembership
+        message: bytes, membership: MixMembership, path_len: int
     ) -> List[Tuple[SphinxPacket, List[NodeInfo]]]:
         return PacketBuilder.__build_packets(
-            MessageFlag.MESSAGE_FLAG_REAL, message, membership
+            MessageFlag.MESSAGE_FLAG_REAL, message, membership, path_len
         )
 
     @staticmethod
     def build_drop_cover_packets(
-        message: bytes, membership: MixMembership
+        message: bytes, membership: MixMembership, path_len: int
     ) -> List[Tuple[SphinxPacket, List[NodeInfo]]]:
         return PacketBuilder.__build_packets(
-            MessageFlag.MESSAGE_FLAG_DROP_COVER, message, membership
+            MessageFlag.MESSAGE_FLAG_DROP_COVER, message, membership, path_len
         )
 
     @staticmethod
     def __build_packets(
-        flag: MessageFlag, message: bytes, membership: MixMembership
+        flag: MessageFlag, message: bytes, membership: MixMembership, path_len: int
     ) -> List[Tuple[SphinxPacket, List[NodeInfo]]]:
+        if path_len <= 0:
+            raise ValueError("path_len must be greater than 0")
+
         last_mix = membership.choose()
 
         msg_with_flag = flag.bytes() + message
@@ -50,7 +53,7 @@ class PacketBuilder:
 
         out = []
         for fragment in fragment_set.fragments:
-            route = membership.generate_route(3, last_mix)
+            route = membership.generate_route(path_len, last_mix)
             packet = SphinxPacket.build(
                 fragment.bytes(),
                 [mixnode.sphinx_node() for mixnode in route],
diff --git a/mixnet/test_packet.py b/mixnet/test_packet.py
index 77e1a03..453e648 100644
--- a/mixnet/test_packet.py
+++ b/mixnet/test_packet.py
@@ -19,7 +19,7 @@ class TestPacket(TestCase):
         global_config, _, key_map = init_mixnet_config(10)
         msg = self.random_bytes(3500)
         packets_and_routes = PacketBuilder.build_real_packets(
-            msg, global_config.membership
+            msg, global_config.membership, 3
         )
         self.assertEqual(4, len(packets_and_routes))
 
@@ -60,7 +60,7 @@ class TestPacket(TestCase):
         global_config, _, key_map = init_mixnet_config(10)
         msg = b"cover"
         packets_and_routes = PacketBuilder.build_drop_cover_packets(
-            msg, global_config.membership
+            msg, global_config.membership, 3
         )
         self.assertEqual(1, len(packets_and_routes))