From 693880ad35afb1d2c58b97b5f514fd5d77967cac Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Sat, 13 Jul 2024 15:53:34 +0900
Subject: [PATCH] Use max_message_size and max_route_len for Sphinx encoding

---
 mixnet/node.py        |  4 +++-
 mixnet/sphinx.py      |  2 ++
 mixnet/test_sphinx.py | 27 +++++++++++++++++++++++++++
 mixnet/test_utils.py  |  5 +++--
 requirements.txt      |  4 ++--
 5 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/mixnet/node.py b/mixnet/node.py
index 06d36cc..d48e997 100644
--- a/mixnet/node.py
+++ b/mixnet/node.py
@@ -53,7 +53,9 @@ class Node:
         """
         A handler to process messages received via gossip channel
         """
-        sphinx_packet = SphinxPacket.from_bytes(msg)
+        sphinx_packet = SphinxPacket.from_bytes(
+            msg, self.global_config.max_mix_path_length
+        )
         result = await self.__process_sphinx_packet(sphinx_packet)
         match result:
             case SphinxPacket():
diff --git a/mixnet/sphinx.py b/mixnet/sphinx.py
index a3bab4a..7294af4 100644
--- a/mixnet/sphinx.py
+++ b/mixnet/sphinx.py
@@ -27,5 +27,7 @@ class SphinxPacketBuilder:
             message,
             route=[mixnode.sphinx_node() for mixnode in route],
             destination=dummy_destination.sphinx_node(),
+            max_route_length=global_config.max_mix_path_length,
+            max_plain_payload_size=global_config.max_message_size,
         )
         return (packet, route)
diff --git a/mixnet/test_sphinx.py b/mixnet/test_sphinx.py
index 3be3d22..60855f3 100644
--- a/mixnet/test_sphinx.py
+++ b/mixnet/test_sphinx.py
@@ -33,6 +33,33 @@ class TestSphinxPacketBuilder(TestCase):
         ).payload.recover_plain_playload()
         self.assertEqual(msg, recovered)
 
+    def test_max_message_size(self):
+        global_config, _, _ = init_mixnet_config(10, max_message_size=2000)
+        mix_path_length = global_config.max_mix_path_length
+
+        packet1, _ = SphinxPacketBuilder.build(
+            self.random_bytes(1500), global_config, mix_path_length
+        )
+        packet2, _ = SphinxPacketBuilder.build(
+            self.random_bytes(2000), global_config, mix_path_length
+        )
+        self.assertEqual(len(packet1.bytes()), len(packet2.bytes()))
+
+        msg = self.random_bytes(2001)
+        with self.assertRaises(ValueError):
+            _ = SphinxPacketBuilder.build(msg, global_config, mix_path_length)
+
+    def test_max_mix_path_length(self):
+        global_config, _, _ = init_mixnet_config(10, max_mix_path_length=2)
+        msg = self.random_bytes(global_config.max_message_size)
+
+        packet1, _ = SphinxPacketBuilder.build(msg, global_config, 1)
+        packet2, _ = SphinxPacketBuilder.build(msg, global_config, 2)
+        self.assertEqual(len(packet1.bytes()), len(packet2.bytes()))
+
+        with self.assertRaises(ValueError):
+            _ = SphinxPacketBuilder.build(msg, global_config, 3)
+
     @staticmethod
     def random_bytes(size: int) -> bytes:
         assert size >= 0
diff --git a/mixnet/test_utils.py b/mixnet/test_utils.py
index b552f54..63aa907 100644
--- a/mixnet/test_utils.py
+++ b/mixnet/test_utils.py
@@ -11,8 +11,9 @@ from mixnet.config import (
 
 def init_mixnet_config(
     num_nodes: int,
+    max_message_size: int = 512,
+    max_mix_path_length: int = 3,
 ) -> tuple[GlobalConfig, list[NodeConfig], dict[bytes, X25519PrivateKey]]:
-    max_mix_path_length = 3
     gossip_config = NomssipConfig(peering_degree=6)
     node_configs = [
         NodeConfig(X25519PrivateKey.generate(), max_mix_path_length, gossip_config)
@@ -26,7 +27,7 @@ def init_mixnet_config(
             ]
         ),
         transmission_rate_per_sec=3,
-        max_message_size=512,
+        max_message_size=max_message_size,
         max_mix_path_length=max_mix_path_length,
     )
     key_map = {
diff --git a/requirements.txt b/requirements.txt
index 84adb0e..fb7551f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,7 +3,7 @@ cffi==1.16.0
 cryptography==41.0.7
 numpy==1.26.3
 pycparser==2.21
-pysphinx==0.0.3
+pysphinx==0.0.5
 scipy==1.11.4
 black==23.12.1
 sympy==1.12
@@ -12,4 +12,4 @@ toml==0.10.2 # used for noir
 portalocker==2.8.2 # portable file locking
 keum==0.2.0 # for CL's use of more obscure curves
 poseidon-hash==0.1.4 # used as the algebraic hash in CL
-hypothesis==6.103.0
\ No newline at end of file
+hypothesis==6.103.0