Mixnet: integrate all the pieces together (#57)

mixnet/README.md

@@ -0,0 +1,21 @@
+# Mixnet Specification
+
+This is the executable specification of Mixnet, which can be used as a networking layer of the Nomos network.
+
+![](structure.png)
+
+## Public Components
+
+- [`mixnet.py`](mixnet.py): A public interface of the Mixnet layer, which can be used by upper layers
+- [`robustness.py`](robustness.py): A public interface of the Robustness layer, which can be on top of the Mixnet layer and used by upper layers
+
+## Private Components
+
+There are two primary components in the Mixnet layer.
+
+- [`client.py`](client.py): A mix client interface, which splits a message into Sphinx packets, sends packets to mix nodes, and receives messages via gossip. Also, this emits cover packets periodically.
+- [`node.py`](node.py): A mix node interface, which receives Sphinx packets from other mix nodes, processes packets, and forwards packets to other mix nodes. This works only when selected by the topology construction.
+
+Each component receives a new topology from the Robustness layer.
+
+There is no interaction between mix client and mix node components.

mixnet/client.py

@@ -1,76 +1,124 @@
 from __future__ import annotations
 
 import asyncio
+from contextlib import suppress
+from typing import Self
 
-from mixnet.mixnet import Mixnet
+from mixnet.config import MixnetConfig
 from mixnet.node import PacketQueue
 from mixnet.packet import PacketBuilder
 from mixnet.poisson import poisson_interval_sec
 
 
-async def mixclient_emitter(
-    mixnet: Mixnet,
-    emission_rate_per_min: int,  # Poisson rate parameter: lambda in the spec
-    redundancy: int,  # b in the spec
-    real_packet_queue: PacketQueue,
-    outbound_socket: PacketQueue,
-):
-    """
-    Emit packets at the Poisson emission_rate_per_min.
+class MixClient:
+    __config: MixnetConfig
 
-    If a real packet is scheduled to be sent, this thread sends the real packet to the mixnet,
-    and schedules redundant real packets to be emitted in the next turns.
+    __real_packet_queue: PacketQueue
+    __outbound_socket: PacketQueue
+    __task: asyncio.Task  # A reference just to prevent task from being garbage collected
 
-    If no real packet is not scheduled, this thread emits a cover packet according to the emission_rate_per_min.
-    """
+    @classmethod
+    async def new(
+        cls,
+        config: MixnetConfig,
+    ) -> Self:
+        self = cls()
+        self.__config = config
+        self.__real_packet_queue = asyncio.Queue()
+        self.__outbound_socket = asyncio.Queue()
+        self.__task = asyncio.create_task(self.__run())
+        return self
 
-    redundant_real_packet_queue: PacketQueue = asyncio.Queue()
+    def set_config(self, config: MixnetConfig) -> None:
+        """
+        Replace the old config with the new config received
 
-    emission_notifier_queue = asyncio.Queue()
-    _ = asyncio.create_task(
-        emission_notifier(emission_rate_per_min, emission_notifier_queue)
-    )
+        In real implementations, this method may be integrated in a long-running task.
+        Here in the spec, this method has been simplified as a setter, assuming the single-thread test environment.
+        """
+        self.__config = config
 
-    while True:
-        # Wait until the next emission time
-        _ = await emission_notifier_queue.get()
-        try:
-            await emit(
-                mixnet,
-                redundancy,
-                real_packet_queue,
-                redundant_real_packet_queue,
-                outbound_socket,
+    def get_config(self) -> MixnetConfig:
+        return self.__config
+
+    async def send_message(self, msg: bytes) -> None:
+        packets_and_routes = PacketBuilder.build_real_packets(
+            msg, self.__config.topology
+        )
+        for packet, route in packets_and_routes:
+            await self.__real_packet_queue.put((route[0].addr, packet))
+
+    def subscribe_messages(self) -> "asyncio.Queue[bytes]":
+        """
+        Subscribe messages, which went through mix nodes and were broadcasted via gossip
+        """
+        return asyncio.Queue()
+
+    @property
+    def outbound_socket(self) -> PacketQueue:
+        return self.__outbound_socket
+
+    async def __run(self):
+        """
+        Emit packets at the Poisson emission_rate_per_min.
+
+        If a real packet is scheduled to be sent, this thread sends the real packet to the mixnet,
+        and schedules redundant real packets to be emitted in the next turns.
+
+        If no real packet is not scheduled, this thread emits a cover packet according to the emission_rate_per_min.
+        """
+
+        redundant_real_packet_queue: PacketQueue = asyncio.Queue()
+
+        emission_notifier_queue = asyncio.Queue()
+        _ = asyncio.create_task(
+            self.__emission_notifier(
+                self.__config.emission_rate_per_min, emission_notifier_queue
             )
-        finally:
-            # Python convention: indicate that the previously enqueued task has been processed
-            emission_notifier_queue.task_done()
+        )
 
+        while True:
+            # Wait until the next emission time
+            _ = await emission_notifier_queue.get()
+            try:
+                await self.__emit(self.__config.redundancy, redundant_real_packet_queue)
+            finally:
+                # Python convention: indicate that the previously enqueued task has been processed
+                emission_notifier_queue.task_done()
 
-async def emit(
-    mixnet: Mixnet,
-    redundancy: int,  # b in the spec
-    real_packet_queue: PacketQueue,
-    redundant_real_packet_queue: PacketQueue,
-    outbound_socket: PacketQueue,
-):
-    if not redundant_real_packet_queue.empty():
-        addr, packet = redundant_real_packet_queue.get_nowait()
-        await outbound_socket.put((addr, packet))
-        return
+    async def __emit(
+        self,
+        redundancy: int,  # b in the spec
+        redundant_real_packet_queue: PacketQueue,
+    ):
+        if not redundant_real_packet_queue.empty():
+            addr, packet = redundant_real_packet_queue.get_nowait()
+            await self.__outbound_socket.put((addr, packet))
+            return
 
-    if not real_packet_queue.empty():
-        addr, packet = real_packet_queue.get_nowait()
-        # Schedule redundant real packets
-        for _ in range(redundancy - 1):
-            redundant_real_packet_queue.put_nowait((addr, packet))
-            await outbound_socket.put((addr, packet))
+        if not self.__real_packet_queue.empty():
+            addr, packet = self.__real_packet_queue.get_nowait()
+            # Schedule redundant real packets
+            for _ in range(redundancy - 1):
+                redundant_real_packet_queue.put_nowait((addr, packet))
+                await self.__outbound_socket.put((addr, packet))
 
-    packet, route = PacketBuilder.drop_cover(b"drop cover", mixnet).next()
-    await outbound_socket.put((route[0].addr, packet))
+        packets_and_routes = PacketBuilder.build_drop_cover_packets(
+            b"drop cover", self.__config.topology
+        )
+        # We have a for loop here, but we expect that the total num of packets is 1
+        # because the dummy message is short.
+        for packet, route in packets_and_routes:
+            await self.__outbound_socket.put((route[0].addr, packet))
 
+    async def __emission_notifier(
+        self, emission_rate_per_min: int, queue: asyncio.Queue
+    ):
+        while True:
+            await asyncio.sleep(poisson_interval_sec(emission_rate_per_min))
+            queue.put_nowait(None)
 
-async def emission_notifier(emission_rate_per_min: int, queue: asyncio.Queue):
-    while True:
-        await asyncio.sleep(poisson_interval_sec(emission_rate_per_min))
-        queue.put_nowait(None)
+    async def cancel(self) -> None:
+        self.__task.cancel()
+        with suppress(asyncio.CancelledError):
+            await self.__task

mixnet/config.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+from typing import List, TypeAlias
+
+from cryptography.hazmat.primitives.asymmetric.x25519 import (
+    X25519PrivateKey,
+    X25519PublicKey,
+)
+from pysphinx.node import Node
+
+from mixnet.bls import BlsPrivateKey, BlsPublicKey
+
+
+@dataclass
+class MixnetConfig:
+    emission_rate_per_min: int  # Poisson rate parameter: lambda
+    redundancy: int
+    delay_rate_per_min: int  # Poisson rate parameter: mu
+    topology: MixnetTopology
+
+
+@dataclass
+class MixnetTopology:
+    # In production, this can be a 1-D array, which is accessible by indexes.
+    # Here, we use a 2-D array for readability.
+    layers: List[List[MixNodeInfo]]
+
+    def generate_route(self, mix_destination: MixNodeInfo) -> list[MixNodeInfo]:
+        """
+        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.
+        """
+        route = [random.choice(layer) for layer in self.layers[:-1]]
+        route.append(mix_destination)
+        return route
+
+    def choose_mix_destination(self) -> MixNodeInfo:
+        """
+        Choose a mix node from the last mix layer as a mix destination
+        that will reconstruct a message from Sphinx packets.
+        """
+        return random.choice(self.layers[-1])
+
+
+# 32-byte that represents an IP address and a port of a mix node.
+NodeAddress: TypeAlias = bytes
+
+
+@dataclass
+class MixNodeInfo:
+    identity_private_key: BlsPrivateKey
+    encryption_private_key: X25519PrivateKey
+    addr: NodeAddress
+
+    def identity_public_key(self) -> BlsPublicKey:
+        return self.identity_private_key.get_g1()
+
+    def encryption_public_key(self) -> X25519PublicKey:
+        return self.encryption_private_key.public_key()
+
+    def sphinx_node(self) -> Node:
+        return Node(self.encryption_private_key, self.addr)

mixnet/mixnet.py

@@ -1,70 +1,51 @@
 from __future__ import annotations
 
-import random
-from dataclasses import dataclass
-from typing import List
+import asyncio
+from typing import Self
 
+from cryptography.hazmat.primitives.asymmetric.x25519 import (
+    X25519PrivateKey,
+)
+
+from mixnet.client import MixClient
+from mixnet.config import MixnetConfig
 from mixnet.node import MixNode
 
 
 class Mixnet:
-    __topology: MixnetTopology | None = None
+    __mixclient: MixClient
+    __mixnode: MixNode
 
-    def get_topology(self) -> MixnetTopology:
-        if self.__topology is None:
-            raise RuntimeError("topology is not set yet")
-        return self.__topology
+    @classmethod
+    async def new(
+        cls,
+        encryption_private_key: X25519PrivateKey,
+        config: MixnetConfig,
+    ) -> Self:
+        self = cls()
+        self.__mixclient = await MixClient.new(config)
+        self.__mixnode = await MixNode.new(encryption_private_key, config)
+        return self
 
-    def set_topology(self, topology: MixnetTopology) -> None:
+    async def publish_message(self, msg: bytes) -> None:
+        await self.__mixclient.send_message(msg)
+
+    def subscribe_messages(self) -> "asyncio.Queue[bytes]":
+        return self.__mixclient.subscribe_messages()
+
+    def set_config(self, config: MixnetConfig) -> None:
         """
-        Replace the old topology with the new topology received, and start establishing new network connections in background.
+        Replace the old config with the new config received.
 
-        In real implementations, this method should be a long-running task, accepting topologies periodically.
+        In real implementations, this method should be a long-running task, accepting configs periodically.
         Here in the spec, this method has been simplified as a setter, assuming the single-thread test environment.
         """
-        self.__topology = topology
-        self.__establish_connections()
+        self.__mixclient.set_config(config)
+        self.__mixnode.set_config(config)
 
-    def __establish_connections(self) -> None:
-        """
-        Establish network connections in advance based on the topology received.
+    def get_config(self) -> MixnetConfig:
+        return self.__mixclient.get_config()
 
-        This is just a preparation to forward subsequent packets as quickly as possible,
-        but this is not a strict requirement.
-
-        In real implementations, this should be a background task.
-        """
-        pass
-
-
-@dataclass
-class MixnetTopology:
-    # In production, this can be a 1-D array, which is accessible by indexes.
-    # Here, we use a 2-D array for readability.
-    layers: List[List[MixNode]]
-
-    def generate_route(self, mix_destination: MixNode) -> list[MixNode]:
-        """
-        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.
-        """
-        route = [random.choice(layer) for layer in self.layers[:-1]]
-        route.append(mix_destination)
-        return route
-
-    def choose_mix_destination(self) -> MixNode:
-        """
-        Choose a mix node from the last mix layer as a mix destination
-        that will reconstruct a message from Sphinx packets.
-        """
-        return random.choice(self.layers[-1])
-
-
-@dataclass
-class MixnetTopologySize:
-    num_layers: int
-    num_mixnodes_per_layer: int
-
-    def num_total_mixnodes(self) -> int:
-        return self.num_layers * self.num_mixnodes_per_layer
+    async def cancel(self) -> None:
+        await self.__mixclient.cancel()
+        await self.__mixnode.cancel()

mixnet/node.py

@@ -1,14 +1,12 @@
 from __future__ import annotations
 
 import asyncio
-from dataclasses import dataclass
-from typing import Tuple, TypeAlias
+from contextlib import suppress
+from typing import Self, Tuple, TypeAlias
 
 from cryptography.hazmat.primitives.asymmetric.x25519 import (
     X25519PrivateKey,
-    X25519PublicKey,
 )
-from pysphinx.node import Node
 from pysphinx.sphinx import (
     Payload,
     ProcessedFinalHopPacket,
@@ -17,51 +15,16 @@ from pysphinx.sphinx import (
     UnknownHeaderTypeError,
 )
 
-from mixnet.bls import BlsPrivateKey, BlsPublicKey
+from mixnet.config import MixnetConfig, NodeAddress
 from mixnet.poisson import poisson_interval_sec
 
-NodeId: TypeAlias = BlsPublicKey
-# 32-byte that represents an IP address and a port of a mix node.
-NodeAddress: TypeAlias = bytes
-
 PacketQueue: TypeAlias = "asyncio.Queue[Tuple[NodeAddress, SphinxPacket]]"
 PacketPayloadQueue: TypeAlias = (
     "asyncio.Queue[Tuple[NodeAddress, SphinxPacket | Payload]]"
 )
 
 
-@dataclass
 class MixNode:
-    identity_private_key: BlsPrivateKey
-    encryption_private_key: X25519PrivateKey
-    addr: NodeAddress
-
-    def identity_public_key(self) -> BlsPublicKey:
-        return self.identity_private_key.get_g1()
-
-    def encryption_public_key(self) -> X25519PublicKey:
-        return self.encryption_private_key.public_key()
-
-    def sphinx_node(self) -> Node:
-        return Node(self.encryption_private_key, self.addr)
-
-    def start(
-        self,
-        delay_rate_per_min: int,  # Poisson rate parameter: mu
-        inbound_socket: PacketQueue,
-        outbound_socket: PacketPayloadQueue,
-    ) -> asyncio.Task:
-        return asyncio.create_task(
-            MixNodeRunner(
-                self.encryption_private_key,
-                delay_rate_per_min,
-                inbound_socket,
-                outbound_socket,
-            ).run()
-        )
-
-
-class MixNodeRunner:
     """
     A class handling incoming packets with delays
 
@@ -69,19 +32,30 @@ class MixNodeRunner:
     in order to define the MixNode as a simple dataclass for clarity.
     """
 
-    def __init__(
+    __config: MixnetConfig
+
+    inbound_socket: PacketQueue
+    outbound_socket: PacketPayloadQueue
+    __task: asyncio.Task  # A reference just to prevent task from being garbage collected
+
+    @classmethod
+    async def new(
+        cls,
+        encryption_private_key: X25519PrivateKey,
+        config: MixnetConfig,
+    ) -> Self:
+        self = cls()
+        self.__config = config
+        self.__establish_connections()
+        self.inbound_socket = asyncio.Queue()
+        self.outbound_socket = asyncio.Queue()
+        self.__task = asyncio.create_task(self.__run(encryption_private_key))
+        return self
+
+    async def __run(
         self,
         encryption_private_key: X25519PrivateKey,
-        delay_rate_per_min: int,  # Poisson rate parameter: mu
-        inbound_socket: PacketQueue,
-        outbound_socket: PacketPayloadQueue,
     ):
-        self.encryption_private_key = encryption_private_key
-        self.delay_rate_per_min = delay_rate_per_min
-        self.inbound_socket = inbound_socket
-        self.outbound_socket = outbound_socket
-
-    async def run(self):
         """
         Read SphinxPackets from inbound socket and spawn a thread for each packet to process it.
 
@@ -95,17 +69,19 @@ class MixNodeRunner:
         while True:
             _, packet = await self.inbound_socket.get()
             task = asyncio.create_task(
-                self.process_packet(
-                    packet,
+                self.__process_packet(
+                    packet, encryption_private_key, self.__config.delay_rate_per_min
                 )
             )
             self.tasks.add(task)
             # To discard the task from the set automatically when it is done.
             task.add_done_callback(self.tasks.discard)
 
-    async def process_packet(
+    async def __process_packet(
         self,
         packet: SphinxPacket,
+        encryption_private_key: X25519PrivateKey,
+        delay_rate_per_min: int,  # Poisson rate parameter: mu
     ):
         """
         Process a single packet with a delay that follows exponential distribution,
@@ -113,10 +89,10 @@ class MixNodeRunner:
 
         This thread is a single server (worker) in a M/M/inf queue that MixNodeRunner approximates.
         """
-        delay_sec = poisson_interval_sec(self.delay_rate_per_min)
+        delay_sec = poisson_interval_sec(delay_rate_per_min)
         await asyncio.sleep(delay_sec)
 
-        processed = packet.process(self.encryption_private_key)
+        processed = packet.process(encryption_private_key)
         match processed:
             case ProcessedForwardHopPacket():
                 await self.outbound_socket.put(
@@ -128,3 +104,31 @@ class MixNodeRunner:
                 )
             case _:
                 raise UnknownHeaderTypeError
+
+    def set_config(self, config: MixnetConfig) -> None:
+        """
+        Replace the old config with the new config received.
+        If topology has been changed, start establishing new network connections in background.
+
+        In real implementations, this method may be integrated in a long-running task.
+        Here in the spec, this method has been simplified as a setter, assuming the single-thread test environment.
+        """
+        if self.__config.topology != config.topology:
+            self.__establish_connections()
+        self.__config = config
+
+    def __establish_connections(self) -> None:
+        """
+        Establish network connections in advance based on the topology received.
+
+        This is just a preparation to forward subsequent packets as quickly as possible,
+        but this is not a strict requirement.
+
+        In real implementations, this should be a background task.
+        """
+        pass
+
+    async def cancel(self) -> None:
+        self.__task.cancel()
+        with suppress(asyncio.CancelledError):
+            await self.__task

mixnet/packet.py

@@ -4,12 +4,12 @@ import uuid
 from dataclasses import dataclass
 from enum import Enum
 from itertools import batched
-from typing import Dict, Iterator, List, Self, Tuple, TypeAlias
+from typing import Dict, List, Self, Tuple, TypeAlias
 
 from pysphinx.payload import Payload
 from pysphinx.sphinx import SphinxPacket
 
-from mixnet.mixnet import Mixnet, MixNode
+from mixnet.config import MixnetTopology, MixNodeInfo
 
 
 class MessageFlag(Enum):
@@ -21,15 +21,26 @@ class MessageFlag(Enum):
 
 
 class PacketBuilder:
-    iter: Iterator[Tuple[SphinxPacket, List[MixNode]]]
+    @staticmethod
+    def build_real_packets(
+        message: bytes, topology: MixnetTopology
+    ) -> List[Tuple[SphinxPacket, List[MixNodeInfo]]]:
+        return PacketBuilder.__build_packets(
+            MessageFlag.MESSAGE_FLAG_REAL, message, topology
+        )
 
-    def __init__(
-        self,
-        flag: MessageFlag,
-        message: bytes,
-        mixnet: Mixnet,
-    ):
-        topology = mixnet.get_topology()
+    @staticmethod
+    def build_drop_cover_packets(
+        message: bytes, topology: MixnetTopology
+    ) -> List[Tuple[SphinxPacket, List[MixNodeInfo]]]:
+        return PacketBuilder.__build_packets(
+            MessageFlag.MESSAGE_FLAG_DROP_COVER, message, topology
+        )
+
+    @staticmethod
+    def __build_packets(
+        flag: MessageFlag, message: bytes, topology: MixnetTopology
+    ) -> List[Tuple[SphinxPacket, List[MixNodeInfo]]]:
         destination = topology.choose_mix_destination()
 
         msg_with_flag = flag.bytes() + message
@@ -37,7 +48,7 @@ class PacketBuilder:
         # If encryption is needed, a shared secret must be appended in front of the message along with the MessageFlag.
         fragment_set = FragmentSet(msg_with_flag)
 
-        packets_and_routes = []
+        out = []
         for fragment in fragment_set.fragments:
             route = topology.generate_route(destination)
             packet = SphinxPacket.build(
@@ -45,20 +56,9 @@ class PacketBuilder:
                 [mixnode.sphinx_node() for mixnode in route],
                 destination.sphinx_node(),
             )
-            packets_and_routes.append((packet, route))
+            out.append((packet, route))
 
-        self.iter = iter(packets_and_routes)
-
-    @classmethod
-    def real(cls, message: bytes, mixnet: Mixnet) -> Self:
-        return cls(MessageFlag.MESSAGE_FLAG_REAL, message, mixnet)
-
-    @classmethod
-    def drop_cover(cls, message: bytes, mixnet: Mixnet) -> Self:
-        return cls(MessageFlag.MESSAGE_FLAG_DROP_COVER, message, mixnet)
-
-    def next(self) -> Tuple[SphinxPacket, List[MixNode]]:
-        return next(self.iter)
+        return out
 
     @staticmethod
     def parse_msg_and_flag(data: bytes) -> Tuple[MessageFlag, bytes]:

mixnet/robustness.py

@@ -1,14 +1,11 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import List
 
+from mixnet.config import MixnetConfig, MixnetTopology, MixNodeInfo
 from mixnet.fisheryates import FisherYates
-from mixnet.mixnet import (
-    Mixnet,
-    MixnetTopology,
-    MixnetTopologySize,
-)
-from mixnet.node import MixNode
+from mixnet.mixnet import Mixnet
 
 
 class Robustness:
@@ -27,37 +24,77 @@ class Robustness:
 
     def __init__(
         self,
-        mixnode_candidates: List[MixNode],
-        mixnet_topology_size: MixnetTopologySize,
+        config: RobustnessConfig,
         mixnet: Mixnet,
     ) -> None:
-        assert mixnet_topology_size.num_total_mixnodes() <= len(mixnode_candidates)
-        self.mixnode_candidates = mixnode_candidates
-        self.mixnet_topology_size = mixnet_topology_size
-        self.mixnet = mixnet
+        self.__config = config
+        self.__mixnet = mixnet
 
     def set_entropy(self, entropy: bytes) -> None:
         """
         Given a entropy received, build a new topology and send it to mixnet.
+        In v1, this doesn't change any mixnet config except topology.
 
         In real implementations, this method should be a long-running task, consuming entropy periodically.
         Here in the spec, this method has been simplified as a setter, assuming the single-thread test environment.
         """
-        topology = self.build_topology(entropy)
-        self.mixnet.set_topology(topology)
+        self.__config.mixnet.mixnet_layer_config.topology = self.build_topology(
+            self.__config.mixnet.mixnode_candidates,
+            self.__config.mixnet.topology_size,
+            entropy,
+        )
+        self.__mixnet.set_config(self.__config.mixnet.mixnet_layer_config)
 
-    def build_topology(self, entropy: bytes) -> MixnetTopology:
+    @staticmethod
+    def build_topology(
+        mixnode_candidates: List[MixNodeInfo],
+        mixnet_topology_size: MixnetTopologySize,
+        entropy: bytes,
+    ) -> MixnetTopology:
         """
         Build a new topology deterministically using an entropy and a given set of candidates.
         """
-        shuffled = FisherYates.shuffle(self.mixnode_candidates, entropy)
-        sampled = shuffled[: self.mixnet_topology_size.num_total_mixnodes()]
+        shuffled = FisherYates.shuffle(mixnode_candidates, entropy)
+        sampled = shuffled[: mixnet_topology_size.num_total_mixnodes()]
 
         layers = []
-        for layer_id in range(self.mixnet_topology_size.num_layers):
-            start = layer_id * self.mixnet_topology_size.num_mixnodes_per_layer
-            layer = sampled[
-                start : start + self.mixnet_topology_size.num_mixnodes_per_layer
-            ]
+        for layer_id in range(mixnet_topology_size.num_layers):
+            start = layer_id * mixnet_topology_size.num_mixnodes_per_layer
+            layer = sampled[start : start + mixnet_topology_size.num_mixnodes_per_layer]
             layers.append(layer)
         return MixnetTopology(layers)
+
+
+@dataclass
+class RobustnessConfig:
+    """In v1, the robustness layer manages configs only for the mixnet layer."""
+
+    mixnet: RobustnessMixnetConfig
+
+
+class RobustnessMixnetConfig:
+    """
+    Configurations for the mixnet layer
+    These configurations are meant to be changed over time according to other parameters from other layers (e.g. consensus).
+    """
+
+    def __init__(
+        self,
+        mixnode_candidates: List[MixNodeInfo],
+        mixnet_topology_size: MixnetTopologySize,
+        mixnet_layer_config: MixnetConfig,
+    ) -> None:
+        assert mixnet_topology_size.num_total_mixnodes() <= len(mixnode_candidates)
+        self.mixnode_candidates = mixnode_candidates
+        self.topology_size = mixnet_topology_size
+        # A config to be injected to the mixnet layer whenever it is updated
+        self.mixnet_layer_config = mixnet_layer_config
+
+
+@dataclass
+class MixnetTopologySize:
+    num_layers: int
+    num_mixnodes_per_layer: int
+
+    def num_total_mixnodes(self) -> int:
+        return self.num_layers * self.num_mixnodes_per_layer

mixnet/test_client.py

@@ -1,58 +1,45 @@
-import asyncio
 from datetime import datetime
+from unittest import IsolatedAsyncioTestCase
 
 import numpy
 
-from mixnet.client import mixclient_emitter
-from mixnet.node import PacketQueue
-from mixnet.packet import PacketBuilder
+from mixnet.client import MixClient
 from mixnet.poisson import poisson_mean_interval_sec
-from mixnet.test_mixnet import TestMixnet
-from mixnet.test_utils import with_test_timeout
+from mixnet.test_utils import (
+    init_robustness_mixnet_config,
+    with_test_timeout,
+)
 from mixnet.utils import random_bytes
 
 
-class TestMixClient(TestMixnet):
+class TestMixClient(IsolatedAsyncioTestCase):
     @with_test_timeout(100)
-    async def test_mixclient_emitter(self):
-        mixnet, _ = self.init()
-        real_packet_queue: PacketQueue = asyncio.Queue()
-        outbound_socket: PacketQueue = asyncio.Queue()
+    async def test_mixclient(self):
+        config = init_robustness_mixnet_config().mixnet_layer_config
+        config.emission_rate_per_min = 30
+        config.redundancy = 3
 
-        emission_rate_per_min = 30
-        redundancy = 3
-        _ = asyncio.create_task(
-            mixclient_emitter(
-                mixnet,
-                emission_rate_per_min,
-                redundancy,
-                real_packet_queue,
-                outbound_socket,
+        mixclient = await MixClient.new(config)
+        try:
+            # Send a 3500-byte msg, expecting that it is split into at least two packets
+            await mixclient.send_message(random_bytes(3500))
+
+            # Calculate intervals between packet emissions from the mix client
+            intervals = []
+            ts = datetime.now()
+            for _ in range(30):
+                _ = await mixclient.outbound_socket.get()
+                now = datetime.now()
+                intervals.append((now - ts).total_seconds())
+                ts = now
+
+            # Check if packets were emitted at the Poisson emission_rate
+            # If emissions follow the Poisson distribution with a rate `lambda`,
+            # a mean interval between emissions must be `1/lambda`.
+            self.assertAlmostEqual(
+                float(numpy.mean(intervals)),
+                poisson_mean_interval_sec(config.emission_rate_per_min),
+                delta=1.0,
             )
-        )
-
-        # Create packets. At least two packets are expected to be generated from a 3500-byte msg
-        builder = PacketBuilder.real(random_bytes(3500), mixnet)
-        # Schedule two packets to the mix client without any interval
-        packet, route = builder.next()
-        await real_packet_queue.put((route[0].addr, packet))
-        packet, route = builder.next()
-        await real_packet_queue.put((route[0].addr, packet))
-
-        # Calculate intervals between packet emissions from the mix client
-        intervals = []
-        ts = datetime.now()
-        for _ in range(30):
-            _ = await outbound_socket.get()
-            now = datetime.now()
-            intervals.append((now - ts).total_seconds())
-            ts = now
-
-        # Check if packets were emitted at the Poisson emission_rate
-        # If emissions follow the Poisson distribution with a rate `lambda`,
-        # a mean interval between emissions must be `1/lambda`.
-        self.assertAlmostEqual(
-            float(numpy.mean(intervals)),
-            poisson_mean_interval_sec(emission_rate_per_min),
-            delta=1.0,
-        )
+        finally:
+            await mixclient.cancel()

mixnet/test_mixnet.py

@@ -1,40 +1,26 @@
-from typing import Tuple
 from unittest import IsolatedAsyncioTestCase
 
-from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
-
-from mixnet.bls import generate_bls
-from mixnet.mixnet import Mixnet, MixnetTopologySize, MixNode
-from mixnet.robustness import Robustness
-from mixnet.utils import random_bytes
+from mixnet.mixnet import Mixnet
+from mixnet.robustness import Robustness, RobustnessConfig
+from mixnet.test_utils import init_robustness_mixnet_config
 
 
 class TestMixnet(IsolatedAsyncioTestCase):
-    @staticmethod
-    def init() -> Tuple[Mixnet, Robustness]:
-        mixnet = Mixnet()
-        robustness = Robustness(
-            [
-                MixNode(
-                    generate_bls(),
-                    X25519PrivateKey.generate(),
-                    random_bytes(32),
-                )
-                for _ in range(12)
-            ],
-            MixnetTopologySize(3, 3),
-            mixnet,
+    async def test_topology_from_robustness(self):
+        robustness_mixnet_config = init_robustness_mixnet_config()
+
+        mixnet = await Mixnet.new(
+            robustness_mixnet_config.mixnode_candidates[0].encryption_private_key,
+            robustness_mixnet_config.mixnet_layer_config,
         )
-        robustness.set_entropy(b"entropy")
+        try:
+            robustness = Robustness(RobustnessConfig(robustness_mixnet_config), mixnet)
+            self.assertEqual(
+                robustness_mixnet_config.mixnet_layer_config, mixnet.get_config()
+            )
 
-        return (mixnet, robustness)
-
-    def test_topology_from_robustness(self):
-        mixnet, robustness = self.init()
-
-        topology1 = mixnet.get_topology()
-
-        robustness.set_entropy(b"new entropy")
-        topology2 = mixnet.get_topology()
-
-        self.assertNotEqual(topology1, topology2)
+            old_topology = robustness_mixnet_config.mixnet_layer_config.topology
+            robustness.set_entropy(b"new entropy")
+            self.assertNotEqual(old_topology, mixnet.get_config().topology)
+        finally:
+            await mixnet.cancel()

mixnet/test_node.py

@@ -1,19 +1,22 @@
 import asyncio
 from datetime import datetime
+from unittest import IsolatedAsyncioTestCase
 
 import numpy
 from pysphinx.sphinx import SphinxPacket
 
-from mixnet.node import NodeAddress, PacketPayloadQueue, PacketQueue
+from mixnet.node import MixNode, NodeAddress, PacketQueue
 from mixnet.packet import PacketBuilder
 from mixnet.poisson import poisson_interval_sec, poisson_mean_interval_sec
-from mixnet.test_mixnet import TestMixnet
-from mixnet.test_utils import with_test_timeout
+from mixnet.test_utils import (
+    init_robustness_mixnet_config,
+    with_test_timeout,
+)
 
 
-class TestMixNodeRunner(TestMixnet):
+class TestMixNodeRunner(IsolatedAsyncioTestCase):
     @with_test_timeout(180)
-    async def test_mixnode_runner_emission_rate(self):
+    async def test_mixnode_emission_rate(self):
         """
         Test if MixNodeRunner works as a M/M/inf queue.
 
@@ -21,71 +24,73 @@ class TestMixNodeRunner(TestMixnet):
         and if processing is delayed according to an exponential distribution with a rate `mu`,
         the rate of outputs should be `lambda`.
         """
-        mixnet, _ = self.init()
-        inbound_socket: PacketQueue = asyncio.Queue()
-        outbound_socket: PacketPayloadQueue = asyncio.Queue()
+        config = init_robustness_mixnet_config().mixnet_layer_config
+        config.emission_rate_per_min = 120  # lambda (= 2msg/sec)
+        config.delay_rate_per_min = 30  # mu (= 2s delay on average)
 
-        packet, route = PacketBuilder.real(b"msg", mixnet).next()
+        packet, route = PacketBuilder.build_real_packets(b"msg", config.topology)[0]
 
-        delay_rate_per_min = 30  # mu (= 2s delay on average)
         # Start only the first mix node for testing
-        _ = route[0].start(delay_rate_per_min, inbound_socket, outbound_socket)
-
-        # Send packets to the first mix node in a Poisson distribution
-        packet_count = 100
-        emission_rate_per_min = 120  # lambda (= 2msg/sec)
-        # This queue is just for counting how many packets have been sent so far.
-        sent_packet_queue: PacketQueue = asyncio.Queue()
-        _ = asyncio.create_task(
-            self.send_packets(
-                inbound_socket,
-                packet,
-                route[0].addr,
-                packet_count,
-                emission_rate_per_min,
-                sent_packet_queue,
+        mixnode = await MixNode.new(route[0].encryption_private_key, config)
+        try:
+            # Send packets to the first mix node in a Poisson distribution
+            packet_count = 100
+            # This queue is just for counting how many packets have been sent so far.
+            sent_packet_queue: PacketQueue = asyncio.Queue()
+            sender_task = asyncio.create_task(
+                self.send_packets(
+                    mixnode.inbound_socket,
+                    packet,
+                    route[0].addr,
+                    packet_count,
+                    config.emission_rate_per_min,
+                    sent_packet_queue,
+                )
             )
-        )
+            try:
+                # Calculate intervals between outputs and gather num_jobs in the first mix node.
+                intervals = []
+                num_jobs = []
+                ts = datetime.now()
+                for _ in range(packet_count):
+                    _ = await mixnode.outbound_socket.get()
+                    now = datetime.now()
+                    intervals.append((now - ts).total_seconds())
 
-        # Calculate intervals between outputs and gather num_jobs in the first mix node.
-        intervals = []
-        num_jobs = []
-        ts = datetime.now()
-        for _ in range(packet_count):
-            _ = await outbound_socket.get()
-            now = datetime.now()
-            intervals.append((now - ts).total_seconds())
+                    # Calculate the current # of jobs staying in the mix node
+                    num_packets_emitted_from_mixnode = len(intervals)
+                    num_packets_sent_to_mixnode = sent_packet_queue.qsize()
+                    num_jobs.append(
+                        num_packets_sent_to_mixnode - num_packets_emitted_from_mixnode
+                    )
 
-            # Calculate the current # of jobs staying in the mix node
-            num_packets_emitted_from_mixnode = len(intervals)
-            num_packets_sent_to_mixnode = sent_packet_queue.qsize()
-            num_jobs.append(
-                num_packets_sent_to_mixnode - num_packets_emitted_from_mixnode
-            )
+                    ts = now
 
-            ts = now
+                # Remove the first interval that would be much larger than other intervals,
+                # because of the delay in mix node.
+                intervals = intervals[1:]
+                num_jobs = num_jobs[1:]
 
-        # Remove the first interval that would be much larger than other intervals,
-        # because of the delay in mix node.
-        intervals = intervals[1:]
-        num_jobs = num_jobs[1:]
-
-        # Check if the emission rate of the first mix node is the same as
-        # the emission rate of the message sender, but with a delay.
-        # If outputs follow the Poisson distribution with a rate `lambda`,
-        # a mean interval between outputs must be `1/lambda`.
-        self.assertAlmostEqual(
-            float(numpy.mean(intervals)),
-            poisson_mean_interval_sec(emission_rate_per_min),
-            delta=1.0,
-        )
-        # If runner is a M/M/inf queue,
-        # a mean number of jobs being processed/scheduled in the runner must be `lambda/mu`.
-        self.assertAlmostEqual(
-            float(numpy.mean(num_jobs)),
-            round(emission_rate_per_min / delay_rate_per_min),
-            delta=1.5,
-        )
+                # Check if the emission rate of the first mix node is the same as
+                # the emission rate of the message sender, but with a delay.
+                # If outputs follow the Poisson distribution with a rate `lambda`,
+                # a mean interval between outputs must be `1/lambda`.
+                self.assertAlmostEqual(
+                    float(numpy.mean(intervals)),
+                    poisson_mean_interval_sec(config.emission_rate_per_min),
+                    delta=1.0,
+                )
+                # If runner is a M/M/inf queue,
+                # a mean number of jobs being processed/scheduled in the runner must be `lambda/mu`.
+                self.assertAlmostEqual(
+                    float(numpy.mean(num_jobs)),
+                    round(config.emission_rate_per_min / config.delay_rate_per_min),
+                    delta=1.5,
+                )
+            finally:
+                await sender_task
+        finally:
+            await mixnode.cancel()
 
     @staticmethod
     async def send_packets(

mixnet/test_packet.py

@@ -1,41 +1,45 @@
 from typing import List
+from unittest import TestCase
 
 from pysphinx.sphinx import ProcessedFinalHopPacket, SphinxPacket
 
-from mixnet.node import MixNode
+from mixnet.config import MixNodeInfo
 from mixnet.packet import (
     Fragment,
     MessageFlag,
     MessageReconstructor,
     PacketBuilder,
 )
-from mixnet.test_mixnet import TestMixnet
+from mixnet.test_utils import init_robustness_mixnet_config
 from mixnet.utils import random_bytes
 
 
-class TestPacket(TestMixnet):
+class TestPacket(TestCase):
     def test_real_packet(self):
-        mixnet, _ = self.init()
-
+        topology = init_robustness_mixnet_config().mixnet_layer_config.topology
         msg = random_bytes(3500)
-        builder = PacketBuilder.real(msg, mixnet)
-        packet0, route0 = builder.next()
-        packet1, route1 = builder.next()
-        packet2, route2 = builder.next()
-        packet3, route3 = builder.next()
-        self.assertRaises(StopIteration, builder.next)
+        packets_and_routes = PacketBuilder.build_real_packets(msg, topology)
+        self.assertEqual(4, len(packets_and_routes))
 
         reconstructor = MessageReconstructor()
         self.assertIsNone(
-            reconstructor.add(self.process_packet(packet1, route1)),
+            reconstructor.add(
+                self.process_packet(packets_and_routes[1][0], packets_and_routes[1][1])
+            ),
         )
         self.assertIsNone(
-            reconstructor.add(self.process_packet(packet3, route3)),
+            reconstructor.add(
+                self.process_packet(packets_and_routes[3][0], packets_and_routes[3][1])
+            ),
         )
         self.assertIsNone(
-            reconstructor.add(self.process_packet(packet2, route2)),
+            reconstructor.add(
+                self.process_packet(packets_and_routes[2][0], packets_and_routes[2][1])
+            ),
+        )
+        msg_with_flag = reconstructor.add(
+            self.process_packet(packets_and_routes[0][0], packets_and_routes[0][1])
         )
-        msg_with_flag = reconstructor.add(self.process_packet(packet0, route0))
         assert msg_with_flag is not None
         self.assertEqual(
             PacketBuilder.parse_msg_and_flag(msg_with_flag),
@@ -43,15 +47,15 @@ class TestPacket(TestMixnet):
         )
 
     def test_cover_packet(self):
-        mixnet, _ = self.init()
-
+        topology = init_robustness_mixnet_config().mixnet_layer_config.topology
         msg = b"cover"
-        builder = PacketBuilder.drop_cover(msg, mixnet)
-        packet, route = builder.next()
-        self.assertRaises(StopIteration, builder.next)
+        packets_and_routes = PacketBuilder.build_drop_cover_packets(msg, topology)
+        self.assertEqual(1, len(packets_and_routes))
 
         reconstructor = MessageReconstructor()
-        msg_with_flag = reconstructor.add(self.process_packet(packet, route))
+        msg_with_flag = reconstructor.add(
+            self.process_packet(packets_and_routes[0][0], packets_and_routes[0][1])
+        )
         assert msg_with_flag is not None
         self.assertEqual(
             PacketBuilder.parse_msg_and_flag(msg_with_flag),
@@ -59,7 +63,7 @@ class TestPacket(TestMixnet):
         )
 
     @staticmethod
-    def process_packet(packet: SphinxPacket, route: List[MixNode]) -> Fragment:
+    def process_packet(packet: SphinxPacket, route: List[MixNodeInfo]) -> Fragment:
         processed = packet.process(route[0].encryption_private_key)
         if isinstance(processed, ProcessedFinalHopPacket):
             return Fragment.from_bytes(processed.payload.recover_plain_playload())

mixnet/test_robustness.py

@@ -1,33 +1,14 @@
 from unittest import TestCase
 
-from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
-
-from mixnet.bls import generate_bls
-from mixnet.mixnet import Mixnet, MixnetTopologySize, MixNode
-from mixnet.robustness import Robustness
-from mixnet.utils import random_bytes
+from mixnet.test_utils import init_robustness_mixnet_config
 
 
 class TestRobustness(TestCase):
     def test_build_topology(self):
-        robustness = Robustness(
-            [
-                MixNode(
-                    generate_bls(),
-                    X25519PrivateKey.generate(),
-                    random_bytes(32),
-                )
-                for _ in range(12)
-            ],
-            MixnetTopologySize(3, 3),
-            Mixnet(),
-        )
+        robustness_mixnet_config = init_robustness_mixnet_config()
+        topology = robustness_mixnet_config.mixnet_layer_config.topology
+        topology_size = robustness_mixnet_config.topology_size
 
-        topology = robustness.build_topology(b"entropy")
-        self.assertEqual(
-            len(topology.layers), robustness.mixnet_topology_size.num_layers
-        )
+        self.assertEqual(len(topology.layers), topology_size.num_layers)
         for layer in topology.layers:
-            self.assertEqual(
-                len(layer), robustness.mixnet_topology_size.num_mixnodes_per_layer
-            )
+            self.assertEqual(len(layer), topology_size.num_mixnodes_per_layer)

mixnet/test_utils.py

@@ -1,5 +1,12 @@
 import asyncio
 
+from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
+
+from mixnet.bls import generate_bls
+from mixnet.config import MixnetConfig, MixNodeInfo
+from mixnet.robustness import MixnetTopologySize, Robustness, RobustnessMixnetConfig
+from mixnet.utils import random_bytes
+
 
 def with_test_timeout(t):
     def wrapper(coroutine):
@@ -10,3 +17,24 @@ def with_test_timeout(t):
         return run
 
     return wrapper
+
+
+def init_robustness_mixnet_config() -> RobustnessMixnetConfig:
+    mixnode_candidates = [
+        MixNodeInfo(
+            generate_bls(),
+            X25519PrivateKey.generate(),
+            random_bytes(32),
+        )
+        for _ in range(12)
+    ]
+    topology_size = MixnetTopologySize(3, 3)
+    mixnet_layer_config = MixnetConfig(
+        30,
+        3,
+        30,
+        Robustness.build_topology(mixnode_candidates, topology_size, b"entropy"),
+    )
+    return RobustnessMixnetConfig(
+        mixnode_candidates, topology_size, mixnet_layer_config
+    )