From c360a4a65f3891f68d0a0e1c7df3a93f9334ffa3 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Thu, 11 Jul 2024 17:46:11 +0900
Subject: [PATCH] add max_message_size global param

---
 mixnet/config.py      | 3 ++-
 mixnet/node.py        | 6 ++++--
 mixnet/sphinx.py      | 9 ++++-----
 mixnet/test_sphinx.py | 2 +-
 mixnet/test_utils.py  | 6 +++---
 5 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/mixnet/config.py b/mixnet/config.py
index c1a6bf2..749038a 100644
--- a/mixnet/config.py
+++ b/mixnet/config.py
@@ -19,7 +19,8 @@ class GlobalConfig:
 
     membership: MixMembership
     transmission_rate_per_sec: int  # Global Transmission Rate
-    # TODO: use this to make the size of Sphinx packet constant
+    # TODO: use these two to make the size of Sphinx packet constant
+    max_message_size: int
     max_mix_path_length: int
 
 
diff --git a/mixnet/node.py b/mixnet/node.py
index f6b011f..06d36cc 100644
--- a/mixnet/node.py
+++ b/mixnet/node.py
@@ -43,7 +43,9 @@ class Node:
         Calculate the actual message size to be gossiped, which depends on the maximum length of mix path.
         """
         sample_sphinx_packet, _ = SphinxPacketBuilder.build(
-            bytes(1), global_config.membership, global_config.max_mix_path_length
+            bytes(global_config.max_message_size),
+            global_config,
+            global_config.max_mix_path_length,
         )
         return len(sample_sphinx_packet.bytes())
 
@@ -98,7 +100,7 @@ class Node:
         # But, in practice, we expect a message to be small enough to fit in a single Sphinx packet.
         sphinx_packet, _ = SphinxPacketBuilder.build(
             msg,
-            self.global_config.membership,
+            self.global_config,
             self.config.mix_path_length,
         )
         await self.nomssip.gossip(sphinx_packet.bytes())
diff --git a/mixnet/sphinx.py b/mixnet/sphinx.py
index c3a6bdf..a3bab4a 100644
--- a/mixnet/sphinx.py
+++ b/mixnet/sphinx.py
@@ -2,23 +2,22 @@ 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
+from mixnet.config import GlobalConfig, NodeInfo
 
 
 class SphinxPacketBuilder:
     @staticmethod
     def build(
-        message: bytes, membership: MixMembership, path_len: int
+        message: bytes, global_config: GlobalConfig, 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():
+        if len(message) > global_config.max_message_size:
             raise ValueError("message is too long")
 
-        route = membership.generate_route(path_len)
+        route = global_config.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.
diff --git a/mixnet/test_sphinx.py b/mixnet/test_sphinx.py
index 13e15de..3be3d22 100644
--- a/mixnet/test_sphinx.py
+++ b/mixnet/test_sphinx.py
@@ -15,7 +15,7 @@ class TestSphinxPacketBuilder(TestCase):
     def test_builder(self):
         global_config, _, key_map = init_mixnet_config(10)
         msg = self.random_bytes(500)
-        packet, route = SphinxPacketBuilder.build(msg, global_config.membership, 3)
+        packet, route = SphinxPacketBuilder.build(msg, global_config, 3)
         self.assertEqual(3, len(route))
 
         processed = packet.process(key_map[route[0].public_key.public_bytes_raw()])
diff --git a/mixnet/test_utils.py b/mixnet/test_utils.py
index 615aea7..b552f54 100644
--- a/mixnet/test_utils.py
+++ b/mixnet/test_utils.py
@@ -12,7 +12,6 @@ from mixnet.config import (
 def init_mixnet_config(
     num_nodes: int,
 ) -> tuple[GlobalConfig, list[NodeConfig], dict[bytes, X25519PrivateKey]]:
-    transmission_rate_per_sec = 3
     max_mix_path_length = 3
     gossip_config = NomssipConfig(peering_degree=6)
     node_configs = [
@@ -26,8 +25,9 @@ def init_mixnet_config(
                 for node_config in node_configs
             ]
         ),
-        transmission_rate_per_sec,
-        max_mix_path_length,
+        transmission_rate_per_sec=3,
+        max_message_size=512,
+        max_mix_path_length=max_mix_path_length,
     )
     key_map = {
         node_config.private_key.public_key().public_bytes_raw(): node_config.private_key