Mixnet PoC base branch

Cargo.toml

@@ -14,8 +14,13 @@ members = [
     "nomos-da/kzg",
     "nomos-da/full-replication",
     "nodes/nomos-node",
+    "nodes/mixnode",
     "simulations",
     "consensus-engine",
     "tests",
+    "mixnet/node",
+    "mixnet/client",
+    "mixnet/protocol",
+    "mixnet/topology",
 ]
 resolver = "2"

mixnet/client/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "mixnet-client"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+serde = { version = "1.0", features = ["derive"] }
+tracing = "0.1.37"
+tokio = { version = "1.29.1", features = ["net"] }
+sphinx-packet = "0.1.0"
+nym-sphinx = { package = "nym-sphinx", git = "https://github.com/nymtech/nym", tag = "v1.1.22" }
+# Using an older version, since `nym-sphinx` depends on `rand` v0.7.3.
+rand = "0.7.3"
+mixnet-protocol = { path = "../protocol" }
+mixnet-topology = { path = "../topology" }
+mixnet-util = { path = "../util" }
+futures = "0.3.28"
+thiserror = "1"

mixnet/client/src/config.rs

@@ -0,0 +1,31 @@
+use std::net::SocketAddr;
+
+use futures::{stream, StreamExt};
+use mixnet_topology::MixnetTopology;
+use serde::{Deserialize, Serialize};
+
+use crate::{receiver::Receiver, MessageStream, MixnetClientError};
+
+#[derive(Serialize, Deserialize, Clone, Debug)]
+pub struct MixnetClientConfig {
+    pub mode: MixnetClientMode,
+    pub topology: MixnetTopology,
+    pub connection_pool_size: usize,
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug)]
+pub enum MixnetClientMode {
+    Sender,
+    SenderReceiver(SocketAddr),
+}
+
+impl MixnetClientMode {
+    pub(crate) async fn run(&self) -> Result<MessageStream, MixnetClientError> {
+        match self {
+            Self::Sender => Ok(stream::empty().boxed()),
+            Self::SenderReceiver(node_address) => {
+                Ok(Receiver::new(*node_address).run().await?.boxed())
+            }
+        }
+    }
+}

mixnet/client/src/lib.rs

@@ -0,0 +1,56 @@
+pub mod config;
+mod receiver;
+mod sender;
+
+use std::error::Error;
+use std::time::Duration;
+
+pub use config::MixnetClientConfig;
+pub use config::MixnetClientMode;
+use futures::stream::BoxStream;
+use mixnet_util::ConnectionPool;
+use rand::Rng;
+use sender::Sender;
+use thiserror::Error;
+
+// A client for sending packets to Mixnet and receiving packets from Mixnet.
+pub struct MixnetClient<R: Rng> {
+    mode: MixnetClientMode,
+    sender: Sender<R>,
+}
+
+pub type MessageStream = BoxStream<'static, Result<Vec<u8>, MixnetClientError>>;
+
+impl<R: Rng> MixnetClient<R> {
+    pub fn new(config: MixnetClientConfig, rng: R) -> Self {
+        let cache = ConnectionPool::new(config.connection_pool_size);
+        Self {
+            mode: config.mode,
+            sender: Sender::new(config.topology, cache, rng),
+        }
+    }
+
+    pub async fn run(&self) -> Result<MessageStream, MixnetClientError> {
+        self.mode.run().await
+    }
+
+    pub fn send(
+        &mut self,
+        msg: Vec<u8>,
+        total_delay: Duration,
+    ) -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
+        self.sender.send(msg, total_delay)
+    }
+}
+
+#[derive(Error, Debug)]
+pub enum MixnetClientError {
+    #[error("mixnet node connect error")]
+    MixnetNodeConnectError,
+    #[error("mixnode stream has been closed")]
+    MixnetNodeStreamClosed,
+    #[error("unexpected stream body received")]
+    UnexpectedStreamBody,
+    #[error("invalid payload")]
+    InvalidPayload,
+}

mixnet/client/src/receiver.rs

@@ -0,0 +1,139 @@
+use std::{error::Error, net::SocketAddr};
+
+use futures::{stream, Stream, StreamExt};
+use mixnet_protocol::Body;
+use nym_sphinx::{
+    chunking::{fragment::Fragment, reconstruction::MessageReconstructor},
+    message::{NymMessage, PaddedMessage},
+    Payload,
+};
+use tokio::net::TcpStream;
+
+use crate::MixnetClientError;
+
+// Receiver accepts TCP connections to receive incoming payloads from the Mixnet.
+pub struct Receiver {
+    node_address: SocketAddr,
+}
+
+impl Receiver {
+    pub fn new(node_address: SocketAddr) -> Self {
+        Self { node_address }
+    }
+
+    pub async fn run(
+        &self,
+    ) -> Result<
+        impl Stream<Item = Result<Vec<u8>, MixnetClientError>> + Send + 'static,
+        MixnetClientError,
+    > {
+        let Ok(socket) = TcpStream::connect(self.node_address).await else {
+            return Err(MixnetClientError::MixnetNodeConnectError);
+        };
+
+        Ok(Self::message_stream(Box::pin(Self::fragment_stream(
+            socket,
+        ))))
+    }
+
+    fn fragment_stream(
+        socket: TcpStream,
+    ) -> impl Stream<Item = Result<Fragment, MixnetClientError>> + Send + 'static {
+        stream::unfold(socket, |mut socket| async move {
+            let Ok(body) = Body::read(&mut socket).await else {
+                // TODO: Maybe this is a hard error and the stream is corrupted? In that case stop the stream
+                return Some((Err(MixnetClientError::MixnetNodeStreamClosed), socket));
+            };
+
+            match body {
+                Body::SphinxPacket(_) => {
+                    Some((Err(MixnetClientError::UnexpectedStreamBody), socket))
+                }
+                Body::FinalPayload(payload) => Some((Self::fragment_from_payload(payload), socket)),
+            }
+        })
+    }
+
+    fn message_stream(
+        fragment_stream: impl Stream<Item = Result<Fragment, MixnetClientError>>
+            + Send
+            + Unpin
+            + 'static,
+    ) -> impl Stream<Item = Result<Vec<u8>, MixnetClientError>> + Send + 'static {
+        // MessageReconstructor buffers all received fragments
+        // and eventually returns reconstructed messages.
+        let message_reconstructor: MessageReconstructor = Default::default();
+
+        stream::unfold(
+            (fragment_stream, message_reconstructor),
+            |(mut fragment_stream, mut message_reconstructor)| async move {
+                let result =
+                    Self::reconstruct_message(&mut fragment_stream, &mut message_reconstructor)
+                        .await;
+                Some((result, (fragment_stream, message_reconstructor)))
+            },
+        )
+    }
+
+    fn fragment_from_payload(payload: Payload) -> Result<Fragment, MixnetClientError> {
+        let Ok(payload_plaintext) = payload.recover_plaintext() else {
+            return Err(MixnetClientError::InvalidPayload);
+        };
+        let Ok(fragment) = Fragment::try_from_bytes(&payload_plaintext) else {
+            return Err(MixnetClientError::InvalidPayload);
+        };
+        Ok(fragment)
+    }
+
+    async fn reconstruct_message(
+        fragment_stream: &mut (impl Stream<Item = Result<Fragment, MixnetClientError>>
+                  + Send
+                  + Unpin
+                  + 'static),
+        message_reconstructor: &mut MessageReconstructor,
+    ) -> Result<Vec<u8>, MixnetClientError> {
+        // Read fragments until at least one message is fully reconstructed.
+        while let Some(next) = fragment_stream.next().await {
+            match next {
+                Ok(fragment) => {
+                    if let Some(message) =
+                        Self::try_reconstruct_message(fragment, message_reconstructor)
+                    {
+                        return Ok(message);
+                    }
+                }
+                Err(e) => {
+                    return Err(e);
+                }
+            }
+        }
+
+        // fragment_stream closed before messages are fully reconstructed
+        Err(MixnetClientError::MixnetNodeStreamClosed)
+    }
+
+    fn try_reconstruct_message(
+        fragment: Fragment,
+        message_reconstructor: &mut MessageReconstructor,
+    ) -> Option<Vec<u8>> {
+        let reconstruction_result = message_reconstructor.insert_new_fragment(fragment);
+        match reconstruction_result {
+            Some((padded_message, _)) => {
+                let message = Self::remove_padding(padded_message).unwrap();
+                Some(message)
+            }
+            None => None,
+        }
+    }
+
+    fn remove_padding(msg: Vec<u8>) -> Result<Vec<u8>, Box<dyn Error>> {
+        let padded_message = PaddedMessage::new_reconstructed(msg);
+        // we need this because PaddedMessage.remove_padding requires it for other NymMessage types.
+        let dummy_num_mix_hops = 0;
+
+        match padded_message.remove_padding(dummy_num_mix_hops)? {
+            NymMessage::Plain(msg) => Ok(msg),
+            _ => todo!("return error"),
+        }
+    }
+}