WIP: Mixnet: Update spec for NomMix

mixnet/config.py

@@ -2,110 +2,77 @@
 
 import random
 from dataclasses import dataclass
-from typing import List, TypeAlias
+from typing import List
 
 from cryptography.hazmat.primitives.asymmetric.x25519 import (
     X25519PrivateKey,
     X25519PublicKey,
 )
-from pysphinx.node import Node
-
-from mixnet.bls import BlsPrivateKey, BlsPublicKey
-from mixnet.fisheryates import FisherYates
+from pysphinx.sphinx import Node as SphinxNode
 
 
 @dataclass
-class MixnetConfig:
-    topology_config: MixnetTopologyConfig
-    mixclient_config: MixClientConfig
-    mixnode_config: MixNodeConfig
-
+class GlobalConfig:
+    """
+    Global parameters used across all nodes in the network
+    """
 
-@dataclass
-class MixnetTopologyConfig:
-    mixnode_candidates: List[MixNodeInfo]
-    size: MixnetTopologySize
-    entropy: bytes
+    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 MixClientConfig:
-    emission_rate_per_min: int  # Poisson rate parameter: lambda
-    redundancy: int
-    topology: MixnetTopology
+class NodeConfig:
+    """
+    Node-specific parameters
+    """
+
+    private_key: X25519PrivateKey
+    mix_path_length: int
+    gossip: GossipConfig
 
 
 @dataclass
-class MixNodeConfig:
-    encryption_private_key: X25519PrivateKey
-    delay_rate_per_min: int  # Poisson rate parameter: mu
+class GossipConfig:
+    # The target number of peers a node should maintain in its p2p network
+    peering_degree: int
 
 
 @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 __init__(
-        self,
-        config: MixnetTopologyConfig,
-    ) -> None:
-        """
-        Build a new topology deterministically using an entropy and a given set of candidates.
-        """
-        shuffled = FisherYates.shuffle(config.mixnode_candidates, config.entropy)
-        sampled = shuffled[: config.size.num_total_mixnodes()]
+class MixMembership:
+    """
+    A list of public information of nodes in the network.
+    We assume that this list is eventually known to all nodes in the network (e.g. via p2p advertising).
+    """
 
-        layers = []
-        for layer_id in range(config.size.num_layers):
-            start = layer_id * config.size.num_mixnodes_per_layer
-            layer = sampled[start : start + config.size.num_mixnodes_per_layer]
-            layers.append(layer)
-        self.layers = layers
+    nodes: List[NodeInfo]
 
-    def generate_route(self, mix_destination: MixNodeInfo) -> list[MixNodeInfo]:
+    def generate_route(self, num_hops: int, last_mix: NodeInfo) -> list[NodeInfo]:
         """
         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
+        return [*(self.choose() for _ in range(num_hops - 1)), last_mix]
 
-    def choose_mix_destination(self) -> MixNodeInfo:
+    def choose(self) -> NodeInfo:
         """
-        Choose a mix node from the last mix layer as a mix destination
-        that will reconstruct a message from Sphinx packets.
+        Choose a mix node as a mix destination that will reconstruct a message from Sphinx packets.
         """
-        return random.choice(self.layers[-1])
+        return random.choice(self.nodes)
 
 
 @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
-
-
-# 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()
+class NodeInfo:
+    """
+    Public information of a node to be shared to all nodes in the network
+    """
 
-    def encryption_public_key(self) -> X25519PublicKey:
-        return self.encryption_private_key.public_key()
+    public_key: X25519PublicKey
 
-    def sphinx_node(self) -> Node:
-        return Node(self.encryption_private_key, self.addr)
+    def sphinx_node(self) -> SphinxNode:
+        # TODO: Use a pre-signed incentive tx, instead of NodeAddress
+        dummy_node_addr = bytes(32)
+        return SphinxNode(self.public_key, dummy_node_addr)

mixnet/connection.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import asyncio
+
+NetworkPacketQueue = asyncio.Queue[bytes]
+SimplexConnection = NetworkPacketQueue
+
+
+class DuplexConnection:
+    """
+    A duplex connection in which data can be transmitted and received simultaneously in both directions.
+    This is to mimic duplex communication in a real network (such as TCP or QUIC).
+    """
+
+    inbound: SimplexConnection
+    outbound: MixSimplexConnection
+
+    def __init__(self, inbound: SimplexConnection, outbound: MixSimplexConnection):
+        self.inbound = inbound
+        self.outbound = outbound
+
+    async def recv(self) -> bytes:
+        return await self.inbound.get()
+
+    async def send(self, packet: bytes):
+        await self.outbound.send(packet)
+
+
+class MixSimplexConnection:
+    """
+    Wraps a SimplexConnection to add a transmission rate and noise to the connection.
+    """
+
+    queue: NetworkPacketQueue
+    conn: SimplexConnection
+    transmission_rate_per_sec: int
+    noise_msg: bytes
+
+    def __init__(
+        self, conn: SimplexConnection, transmission_rate_per_sec: int, noise_msg: bytes
+    ):
+        self.queue = asyncio.Queue()
+        self.conn = conn
+        self.transmission_rate_per_sec = transmission_rate_per_sec
+        self.noise_msg = noise_msg
+        self.task = asyncio.create_task(self.__run())
+
+    async def __run(self):
+        while True:
+            await asyncio.sleep(1 / self.transmission_rate_per_sec)
+            # TODO: temporal mixing
+            if self.queue.empty():
+                # To guarantee GTR, send noise if there is no message to send
+                msg = self.noise_msg
+            else:
+                msg = self.queue.get_nowait()
+            await self.conn.put(msg)
+
+    async def send(self, msg: bytes):
+        await self.queue.put(msg)