AsyncWrite
This commit is contained in:
@@ -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);
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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}")]
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user