2024-07-11 17:46:11 +09:00

107 lines
3.7 KiB
Python

from __future__ import annotations
import asyncio
from typing import TypeAlias
from pysphinx.sphinx import (
ProcessedFinalHopPacket,
ProcessedForwardHopPacket,
SphinxPacket,
)
from mixnet.config import GlobalConfig, NodeConfig
from mixnet.nomssip import Nomssip
from mixnet.sphinx import SphinxPacketBuilder
BroadcastChannel: TypeAlias = asyncio.Queue[bytes]
class Node:
"""
This represents any node in the network, which:
- generates/gossips mix messages (Sphinx packets)
- performs cryptographic mix (unwrapping Sphinx packets)
- generates noise
"""
def __init__(self, config: NodeConfig, global_config: GlobalConfig):
self.config = config
self.global_config = global_config
self.nomssip = Nomssip(
Nomssip.Config(
global_config.transmission_rate_per_sec,
config.nomssip.peering_degree,
self.__calculate_message_size(global_config),
),
self.__process_msg,
)
self.broadcast_channel = asyncio.Queue()
@staticmethod
def __calculate_message_size(global_config: GlobalConfig) -> int:
"""
Calculate the actual message size to be gossiped, which depends on the maximum length of mix path.
"""
sample_sphinx_packet, _ = SphinxPacketBuilder.build(
bytes(global_config.max_message_size),
global_config,
global_config.max_mix_path_length,
)
return len(sample_sphinx_packet.bytes())
async def __process_msg(self, msg: bytes) -> None:
"""
A handler to process messages received via gossip channel
"""
sphinx_packet = SphinxPacket.from_bytes(msg)
result = await self.__process_sphinx_packet(sphinx_packet)
match result:
case SphinxPacket():
# Gossip the next Sphinx packet
await self.nomssip.gossip(result.bytes())
case bytes():
# Broadcast the message fully recovered from Sphinx packets
await self.broadcast_channel.put(result)
case None:
return
async def __process_sphinx_packet(
self, packet: SphinxPacket
) -> SphinxPacket | bytes | None:
"""
Unwrap the Sphinx packet and process the next Sphinx packet or the payload if possible
"""
try:
processed = packet.process(self.config.private_key)
match processed:
case ProcessedForwardHopPacket():
return processed.next_packet
case ProcessedFinalHopPacket():
return processed.payload.recover_plain_playload()
except ValueError:
# Return nothing, if it cannot be unwrapped by the private key of this node.
return None
def connect(self, peer: Node):
"""
Establish a duplex connection with a peer node.
"""
inbound_conn, outbound_conn = asyncio.Queue(), asyncio.Queue()
# Register a duplex connection for its own use
self.nomssip.add_conn(inbound_conn, outbound_conn)
# Register a duplex connection for the peer
peer.nomssip.add_conn(outbound_conn, inbound_conn)
async def send_message(self, msg: bytes):
"""
Build a Sphinx packet and gossip it to all connected peers.
"""
# Here, we handle the case in which a msg is split into multiple Sphinx packets.
# 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,
self.config.mix_path_length,
)
await self.nomssip.gossip(sphinx_packet.bytes())