AsyncWrite

This commit is contained in:
durch
2024-06-13 17:33:03 +02:00
parent 4adcd32ebf
commit cefb217c25
6 changed files with 199 additions and 29 deletions
@@ -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<MixPacket>,
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);
+6 -3
View File
@@ -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<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
deserializer.deserialize_bytes(MixPacketVisitor)
}
+27 -8
View File
@@ -63,11 +63,16 @@ impl From<PlainMessage> for ReconstructedMessage {
}
pub struct ReconstructedMessageCodec;
const OFFSET: usize = 4;
impl Encoder<ReconstructedMessage> 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<Option<Self::Item>, 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}")]
+148 -12
View File
@@ -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<Result<usize, std::io::Error>> {
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<std::prelude::v1::Result<(), std::io::Error>> {
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<std::prelude::v1::Result<(), std::io::Error>> {
AsyncWrite::poll_flush(self, cx)
}
}
impl Sink<InputMessage> for MixnetClient {
type Error = Error;
fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
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<InputMessage> 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<Result<()>> {
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<Result<()>> {
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<Result<usize, std::io::Error>> {
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<std::prelude::v1::Result<(), std::io::Error>> {
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<std::prelude::v1::Result<(), std::io::Error>> {
AsyncWrite::poll_flush(self, cx)
}
}
impl Sink<InputMessage> for MixnetClientSender {
type Error = Error;
fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
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<Result<()>> {
self.sender()
.poll_flush_unpin(cx)
.map_err(|_| Error::MessageSendingFailure)
}
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
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<InputMessage> {
&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<InputMessage> {
&mut self.client_input.input_sender
}
}
+2
View File
@@ -16,6 +16,8 @@ pub trait MixnetMessageSender {
None
}
fn sender(&mut self) -> &mut tokio_util::sync::PollSender<InputMessage>;
/// 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<()>;
@@ -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;