simplified keeping track of message sequence

This commit is contained in:
Jędrzej Stuczyński
2023-06-08 16:04:56 +01:00
parent 7228331db6
commit f4cd372808
7 changed files with 110 additions and 118 deletions
Generated
-1
View File
@@ -3949,7 +3949,6 @@ dependencies = [
"nym-credential-storage",
"nym-crypto",
"nym-network-defaults",
"nym-ordered-buffer",
"nym-sdk",
"nym-service-providers-common",
"nym-socks5-proxy-helpers",
-15
View File
@@ -1,18 +1,3 @@
mod buffer;
#[deprecated]
mod sender;
pub use buffer::{OrderedMessageBuffer, ReadContiguousData};
pub use sender::OrderedMessageSender;
use std::sync::atomic::AtomicU64;
use std::sync::Arc;
#[derive(Debug, Clone)]
pub struct Socks5DataStream {
inner: Arc<Socks5DataStreamInner>,
}
#[derive(Debug)]
pub struct Socks5DataStreamInner {
next_seq: AtomicU64,
}
+2
View File
@@ -4,3 +4,5 @@
pub mod available_reader;
pub mod connection_controller;
pub mod proxy_runner;
pub mod ordered_sender;
@@ -0,0 +1,99 @@
// Copyright 2023 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
use crate::proxy_runner::MixProxySender;
use bytes::Bytes;
use log::{debug, error};
use nym_socks5_requests::{ConnectionId, SocketData};
use std::io;
pub struct OrderedMessageSender<S, F> {
connection_id: ConnectionId,
mixnet_sender: MixProxySender<S>,
next_message_seq: u64,
mix_message_adapter: F,
}
impl<S, F> OrderedMessageSender<S, F>
where
F: Fn(SocketData) -> S,
{
pub fn new(
connection_id: ConnectionId,
mixnet_sender: MixProxySender<S>,
mix_message_adapter: F,
) -> Self {
OrderedMessageSender {
connection_id,
mixnet_sender,
next_message_seq: 0,
mix_message_adapter,
}
}
fn sequence(&mut self) -> u64 {
let next = self.next_message_seq;
self.next_message_seq += 1;
next
}
fn construct_message(&mut self, data: Vec<u8>, local_socket_closed: bool) -> S {
let data = SocketData::new(
self.sequence(),
self.connection_id,
local_socket_closed,
data,
);
(self.mix_message_adapter)(data)
}
async fn send_message(&self, message: S) {
if self.mixnet_sender.send(message).await.is_err() {
panic!("BatchRealMessageReceiver has stopped receiving!")
}
}
pub async fn send_empty_close(&mut self) {
let message = self.construct_message(Vec::new(), true);
self.send_message(message).await
}
pub async fn send_empty_keepalive(&mut self) {
log::trace!("Sending keepalive for connection: {}", self.connection_id);
let message = self.construct_message(Vec::new(), false);
self.send_message(message).await
}
pub async fn deal_with_data(
&mut self,
read_data: Option<io::Result<Bytes>>,
local_destination_address: &str,
remote_source_address: &str,
) -> bool {
let connection_id = self.connection_id;
let (read_data, is_finished) = match read_data {
Some(data) => match data {
Ok(data) => (data, false),
Err(err) => {
error!(target: &*format!("({connection_id}) socks5 inbound"), "failed to read request from the socket - {err}");
(Default::default(), true)
}
},
None => (Default::default(), true),
};
debug!(
target: &*format!("({connection_id}) socks5 inbound"),
"[{} bytes]\t{} → local → mixnet → remote → {}. Local closed: {}",
read_data.len(),
local_destination_address,
remote_source_address,
is_finished
);
let message = self.construct_message(read_data.into(), is_finished);
self.send_message(message).await;
is_finished
}
}
@@ -4,104 +4,21 @@
use super::MixProxySender;
use super::SHUTDOWN_TIMEOUT;
use crate::available_reader::AvailableReader;
use crate::ordered_sender::OrderedMessageSender;
use crate::proxy_runner::KEEPALIVE_INTERVAL;
use bytes::Bytes;
use futures::FutureExt;
use futures::StreamExt;
use log::*;
use nym_ordered_buffer::OrderedMessageSender;
use nym_socks5_requests::{ConnectionId, SocketData};
use nym_task::connections::LaneQueueLengths;
use nym_task::connections::TransmissionLane;
use nym_task::TaskClient;
use std::fmt::Debug;
use std::sync::Arc;
use std::time::Duration;
use std::{io, sync::Arc};
use tokio::select;
use tokio::{net::tcp::OwnedReadHalf, sync::Notify, time::sleep};
async fn send_empty_close<F, S>(
connection_id: ConnectionId,
message_sender: &mut OrderedMessageSender,
mix_sender: &MixProxySender<S>,
adapter_fn: F,
) where
F: Fn(SocketData) -> S,
S: Debug,
{
let message = SocketData::new(message_sender.sequence(), connection_id, true, Vec::new());
mix_sender
.send(adapter_fn(message))
.await
.expect("BatchRealMessageReceiver has stopped receiving!");
}
async fn send_empty_keepalive<F, S>(
connection_id: ConnectionId,
message_sender: &mut OrderedMessageSender,
mix_sender: &MixProxySender<S>,
adapter_fn: F,
) where
F: Fn(SocketData) -> S,
S: Debug,
{
log::trace!("Sending keepalive for connection: {connection_id}");
let message = SocketData::new(message_sender.sequence(), connection_id, false, Vec::new());
mix_sender
.send(adapter_fn(message))
.await
.expect("BatchRealMessageReceiver has stopped receiving!");
}
async fn deal_with_data<F, S>(
read_data: Option<io::Result<Bytes>>,
local_destination_address: &str,
remote_source_address: &str,
connection_id: ConnectionId,
message_sender: &mut OrderedMessageSender,
mix_sender: &MixProxySender<S>,
adapter_fn: F,
) -> bool
where
F: Fn(SocketData) -> S,
S: Debug,
{
let (read_data, is_finished) = match read_data {
Some(data) => match data {
Ok(data) => (data, false),
Err(err) => {
error!(target: &*format!("({connection_id}) socks5 inbound"), "failed to read request from the socket - {err}");
(Default::default(), true)
}
},
None => (Default::default(), true),
};
debug!(
target: &*format!("({connection_id}) socks5 inbound"),
"[{} bytes]\t{} → local → mixnet → remote → {}. Local closed: {}",
read_data.len(),
local_destination_address,
remote_source_address,
is_finished
);
let message = SocketData::new(
message_sender.sequence(),
connection_id,
is_finished,
read_data.into(),
);
mix_sender
.send(adapter_fn(message))
.await
.expect("InputMessageReceiver has stopped receiving!");
is_finished
}
async fn wait_until_lane_empty(lane_queue_lengths: &Option<LaneQueueLengths>, connection_id: u64) {
if let Some(lane_queue_lengths) = lane_queue_lengths {
if tokio::time::timeout(
@@ -179,7 +96,7 @@ where
// TODO: this multiplication by 4 is completely arbitrary here
let mut available_reader =
AvailableReader::new(&mut reader, Some(available_plaintext_per_mix_packet * 4));
let mut message_sender = OrderedMessageSender::new();
let mut message_sender = OrderedMessageSender::new(connection_id, mix_sender, adapter_fn);
// Shutdown if outbound signled to shutdown
let shutdown_future = shutdown_notify.notified().then(|_| sleep(SHUTDOWN_TIMEOUT));
@@ -218,7 +135,7 @@ where
);
// inform remote just in case it was closed because of lack of heartbeat.
// worst case the remote will just have couple of false negatives
send_empty_close(connection_id, &mut message_sender, &mix_sender, &adapter_fn).await;
message_sender.send_empty_close().await;
break;
}
_ = shutdown_listener.recv() => {
@@ -234,7 +151,7 @@ where
break;
}
_ = keepalive_timer.tick() => {
send_empty_keepalive(connection_id, &mut message_sender, &mix_sender, &adapter_fn).await;
message_sender.send_empty_keepalive().await;
}
// Read the next data when there is space in the lane.
// The purpose of chaining the wait here is that it makes sure we can cancel the
@@ -242,15 +159,8 @@ where
read_data = wait_until_lane_almost_empty(&lane_queue_lengths, connection_id)
.then(|_| available_reader.next()), if !we_are_closed =>
{
if deal_with_data(
read_data,
&local_destination_address,
&remote_source_address,
connection_id,
&mut message_sender,
&mix_sender,
&adapter_fn,
).await {
if message_sender.deal_with_data(read_data, &local_destination_address,
&remote_source_address).await {
// After reading the last data, notify the closing_future to wait until the
// lane is clear before exiting.
// We don't wait here since we want to be able to cancel the wait on close or
@@ -43,7 +43,6 @@ nym-bin-common = { path = "../../common/bin-common", features = ["output_format"
nym-network-defaults = { path = "../../common/network-defaults" }
nym-sdk = { path = "../../sdk/rust/nym-sdk" }
nym-sphinx = { path = "../../common/nymsphinx" }
nym-ordered-buffer = {path = "../../common/socks5/ordered-buffer"}
nym-socks5-proxy-helpers = { path = "../../common/socks5/proxy-helpers" }
nym-service-providers-common = { path = "../common" }
nym-socks5-requests = { path = "../../common/socks5/requests" }
@@ -6,7 +6,6 @@ use crate::core::new_legacy_request_version;
use crate::reply::MixnetMessage;
use async_trait::async_trait;
use log::*;
use nym_ordered_buffer::OrderedMessageSender;
use nym_service_providers_common::interface::RequestVersion;
use nym_socks5_proxy_helpers::proxy_runner::MixProxySender;
use nym_socks5_requests::{
@@ -192,9 +191,8 @@ impl StatisticsCollector for ServiceStatisticsCollector {
.expect("MixProxyReader has stopped receiving!");
trace!("Sending data to statistics service");
let mut message_sender = OrderedMessageSender::new();
// let ordered_msg = message_sender.wrap_message(msg).into_bytes();
let message = SocketData::new(message_sender.sequence(), conn_id, true, msg);
let message = SocketData::new(0, conn_id, true, msg);
let send_req = Socks5RequestContent::new_send(message);
let mixnet_message = MixnetMessage::new_network_data_request(