From cefb217c253b6a34386d14f5e664e31f2cd36f9e Mon Sep 17 00:00:00 2001 From: durch Date: Thu, 13 Jun 2024 17:33:03 +0200 Subject: [PATCH] AsyncWrite --- .../src/client/inbound_messages.rs | 18 +- common/nymsphinx/forwarding/src/packet.rs | 9 +- common/nymsphinx/src/receiver.rs | 35 +++- sdk/rust/nym-sdk/src/mixnet/native_client.rs | 160 ++++++++++++++++-- sdk/rust/nym-sdk/src/mixnet/traits.rs | 2 + .../network-requester/src/core.rs | 4 +- 6 files changed, 199 insertions(+), 29 deletions(-) diff --git a/common/client-core/src/client/inbound_messages.rs b/common/client-core/src/client/inbound_messages.rs index c6fc5b03b5..af251c2373 100644 --- a/common/client-core/src/client/inbound_messages.rs +++ b/common/client-core/src/client/inbound_messages.rs @@ -7,7 +7,10 @@ use nym_sphinx::forwarding::packet::MixPacket; use nym_sphinx::params::PacketType; use nym_task::connections::TransmissionLane; use serde::{Deserialize, Serialize}; -use tokio_util::{bytes::BytesMut, codec::{Decoder, Encoder}}; +use tokio_util::{ + bytes::BytesMut, + codec::{Decoder, Encoder}, +}; use crate::error::ClientCoreError; @@ -68,6 +71,10 @@ pub enum InputMessage { } impl InputMessage { + pub fn simple(data: &[u8], recipient: Recipient) -> Self { + InputMessage::new_regular(recipient, data.to_vec(), TransmissionLane::General, None) + } + pub fn new_premade( msgs: Vec, lane: TransmissionLane, @@ -201,8 +208,11 @@ impl InputMessage { InputMessage::MessageWrapper { message, .. } => message.lane(), } } -} + pub fn serialized_size(&self) -> u64 { + bincode::serialized_size(self).expect("failed to get serialized InputMessage size") + 4 + } +} // TODO: Tests pub struct InputMessageCodec; @@ -235,9 +245,9 @@ impl Decoder for InputMessageCodec { return Ok(None); } - let decoded = match bincode::deserialize(&buf[4..len]) { + let decoded = match bincode::deserialize(&buf[4..len + 4]) { Ok(decoded) => decoded, - Err(_) => return Ok(None) + Err(_) => return Ok(None), }; buf.advance(len + 4); diff --git a/common/nymsphinx/forwarding/src/packet.rs b/common/nymsphinx/forwarding/src/packet.rs index 7846e1e067..534f6b64dd 100644 --- a/common/nymsphinx/forwarding/src/packet.rs +++ b/common/nymsphinx/forwarding/src/packet.rs @@ -4,7 +4,10 @@ use nym_sphinx_addressing::nodes::{NymNodeRoutingAddress, NymNodeRoutingAddressError}; use nym_sphinx_params::{PacketSize, PacketType}; use nym_sphinx_types::{NymPacket, NymPacketError}; -use serde::{de::{self, Visitor}, Deserialize, Deserializer, Serialize, Serializer}; +use serde::{ + de::{self, Visitor}, + Deserialize, Deserializer, Serialize, Serializer, +}; use std::fmt::{self, Debug, Formatter}; use thiserror::Error; @@ -128,7 +131,7 @@ impl Serialize for MixPacket { struct MixPacketVisitor; -impl <'de> Visitor<'de> for MixPacketVisitor { +impl<'de> Visitor<'de> for MixPacketVisitor { type Value = MixPacket; fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { @@ -140,7 +143,7 @@ impl <'de> Visitor<'de> for MixPacketVisitor { } } -impl <'de> Deserialize <'de> for MixPacket { +impl<'de> Deserialize<'de> for MixPacket { fn deserialize>(deserializer: D) -> Result { deserializer.deserialize_bytes(MixPacketVisitor) } diff --git a/common/nymsphinx/src/receiver.rs b/common/nymsphinx/src/receiver.rs index 40f4c9b7c4..691ace4da2 100644 --- a/common/nymsphinx/src/receiver.rs +++ b/common/nymsphinx/src/receiver.rs @@ -63,11 +63,16 @@ impl From for ReconstructedMessage { } pub struct ReconstructedMessageCodec; +const OFFSET: usize = 4; impl Encoder for ReconstructedMessageCodec { type Error = MessageRecoveryError; - fn encode(&mut self, item: ReconstructedMessage, buf: &mut BytesMut) -> Result<(), Self::Error> { + fn encode( + &mut self, + item: ReconstructedMessage, + buf: &mut BytesMut, + ) -> Result<(), Self::Error> { let encoded = bincode::serialize(&item).expect("failed to serialize ReconstructedMessage"); let encoded_len = encoded.len() as u32; let mut encoded_with_len = encoded_len.to_le_bytes().to_vec(); @@ -83,18 +88,33 @@ impl Decoder for ReconstructedMessageCodec { type Error = MessageRecoveryError; fn decode(&mut self, buf: &mut BytesMut) -> Result, Self::Error> { - if buf.len() < 4 { + println!("decoder called with buf: {:?}", buf.to_vec()); + if buf.len() < OFFSET { return Ok(None); } - let len = u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]) as usize; - if buf.len() < len + 4 { + let len = u32::from_le_bytes( + buf[0..OFFSET] + .try_into() + .expect("We know that we have at least OFFSET bytes in there"), + ) as usize; + + println!("len to decode {}", len); + println!("buf len {}", buf.len()); + + if buf.len() < len + OFFSET { return Ok(None); } - let decoded = match bincode::deserialize(&buf[4..len]) { - Ok(decoded) => decoded, - Err(_) => return Ok(None) + let decoded = match bincode::deserialize(&buf[OFFSET..len + OFFSET]) { + Ok(decoded) => { + println!("decoded: {:?}", decoded); + decoded + } + Err(e) => { + println!("error: {:?}", e); + return Ok(None); + } }; buf.advance(len + 4); @@ -103,7 +123,6 @@ impl Decoder for ReconstructedMessageCodec { } } - #[derive(Debug, Error)] pub enum MessageRecoveryError { #[error("The received message did not contain enough bytes to recover the ephemeral public key. Got {provided}. required: {required}")] diff --git a/sdk/rust/nym-sdk/src/mixnet/native_client.rs b/sdk/rust/nym-sdk/src/mixnet/native_client.rs index 0643dda0a3..26993ca64c 100644 --- a/sdk/rust/nym-sdk/src/mixnet/native_client.rs +++ b/sdk/rust/nym-sdk/src/mixnet/native_client.rs @@ -2,11 +2,11 @@ use crate::mixnet::client::MixnetClientBuilder; use crate::mixnet::traits::MixnetMessageSender; use crate::{Error, Result}; use async_trait::async_trait; -use bytecodec::io::WriteBuf; use bytes::{Buf as _, BytesMut}; -use futures::{ready, Sink, SinkExt, Stream, StreamExt}; +use futures::{ready, FutureExt, Sink, SinkExt, Stream, StreamExt}; use log::error; use nym_client_core::client::base_client::GatewayConnection; +use nym_client_core::client::inbound_messages::InputMessageCodec; use nym_client_core::client::{ base_client::{ClientInput, ClientOutput, ClientState}, inbound_messages::InputMessage, @@ -21,11 +21,11 @@ use nym_task::{ TaskHandle, }; use nym_topology::NymTopology; -use std::pin::Pin; +use std::pin::{pin, Pin}; use std::sync::Arc; use std::task::{Context, Poll}; -use tokio::io::{AsyncRead, ReadBuf}; -use tokio_util::codec::{Encoder, FramedWrite}; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; +use tokio_util::codec::{Encoder, FramedRead}; /// Client connected to the Nym mixnet. pub struct MixnetClient { @@ -288,11 +288,61 @@ impl AsyncRead for MixnetClient { } } +impl AsyncWrite for MixnetClient { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let codec = InputMessageCodec {}; + let mut reader = FramedRead::new(buf, codec); + let mut fut = reader.next(); + let msg = match fut.poll_unpin(cx) { + Poll::Ready(Some(Ok(msg))) => msg, + Poll::Ready(Some(Err(_))) => { + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::Other, + "failed to read message from input", + ))) + } + Poll::Pending => return Poll::Pending, + Poll::Ready(None) => return Poll::Ready(Ok(0)), + }; + + let msg_size = msg.serialized_size(); + + let mut fut = pin!(self.client_input.send(msg)); + match fut.poll_unpin(cx) { + Poll::Ready(Ok(())) => Poll::Ready(Ok(msg_size as usize)), + Poll::Ready(Err(_)) => Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::Other, + "failed to send message to mixnet", + ))), + Poll::Pending => Poll::Pending, + } + } + + fn poll_flush( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + Sink::poll_flush(self, cx) + .map_err(|_| std::io::Error::new(std::io::ErrorKind::Other, "failed to flush the sink")) + } + + fn poll_shutdown( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + AsyncWrite::poll_flush(self, cx) + } +} + impl Sink for MixnetClient { type Error = Error; fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - match self.client_input.input_sender.poll_ready_unpin(cx) { + match self.sender().poll_ready_unpin(cx) { Poll::Ready(Ok(())) => Poll::Ready(Ok(())), Poll::Ready(Err(_)) => Poll::Ready(Err(Error::MessageSendingFailure)), Poll::Pending => Poll::Pending, @@ -300,22 +350,100 @@ impl Sink for MixnetClient { } fn start_send(mut self: Pin<&mut Self>, item: InputMessage) -> Result<()> { - self.client_input - .input_sender + self.sender() .start_send_unpin(item) .map_err(|_| Error::MessageSendingFailure) } fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.client_input - .input_sender + self.sender() .poll_flush_unpin(cx) .map_err(|_| Error::MessageSendingFailure) } fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.client_input - .input_sender + self.sender() + .poll_close_unpin(cx) + .map_err(|_| Error::MessageSendingFailure) + } +} + +// TODO: there should be a better way of implementing Sink and AsyncWrite over T: MixnetMessageSender +impl AsyncWrite for MixnetClientSender { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let codec = InputMessageCodec {}; + let mut reader = FramedRead::new(buf, codec); + let mut fut = reader.next(); + let msg = match fut.poll_unpin(cx) { + Poll::Ready(Some(Ok(msg))) => msg, + Poll::Ready(Some(Err(_))) => { + return Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::Other, + "failed to read message from input", + ))) + } + Poll::Pending => return Poll::Pending, + Poll::Ready(None) => return Poll::Ready(Ok(0)), + }; + + let msg_size = msg.serialized_size(); + + let mut fut = pin!(self.client_input.send(msg)); + match fut.poll_unpin(cx) { + Poll::Ready(Ok(())) => Poll::Ready(Ok(msg_size as usize)), + Poll::Ready(Err(_)) => Poll::Ready(Err(std::io::Error::new( + std::io::ErrorKind::Other, + "failed to send message to mixnet", + ))), + Poll::Pending => Poll::Pending, + } + } + + fn poll_flush( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + Sink::poll_flush(self, cx) + .map_err(|_| std::io::Error::new(std::io::ErrorKind::Other, "failed to flush the sink")) + } + + fn poll_shutdown( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + AsyncWrite::poll_flush(self, cx) + } +} + +impl Sink for MixnetClientSender { + type Error = Error; + + fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.sender().poll_ready_unpin(cx) { + Poll::Ready(Ok(())) => Poll::Ready(Ok(())), + Poll::Ready(Err(_)) => Poll::Ready(Err(Error::MessageSendingFailure)), + Poll::Pending => Poll::Pending, + } + } + + fn start_send(mut self: Pin<&mut Self>, item: InputMessage) -> Result<()> { + self.sender() + .start_send_unpin(item) + .map_err(|_| Error::MessageSendingFailure) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.sender() + .poll_flush_unpin(cx) + .map_err(|_| Error::MessageSendingFailure) + } + + fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.sender() .poll_close_unpin(cx) .map_err(|_| Error::MessageSendingFailure) } @@ -363,6 +491,10 @@ impl MixnetMessageSender for MixnetClient { .await .map_err(|_| Error::MessageSendingFailure) } + + fn sender(&mut self) -> &mut tokio_util::sync::PollSender { + &mut self.client_input.input_sender + } } #[async_trait] @@ -377,4 +509,8 @@ impl MixnetMessageSender for MixnetClientSender { .await .map_err(|_| Error::MessageSendingFailure) } + + fn sender(&mut self) -> &mut tokio_util::sync::PollSender { + &mut self.client_input.input_sender + } } diff --git a/sdk/rust/nym-sdk/src/mixnet/traits.rs b/sdk/rust/nym-sdk/src/mixnet/traits.rs index a20a7cb996..d8649dfa7b 100644 --- a/sdk/rust/nym-sdk/src/mixnet/traits.rs +++ b/sdk/rust/nym-sdk/src/mixnet/traits.rs @@ -16,6 +16,8 @@ pub trait MixnetMessageSender { None } + fn sender(&mut self) -> &mut tokio_util::sync::PollSender; + /// Sends a [`InputMessage`] to the mixnet. This is the most low-level sending function, for /// full customization. async fn send(&mut self, message: InputMessage) -> Result<()>; diff --git a/service-providers/network-requester/src/core.rs b/service-providers/network-requester/src/core.rs index f4a349efeb..cf4b75794a 100644 --- a/service-providers/network-requester/src/core.rs +++ b/service-providers/network-requester/src/core.rs @@ -17,7 +17,7 @@ use nym_client_core::client::mix_traffic::transceiver::GatewayTransceiver; use nym_client_core::config::disk_persistence::CommonClientPaths; use nym_client_core::HardcodedTopologyProvider; use nym_network_defaults::NymNetworkDetails; -use nym_sdk::mixnet::{MixnetMessageSender, TopologyProvider}; +use nym_sdk::mixnet::TopologyProvider; use nym_service_providers_common::interface::{ BinaryInformation, ProviderInterfaceVersion, Request, RequestVersion, }; @@ -429,7 +429,7 @@ impl NRServiceProvider { } let response_message = msg.into_input_message(packet_type); - mixnet_client_sender.send(response_message).await.unwrap(); + nym_sdk::mixnet::MixnetMessageSender::send(&mut mixnet_client_sender, response_message).await.unwrap(); } else { log::error!("Exiting: channel closed!"); break;