diff --git a/sdk/rust/nym-sdk/Cargo.toml b/sdk/rust/nym-sdk/Cargo.toml index 98318604ae..4253c92e2f 100644 --- a/sdk/rust/nym-sdk/Cargo.toml +++ b/sdk/rust/nym-sdk/Cargo.toml @@ -42,6 +42,8 @@ tap = { workspace = true } thiserror = { workspace = true } url = { workspace = true } toml = "0.5.10" +tokio = { workspace = true } +tokio-util = { workspace = true, features = ["codec"] } [dev-dependencies] anyhow = { workspace = true } diff --git a/sdk/rust/nym-sdk/src/mixnet/native_client.rs b/sdk/rust/nym-sdk/src/mixnet/native_client.rs index 5620d58980..0643dda0a3 100644 --- a/sdk/rust/nym-sdk/src/mixnet/native_client.rs +++ b/sdk/rust/nym-sdk/src/mixnet/native_client.rs @@ -2,7 +2,9 @@ use crate::mixnet::client::MixnetClientBuilder; use crate::mixnet::traits::MixnetMessageSender; use crate::{Error, Result}; use async_trait::async_trait; -use futures::{ready, AsyncRead, Sink, SinkExt, Stream, StreamExt}; +use bytecodec::io::WriteBuf; +use bytes::{Buf as _, BytesMut}; +use futures::{ready, Sink, SinkExt, Stream, StreamExt}; use log::error; use nym_client_core::client::base_client::GatewayConnection; use nym_client_core::client::{ @@ -12,6 +14,7 @@ use nym_client_core::client::{ }; use nym_crypto::asymmetric::identity; use nym_sphinx::addressing::clients::Recipient; +use nym_sphinx::receiver::ReconstructedMessageCodec; use nym_sphinx::{params::PacketType, receiver::ReconstructedMessage}; use nym_task::{ connections::{ConnectionCommandSender, LaneQueueLengths}, @@ -21,6 +24,8 @@ use nym_topology::NymTopology; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use tokio::io::{AsyncRead, ReadBuf}; +use tokio_util::codec::{Encoder, FramedWrite}; /// Client connected to the Nym mixnet. pub struct MixnetClient { @@ -58,7 +63,7 @@ pub struct MixnetClient { #[derive(Debug, Default)] struct ReadBuffer { - buffer: Vec, + buffer: BytesMut, } impl ReadBuffer { @@ -210,22 +215,41 @@ impl MixnetClient { // it must be dropped to finalize the shutdown } + // fn read_buffer_to_slice( + // &mut self, + // buf: &mut [u8], + // cx: &mut Context<'_>, + // ) -> Poll> { + // if self._read.buffer.len() < buf.len() { + // let written = self._read.buffer.len(); + // buf[..written].copy_from_slice(&self._read.buffer); + // self._read.clear(); + // Poll::Ready(Ok(written)) + // } else { + // let written = buf.len(); + // buf.copy_from_slice(&self._read.buffer[..written]); + // self._read.buffer = self._read.buffer[written..].to_vec(); + // cx.waker().wake_by_ref(); + // Poll::Ready(Ok(written)) + // } + // } + fn read_buffer_to_slice( &mut self, - buf: &mut [u8], + buf: &mut ReadBuf, cx: &mut Context<'_>, - ) -> Poll> { - if self._read.buffer.len() < buf.len() { - let written = self._read.buffer.len(); - buf[..written].copy_from_slice(&self._read.buffer); + ) -> Poll> { + if self._read.buffer.len() < buf.capacity() { + // let written = self._read.buffer.len(); + buf.put_slice(&self._read.buffer); self._read.clear(); - Poll::Ready(Ok(written)) + Poll::Ready(Ok(())) } else { - let written = buf.len(); - buf.copy_from_slice(&self._read.buffer[..written]); - self._read.buffer = self._read.buffer[written..].to_vec(); + let written = buf.capacity(); + buf.put_slice(&self._read.buffer[..written]); + self._read.buffer.advance(written); cx.waker().wake_by_ref(); - Poll::Ready(Ok(written)) + Poll::Ready(Ok(())) } } } @@ -240,19 +264,25 @@ impl AsyncRead for MixnetClient { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, - buf: &mut [u8], - ) -> Poll> { + buf: &mut ReadBuf, + ) -> Poll> { + let mut codec = ReconstructedMessageCodec {}; + if self._read.pending() { return self.read_buffer_to_slice(buf, cx); } let msg = match self.as_mut().poll_next(cx) { Poll::Ready(Some(msg)) => msg, - Poll::Ready(None) => return Poll::Ready(Ok(0)), + Poll::Ready(None) => return Poll::Ready(Ok(())), Poll::Pending => return Poll::Pending, }; - self._read.buffer = msg.message; + // let mut buffer = BytesMut::new(); + + codec.encode(msg, &mut self._read.buffer).unwrap(); + + // = buffer.to_vec(); self.read_buffer_to_slice(buf, cx) } @@ -270,15 +300,24 @@ impl Sink for MixnetClient { } fn start_send(mut self: Pin<&mut Self>, item: InputMessage) -> Result<()> { - self.client_input.input_sender.start_send_unpin(item).map_err(|_| Error::MessageSendingFailure) + self.client_input + .input_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.poll_flush_unpin(cx).map_err(|_| Error::MessageSendingFailure) + self.client_input + .input_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.poll_close_unpin(cx).map_err(|_| Error::MessageSendingFailure) + self.client_input + .input_sender + .poll_close_unpin(cx) + .map_err(|_| Error::MessageSendingFailure) } }