From 3585d3cf867853366dade70fb235ad8896430115 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Thu, 27 Jun 2024 17:32:22 +0900
Subject: [PATCH] separate global config

---
 mixnet/config.py      |  8 +++++---
 mixnet/node.py        | 14 ++++++++------
 mixnet/test_node.py   |  6 ++----
 mixnet/test_packet.py |  4 ++--
 mixnet/test_utils.py  | 17 +++++++++++------
 5 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/mixnet/config.py b/mixnet/config.py
index e09983d..6360cc3 100644
--- a/mixnet/config.py
+++ b/mixnet/config.py
@@ -12,15 +12,17 @@ from pysphinx.sphinx import Node as SphinxNode
 
 
 @dataclass
-class MixnetConfig:
-    node_configs: List[NodeConfig]
+class GlobalConfig:
     membership: MixMembership
+    transmission_rate_per_sec: int  # Global Transmission Rate
+    # TODO: use this to make the size of Sphinx packet constant
+    max_mix_path_length: int
 
 
 @dataclass
 class NodeConfig:
     private_key: X25519PrivateKey
-    transmission_rate_per_sec: int  # Global Transmission Rate
+    mix_path_length: int  # TODO: use this when creating Sphinx packets
 
 
 @dataclass
diff --git a/mixnet/node.py b/mixnet/node.py
index 3fc5c7a..42fe6d2 100644
--- a/mixnet/node.py
+++ b/mixnet/node.py
@@ -11,7 +11,7 @@ from pysphinx.sphinx import (
     SphinxPacket,
 )
 
-from mixnet.config import MixMembership, NodeConfig
+from mixnet.config import GlobalConfig, NodeConfig
 from mixnet.packet import Fragment, MessageFlag, MessageReconstructor, PacketBuilder
 
 NetworkPacket: TypeAlias = "SphinxPacket | bytes"
@@ -22,14 +22,14 @@ BroadcastChannel: TypeAlias = "asyncio.Queue[bytes]"
 
 class Node:
     config: NodeConfig
-    membership: MixMembership
+    global_config: GlobalConfig
     mixgossip_channel: MixGossipChannel
     reconstructor: MessageReconstructor
     broadcast_channel: BroadcastChannel
 
-    def __init__(self, config: NodeConfig, membership: MixMembership):
+    def __init__(self, config: NodeConfig, global_config: GlobalConfig):
         self.config = config
-        self.membership = membership
+        self.global_config = global_config
         self.mixgossip_channel = MixGossipChannel(self.__process_sphinx_packet)
         self.reconstructor = MessageReconstructor()
         self.broadcast_channel = asyncio.Queue()
@@ -61,11 +61,13 @@ class Node:
         conn = asyncio.Queue()
         peer.mixgossip_channel.add_inbound(conn)
         self.mixgossip_channel.add_outbound(
-            MixOutboundConnection(conn, self.config.transmission_rate_per_sec)
+            MixOutboundConnection(conn, self.global_config.transmission_rate_per_sec)
         )
 
     async def send_message(self, msg: bytes):
-        for packet, _ in PacketBuilder.build_real_packets(msg, self.membership):
+        for packet, _ in PacketBuilder.build_real_packets(
+            msg, self.global_config.membership
+        ):
             await self.mixgossip_channel.gossip(packet)
 
 
diff --git a/mixnet/test_node.py b/mixnet/test_node.py
index 89ddf6a..9a67e25 100644
--- a/mixnet/test_node.py
+++ b/mixnet/test_node.py
@@ -9,10 +9,8 @@ from mixnet.test_utils import (
 
 class TestNode(IsolatedAsyncioTestCase):
     async def test_node(self):
-        config = init_mixnet_config(10)
-        nodes = [
-            Node(node_config, config.membership) for node_config in config.node_configs
-        ]
+        global_config, node_configs = init_mixnet_config(10)
+        nodes = [Node(node_config, global_config) for node_config in node_configs]
         for i, node in enumerate(nodes):
             node.connect(nodes[(i + 1) % len(nodes)])
 
diff --git a/mixnet/test_packet.py b/mixnet/test_packet.py
index 81e3a25..5cf8098 100644
--- a/mixnet/test_packet.py
+++ b/mixnet/test_packet.py
@@ -16,7 +16,7 @@ from mixnet.test_utils import init_mixnet_config
 
 class TestPacket(TestCase):
     def test_real_packet(self):
-        membership = init_mixnet_config(10).membership
+        membership = init_mixnet_config(10)[0].membership
         msg = self.random_bytes(3500)
         packets_and_routes = PacketBuilder.build_real_packets(msg, membership)
         self.assertEqual(4, len(packets_and_routes))
@@ -47,7 +47,7 @@ class TestPacket(TestCase):
         )
 
     def test_cover_packet(self):
-        membership = init_mixnet_config(10).membership
+        membership = init_mixnet_config(10)[0].membership
         msg = b"cover"
         packets_and_routes = PacketBuilder.build_drop_cover_packets(msg, membership)
         self.assertEqual(1, len(packets_and_routes))
diff --git a/mixnet/test_utils.py b/mixnet/test_utils.py
index 4ff1d4f..151a693 100644
--- a/mixnet/test_utils.py
+++ b/mixnet/test_utils.py
@@ -1,20 +1,25 @@
 from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
 
 from mixnet.config import (
+    GlobalConfig,
     MixMembership,
-    MixnetConfig,
     NodeConfig,
     NodeInfo,
 )
 
 
-def init_mixnet_config(num_nodes: int) -> MixnetConfig:
+def init_mixnet_config(num_nodes: int) -> tuple[GlobalConfig, list[NodeConfig]]:
     transmission_rate_per_sec = 3
+    max_mix_path_length = 3
     node_configs = [
-        NodeConfig(X25519PrivateKey.generate(), transmission_rate_per_sec)
+        NodeConfig(X25519PrivateKey.generate(), max_mix_path_length)
         for _ in range(num_nodes)
     ]
-    membership = MixMembership(
-        [NodeInfo(node_config.private_key) for node_config in node_configs]
+    global_config = GlobalConfig(
+        MixMembership(
+            [NodeInfo(node_config.private_key) for node_config in node_configs]
+        ),
+        transmission_rate_per_sec,
+        max_mix_path_length,
     )
-    return MixnetConfig(node_configs, membership)
+    return (global_config, node_configs)