Use tokio AsyncRead

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 }

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<u8>,
+    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<std::result::Result<usize, std::io::Error>> {
+    //     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<std::result::Result<usize, std::io::Error>> {
-        if self._read.buffer.len() < buf.len() {
-            let written = self._read.buffer.len();
-            buf[..written].copy_from_slice(&self._read.buffer);
+    ) -> Poll<tokio::io::Result<()>> {
+        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<std::result::Result<usize, std::io::Error>> {
+        buf: &mut ReadBuf,
+    ) -> Poll<tokio::io::Result<()>> {
+        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<InputMessage> 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<Result<()>> {
-        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<Result<()>> {
-        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)
     }
 }