nomos-specs/mixnet/node.py

108 lines
3.4 KiB
Python
Raw Normal View History

2024-01-23 01:29:14 +00:00
from __future__ import annotations
2024-01-25 09:04:55 +00:00
import asyncio
from contextlib import suppress
from typing import Self, Tuple, TypeAlias
2024-01-23 01:29:14 +00:00
from cryptography.hazmat.primitives.asymmetric.x25519 import (
X25519PrivateKey,
)
from pysphinx.sphinx import (
Payload,
ProcessedFinalHopPacket,
ProcessedForwardHopPacket,
SphinxPacket,
UnknownHeaderTypeError,
)
from mixnet.config import MixNodeConfig, NodeAddress
2024-01-23 01:29:14 +00:00
from mixnet.poisson import poisson_interval_sec
2024-01-25 09:04:55 +00:00
PacketQueue: TypeAlias = "asyncio.Queue[Tuple[NodeAddress, SphinxPacket]]"
2024-01-23 01:29:14 +00:00
PacketPayloadQueue: TypeAlias = (
2024-01-25 09:04:55 +00:00
"asyncio.Queue[Tuple[NodeAddress, SphinxPacket | Payload]]"
2024-01-23 01:29:14 +00:00
)
class MixNode:
"""
2024-01-25 09:04:55 +00:00
A class handling incoming packets with delays
2024-01-23 01:29:14 +00:00
2024-01-25 09:04:55 +00:00
This class is defined separated with the MixNode class,
in order to define the MixNode as a simple dataclass for clarity.
2024-01-23 01:29:14 +00:00
"""
config: MixNodeConfig
inbound_socket: PacketQueue
outbound_socket: PacketPayloadQueue
task: asyncio.Task # A reference just to prevent task from being garbage collected
@classmethod
async def new(
cls,
config: MixNodeConfig,
) -> Self:
self = cls()
self.config = config
self.inbound_socket = asyncio.Queue()
self.outbound_socket = asyncio.Queue()
self.task = asyncio.create_task(self.__run())
return self
async def __run(self):
2024-01-23 01:29:14 +00:00
"""
2024-01-25 09:04:55 +00:00
Read SphinxPackets from inbound socket and spawn a thread for each packet to process it.
2024-01-23 01:29:14 +00:00
2024-01-25 09:04:55 +00:00
This thread approximates a M/M/inf queue.
2024-01-23 01:29:14 +00:00
"""
2024-01-25 09:04:55 +00:00
# A set just for gathering a reference of tasks to prevent them from being garbage collected.
# https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
self.tasks = set()
2024-01-23 01:29:14 +00:00
2024-01-25 09:04:55 +00:00
while True:
_, packet = await self.inbound_socket.get()
task = asyncio.create_task(
self.__process_packet(
packet,
self.config.encryption_private_key,
self.config.delay_rate_per_min,
2024-01-25 09:04:55 +00:00
)
)
self.tasks.add(task)
# To discard the task from the set automatically when it is done.
task.add_done_callback(self.tasks.discard)
2024-01-23 01:29:14 +00:00
async def __process_packet(
2024-01-23 01:29:14 +00:00
self,
packet: SphinxPacket,
encryption_private_key: X25519PrivateKey,
delay_rate_per_min: int, # Poisson rate parameter: mu
2024-01-23 01:29:14 +00:00
):
2024-01-25 09:04:55 +00:00
"""
Process a single packet with a delay that follows exponential distribution,
and forward it to the next mix node or the mix destination
2024-01-23 01:29:14 +00:00
2024-01-25 09:04:55 +00:00
This thread is a single server (worker) in a M/M/inf queue that MixNodeRunner approximates.
"""
delay_sec = poisson_interval_sec(delay_rate_per_min)
2024-01-25 09:04:55 +00:00
await asyncio.sleep(delay_sec)
2024-01-23 01:29:14 +00:00
processed = packet.process(encryption_private_key)
2024-01-23 01:29:14 +00:00
match processed:
case ProcessedForwardHopPacket():
2024-01-25 09:04:55 +00:00
await self.outbound_socket.put(
2024-01-23 01:29:14 +00:00
(processed.next_node_address, processed.next_packet)
)
case ProcessedFinalHopPacket():
2024-01-25 09:04:55 +00:00
await self.outbound_socket.put(
2024-01-23 01:29:14 +00:00
(processed.destination_node_address, processed.payload)
)
case _:
raise UnknownHeaderTypeError
async def cancel(self) -> None:
self.task.cancel()
with suppress(asyncio.CancelledError):
await self.task