From fa95d15eac1be606b243f7eab0d5e2080b82e2f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= Date: Mon, 21 Nov 2022 23:52:30 +0100 Subject: [PATCH] socks5-client: throttle connection inbound from application until data is sent (#1783) * socks5: throttle send * client-connections: add additional methods * WIP * Update * Input message sender bounded * WIP * Remove the delay that is no longer needed * rustfmt * clippy * Fix wasm build * clippy * Try to use MixProxySender/Reader type alias * Extract out wait function * Wait on every msg * changelog: add note * rustfmt --- CHANGELOG.md | 2 + .../src/client/inbound_messages.rs | 5 +- .../input_message_listener.rs | 16 ++-- .../retransmission_request_listener.rs | 11 ++- .../src/client/real_messages_control/mod.rs | 2 +- .../real_traffic_stream.rs | 9 +- clients/native/src/client/mod.rs | 33 +++++-- clients/native/src/websocket/handler.rs | 41 +++++--- clients/socks5/src/client/mod.rs | 7 +- clients/socks5/src/socks/client.rs | 16 +++- clients/socks5/src/socks/server.rs | 6 +- clients/webassembly/Cargo.toml | 9 +- clients/webassembly/src/client/mod.rs | 13 ++- common/client-connections/src/lib.rs | 32 ++++++- .../proxy-helpers/src/proxy_runner/inbound.rs | 96 ++++++++++++++++--- .../proxy-helpers/src/proxy_runner/mod.rs | 12 ++- .../network-requester/src/connection.rs | 6 +- .../network-requester/src/core.rs | 36 ++++--- .../src/statistics/collector.rs | 28 ++++-- 19 files changed, 282 insertions(+), 98 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d21a49d7fd..7cd8056355 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,8 +12,10 @@ Post 1.0.0 release, the changelog format is based on [Keep a Changelog](https:// ### Changed - clients: add concept of transmission lanes to better handle multiple data streams ([#1720]) +- socks5-client: wait closing inbound connection until data is sent, and throttle incoming data in general ([#1783]) [#1720]: https://github.com/nymtech/nym/pull/1720 +[#1783]: https://github.com/nymtech/nym/pull/1783 ## [v1.1.0](https://github.com/nymtech/nym/tree/v1.1.0) (2022-11-09) diff --git a/clients/client-core/src/client/inbound_messages.rs b/clients/client-core/src/client/inbound_messages.rs index 3cc71d4460..909e3720f9 100644 --- a/clients/client-core/src/client/inbound_messages.rs +++ b/clients/client-core/src/client/inbound_messages.rs @@ -1,10 +1,9 @@ use client_connections::TransmissionLane; -use futures::channel::mpsc; use nymsphinx::addressing::clients::Recipient; use nymsphinx::anonymous_replies::ReplySurb; -pub type InputMessageSender = mpsc::UnboundedSender; -pub type InputMessageReceiver = mpsc::UnboundedReceiver; +pub type InputMessageSender = tokio::sync::mpsc::Sender; +pub type InputMessageReceiver = tokio::sync::mpsc::Receiver; #[derive(Debug)] pub enum InputMessage { diff --git a/clients/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs b/clients/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs index 20cbc10de5..58fa65af56 100644 --- a/clients/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs +++ b/clients/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs @@ -9,7 +9,6 @@ use crate::client::{ topology_control::TopologyAccessor, }; use client_connections::TransmissionLane; -use futures::StreamExt; use log::*; use nymsphinx::anonymous_replies::ReplySurb; use nymsphinx::preparer::MessagePreparer; @@ -188,9 +187,14 @@ where // there's no point in trying to send nothing if let Some(real_messages) = real_messages { // tells real message sender (with the poisson timer) to send this to the mix network - self.real_message_sender - .unbounded_send((real_messages, lane)) - .unwrap(); + if self + .real_message_sender + .send((real_messages, lane)) + .await + .is_err() + { + panic!(); + } } } @@ -200,7 +204,7 @@ where while !shutdown.is_shutdown() { tokio::select! { - input_msg = self.input_receiver.next() => match input_msg { + input_msg = self.input_receiver.recv() => match input_msg { Some(input_msg) => { self.on_input_message(input_msg).await; }, @@ -221,7 +225,7 @@ where #[cfg(target_arch = "wasm32")] pub(super) async fn run(&mut self) { debug!("Started InputMessageListener without graceful shutdown support"); - while let Some(input_msg) = self.input_receiver.next().await { + while let Some(input_msg) = self.input_receiver.recv().await { self.on_input_message(input_msg).await; } } diff --git a/clients/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs b/clients/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs index 8a26809c79..2fbab44b61 100644 --- a/clients/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs +++ b/clients/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs @@ -116,12 +116,17 @@ where .unwrap(); // send to `OutQueueControl` to eventually send to the mix network - self.real_message_sender - .unbounded_send(( + if self + .real_message_sender + .send(( vec![RealMessage::new(prepared_fragment.mix_packet, frag_id)], TransmissionLane::Retransmission, )) - .unwrap(); + .await + .is_err() + { + panic!(); + } } #[cfg(not(target_arch = "wasm32"))] diff --git a/clients/client-core/src/client/real_messages_control/mod.rs b/clients/client-core/src/client/real_messages_control/mod.rs index a1e9290c64..d8b9dd34a6 100644 --- a/clients/client-core/src/client/real_messages_control/mod.rs +++ b/clients/client-core/src/client/real_messages_control/mod.rs @@ -119,7 +119,7 @@ impl RealMessagesController { ) -> Self { let rng = OsRng; - let (real_message_sender, real_message_receiver) = mpsc::unbounded(); + let (real_message_sender, real_message_receiver) = tokio::sync::mpsc::channel(1); let (sent_notifier_tx, sent_notifier_rx) = mpsc::unbounded(); let ack_controller_connectors = AcknowledgementControllerConnectors::new( diff --git a/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs b/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs index 43be7524b3..9f42232db5 100644 --- a/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs +++ b/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs @@ -7,7 +7,6 @@ use crate::client::topology_control::TopologyAccessor; use client_connections::{ ClosedConnectionReceiver, ConnectionId, LaneQueueLengths, TransmissionLane, }; -use futures::channel::mpsc; use futures::task::{Context, Poll}; use futures::{Future, Stream, StreamExt}; use log::*; @@ -158,8 +157,8 @@ impl RealMessage { // messages are already prepared, etc. the real point of it is to forward it to mix_traffic // after sufficient delay pub(crate) type BatchRealMessageSender = - mpsc::UnboundedSender<(Vec, TransmissionLane)>; -type BatchRealMessageReceiver = mpsc::UnboundedReceiver<(Vec, TransmissionLane)>; + tokio::sync::mpsc::Sender<(Vec, TransmissionLane)>; +type BatchRealMessageReceiver = tokio::sync::mpsc::Receiver<(Vec, TransmissionLane)>; pub(crate) enum StreamMessage { Cover, @@ -368,7 +367,7 @@ where // in `Vec`, this ensures that on average we will fetch messages faster than we can // send, which is a condition for being able to multiplex sphinx packets from multiple // data streams. - match Pin::new(&mut self.real_receiver).poll_next(cx) { + match Pin::new(&mut self.real_receiver).poll_recv(cx) { // in the case our real message channel stream was closed, we should also indicate we are closed // (and whoever is using the stream should panic) Poll::Ready(None) => Poll::Ready(None), @@ -416,7 +415,7 @@ where self.on_close_connection(id); } - match Pin::new(&mut self.real_receiver).poll_next(cx) { + match Pin::new(&mut self.real_receiver).poll_recv(cx) { // in the case our real message channel stream was closed, we should also indicate we are closed // (and whoever is using the stream should panic) Poll::Ready(None) => Poll::Ready(None), diff --git a/clients/native/src/client/mod.rs b/clients/native/src/client/mod.rs index 91fb13bcfa..ecc057e62b 100644 --- a/clients/native/src/client/mod.rs +++ b/clients/native/src/client/mod.rs @@ -304,28 +304,43 @@ impl NymClient { /// EXPERIMENTAL DIRECT RUST API /// It's untested and there are absolutely no guarantees about it (but seems to have worked /// well enough in local tests) - pub fn send_message(&mut self, recipient: Recipient, message: Vec, with_reply_surb: bool) { + pub async fn send_message( + &mut self, + recipient: Recipient, + message: Vec, + with_reply_surb: bool, + ) { let lane = TransmissionLane::General; let input_msg = InputMessage::new_fresh(recipient, message, with_reply_surb, lane); - self.input_tx + if self + .input_tx .as_ref() .expect("start method was not called before!") - .unbounded_send(input_msg) - .unwrap(); + .send(input_msg) + .await + .is_err() + { + panic!(); + } } /// EXPERIMENTAL DIRECT RUST API /// It's untested and there are absolutely no guarantees about it (but seems to have worked /// well enough in local tests) - pub fn send_reply(&mut self, reply_surb: ReplySurb, message: Vec) { + pub async fn send_reply(&mut self, reply_surb: ReplySurb, message: Vec) { let input_msg = InputMessage::new_reply(reply_surb, message); - self.input_tx + if self + .input_tx .as_ref() .expect("start method was not called before!") - .unbounded_send(input_msg) - .unwrap(); + .send(input_msg) + .await + .is_err() + { + panic!(); + } } /// EXPERIMENTAL DIRECT RUST API @@ -382,7 +397,7 @@ impl NymClient { let (received_buffer_request_sender, received_buffer_request_receiver) = mpsc::unbounded(); // channels responsible for controlling real messages - let (input_sender, input_receiver) = mpsc::unbounded::(); + let (input_sender, input_receiver) = tokio::sync::mpsc::channel::(1); // channels responsible for controlling ack messages let (ack_sender, ack_receiver) = mpsc::unbounded(); diff --git a/clients/native/src/websocket/handler.rs b/clients/native/src/websocket/handler.rs index 9830b6fdd0..07b196354d 100644 --- a/clients/native/src/websocket/handler.rs +++ b/clients/native/src/websocket/handler.rs @@ -81,7 +81,7 @@ impl Handler { } } - fn handle_send( + async fn handle_send( &mut self, recipient: &Recipient, message: Vec, @@ -91,18 +91,26 @@ impl Handler { // the ack control is now responsible for chunking, etc. let lane = TransmissionLane::ConnectionId(connection_id); let input_msg = InputMessage::new_fresh(*recipient, message, with_reply_surb, lane); - self.msg_input.unbounded_send(input_msg).unwrap(); + if self.msg_input.send(input_msg).await.is_err() { + panic!(); + } None } - fn handle_reply(&mut self, reply_surb: ReplySurb, message: Vec) -> Option { + async fn handle_reply( + &mut self, + reply_surb: ReplySurb, + message: Vec, + ) -> Option { if message.len() > ReplySurb::max_msg_len(Default::default()) { return Some(ServerResponse::new_error(format!("too long message to put inside a reply SURB. Received: {} bytes and maximum is {} bytes", message.len(), ReplySurb::max_msg_len(Default::default())))); } let input_msg = InputMessage::new_reply(reply_surb, message); - self.msg_input.unbounded_send(input_msg).unwrap(); + if self.msg_input.send(input_msg).await.is_err() { + panic!(); + } None } @@ -118,24 +126,27 @@ impl Handler { None } - fn handle_request(&mut self, request: ClientRequest) -> Option { + async fn handle_request(&mut self, request: ClientRequest) -> Option { match request { ClientRequest::Send { recipient, message, with_reply_surb, connection_id, - } => self.handle_send(&recipient, message, with_reply_surb, connection_id), + } => { + self.handle_send(&recipient, message, with_reply_surb, connection_id) + .await + } ClientRequest::Reply { message, reply_surb, - } => self.handle_reply(reply_surb, message), + } => self.handle_reply(reply_surb, message).await, ClientRequest::SelfAddress => Some(self.handle_self_address()), ClientRequest::ClosedConnection(id) => self.handle_closed_connection(id), } } - fn handle_text_message(&mut self, msg: String) -> Option { + async fn handle_text_message(&mut self, msg: String) -> Option { debug!("Handling text message request"); trace!("Content: {:?}", msg); @@ -144,13 +155,13 @@ impl Handler { let response = match client_request { Err(err) => Some(ServerResponse::Error(err)), - Ok(req) => self.handle_request(req), + Ok(req) => self.handle_request(req).await, }; response.map(|resp| WsMessage::text(resp.into_text())) } - fn handle_binary_message(&mut self, msg: &[u8]) -> Option { + async fn handle_binary_message(&mut self, msg: &[u8]) -> Option { debug!("Handling binary message request"); self.received_response_type = ReceivedResponseType::Binary; @@ -158,19 +169,19 @@ impl Handler { let response = match client_request { Err(err) => Some(ServerResponse::Error(err)), - Ok(req) => self.handle_request(req), + Ok(req) => self.handle_request(req).await, }; response.map(|resp| WsMessage::Binary(resp.into_binary())) } - fn handle_ws_request(&mut self, raw_request: WsMessage) -> Option { + async fn handle_ws_request(&mut self, raw_request: WsMessage) -> Option { // apparently tungstenite auto-handles ping/pong/close messages so for now let's ignore // them and let's test that claim. If that's not the case, just copy code from // old version of this file. match raw_request { - WsMessage::Text(text_message) => self.handle_text_message(text_message), - WsMessage::Binary(binary_message) => self.handle_binary_message(&binary_message), + WsMessage::Text(text_message) => self.handle_text_message(text_message).await, + WsMessage::Binary(binary_message) => self.handle_binary_message(&binary_message).await, _ => None, } } @@ -232,7 +243,7 @@ impl Handler { break; } - if let Some(response) = self.handle_ws_request(socket_msg) { + if let Some(response) = self.handle_ws_request(socket_msg).await { if let Err(err) = self.send_websocket_response(response).await { warn!( "Failed to send message over websocket: {}. Assuming the connection is dead.", diff --git a/clients/socks5/src/client/mod.rs b/clients/socks5/src/client/mod.rs index b35d7fd9fe..5f6a261285 100644 --- a/clients/socks5/src/client/mod.rs +++ b/clients/socks5/src/client/mod.rs @@ -286,6 +286,7 @@ impl NymClient { buffer_requester: ReceivedBufferRequestSender, msg_input: InputMessageSender, closed_connection_tx: ClosedConnectionSender, + lane_queue_lengths: LaneQueueLengths, shutdown: ShutdownListener, ) { info!("Starting socks5 listener..."); @@ -298,6 +299,7 @@ impl NymClient { authenticator, self.config.get_provider_mix_address(), self.as_mix_recipient(), + lane_queue_lengths, shutdown, ); tokio::spawn(async move { @@ -372,7 +374,7 @@ impl NymClient { let (received_buffer_request_sender, received_buffer_request_receiver) = mpsc::unbounded(); // channels responsible for controlling real messages - let (input_sender, input_receiver) = mpsc::unbounded::(); + let (input_sender, input_receiver) = tokio::sync::mpsc::channel::(1); // channels responsible for controlling ack messages let (ack_sender, ack_receiver) = mpsc::unbounded(); @@ -422,7 +424,7 @@ impl NymClient { input_receiver, sphinx_message_sender.clone(), closed_connection_rx, - shared_lane_queue_lengths, + shared_lane_queue_lengths.clone(), shutdown.subscribe(), ); @@ -442,6 +444,7 @@ impl NymClient { received_buffer_request_sender, input_sender, closed_connection_tx, + shared_lane_queue_lengths, shutdown.subscribe(), ); diff --git a/clients/socks5/src/socks/client.rs b/clients/socks5/src/socks/client.rs index 5188f7a15d..e77dba2975 100644 --- a/clients/socks5/src/socks/client.rs +++ b/clients/socks5/src/socks/client.rs @@ -4,7 +4,7 @@ use super::authentication::{AuthenticationMethods, Authenticator, User}; use super::request::{SocksCommand, SocksRequest}; use super::types::{ResponseCode, SocksProxyError}; use super::{RESERVED, SOCKS_VERSION}; -use client_connections::TransmissionLane; +use client_connections::{LaneQueueLengths, TransmissionLane}; use client_core::client::inbound_messages::{InputMessage, InputMessageSender}; use futures::channel::mpsc; use futures::task::{Context, Poll}; @@ -141,6 +141,7 @@ pub(crate) struct SocksClient { service_provider: Recipient, self_address: Recipient, started_proxy: bool, + lane_queue_lengths: LaneQueueLengths, shutdown_listener: ShutdownListener, } @@ -158,6 +159,7 @@ impl Drop for SocksClient { impl SocksClient { /// Create a new SOCKClient + #[allow(clippy::too_many_arguments)] pub fn new( stream: TcpStream, authenticator: Authenticator, @@ -165,6 +167,7 @@ impl SocksClient { service_provider: Recipient, controller_sender: ControllerSender, self_address: Recipient, + lane_queue_lengths: LaneQueueLengths, shutdown_listener: ShutdownListener, ) -> Self { let connection_id = Self::generate_random(); @@ -179,6 +182,7 @@ impl SocksClient { service_provider, self_address, started_proxy: false, + lane_queue_lengths, shutdown_listener, } } @@ -226,7 +230,7 @@ impl SocksClient { } } - fn send_connect_to_mixnet(&mut self, remote_address: RemoteAddress) { + async fn send_connect_to_mixnet(&mut self, remote_address: RemoteAddress) { let req = Request::new_connect(self.connection_id, remote_address, self.self_address); let msg = Message::Request(req); @@ -236,11 +240,14 @@ impl SocksClient { false, TransmissionLane::ConnectionId(self.connection_id), ); - self.input_sender.unbounded_send(input_message).unwrap(); + if self.input_sender.send(input_message).await.is_err() { + panic!(); + } } async fn run_proxy(&mut self, conn_receiver: ConnectionReceiver, remote_proxy_target: String) { - self.send_connect_to_mixnet(remote_proxy_target.clone()); + self.send_connect_to_mixnet(remote_proxy_target.clone()) + .await; let stream = self.stream.run_proxy(); let local_stream_remote = stream @@ -258,6 +265,7 @@ impl SocksClient { conn_receiver, input_sender, connection_id, + Some(self.lane_queue_lengths.clone()), self.shutdown_listener.clone(), ) .run(move |conn_id, read_data, socket_closed| { diff --git a/clients/socks5/src/socks/server.rs b/clients/socks5/src/socks/server.rs index 11912d0a0d..af5a6f9fed 100644 --- a/clients/socks5/src/socks/server.rs +++ b/clients/socks5/src/socks/server.rs @@ -4,7 +4,7 @@ use super::{ mixnet_responses::MixnetResponseListener, types::{ResponseCode, SocksProxyError}, }; -use client_connections::ClosedConnectionSender; +use client_connections::{ClosedConnectionSender, LaneQueueLengths}; use client_core::client::{ inbound_messages::InputMessageSender, received_buffer::ReceivedBufferRequestSender, }; @@ -21,6 +21,7 @@ pub struct SphinxSocksServer { listening_address: SocketAddr, service_provider: Recipient, self_address: Recipient, + lane_queue_lengths: LaneQueueLengths, shutdown: ShutdownListener, } @@ -31,6 +32,7 @@ impl SphinxSocksServer { authenticator: Authenticator, service_provider: Recipient, self_address: Recipient, + lane_queue_lengths: LaneQueueLengths, shutdown: ShutdownListener, ) -> Self { // hardcode ip as we (presumably) ONLY want to listen locally. If we change it, we can @@ -42,6 +44,7 @@ impl SphinxSocksServer { listening_address: format!("{}:{}", ip, port).parse().unwrap(), service_provider, self_address, + lane_queue_lengths, shutdown, } } @@ -85,6 +88,7 @@ impl SphinxSocksServer { self.service_provider, controller_sender.clone(), self.self_address, + self.lane_queue_lengths.clone(), self.shutdown.clone(), ); diff --git a/clients/webassembly/Cargo.toml b/clients/webassembly/Cargo.toml index 708c2aeb1c..2833598a16 100644 --- a/clients/webassembly/Cargo.toml +++ b/clients/webassembly/Cargo.toml @@ -19,13 +19,14 @@ coconut = ["coconut-interface", "credentials", "gateway-client/coconut"] [dependencies] futures = "0.3" -serde = { version = "1.0", features = ["derive"] } -serde-wasm-bindgen = "0.4" -wasm-bindgen = { version = "=0.2.83", features = ["serde-serialize"] } -wasm-bindgen-futures = "0.4" js-sys = "0.3" rand = { version = "0.7.3", features = ["wasm-bindgen"] } +serde = { version = "1.0", features = ["derive"] } +serde-wasm-bindgen = "0.4" +tokio = { version = "1.21.2", features = ["sync"] } url = "2.2" +wasm-bindgen = { version = "=0.2.83", features = ["serde-serialize"] } +wasm-bindgen-futures = "0.4" # internal client-core = { path = "../client-core", default-features = false, features = ["wasm"] } diff --git a/clients/webassembly/src/client/mod.rs b/clients/webassembly/src/client/mod.rs index e49d6199db..654989644e 100644 --- a/clients/webassembly/src/client/mod.rs +++ b/clients/webassembly/src/client/mod.rs @@ -326,7 +326,7 @@ impl NymClient { let (received_buffer_request_sender, received_buffer_request_receiver) = mpsc::unbounded(); // channels responsible for controlling real messages - let (input_sender, input_receiver) = mpsc::unbounded::(); + let (input_sender, input_receiver) = tokio::sync::mpsc::channel::(1); // channels responsible for controlling ack messages let (ack_sender, ack_receiver) = mpsc::unbounded(); @@ -395,11 +395,16 @@ impl NymClient { let input_msg = InputMessage::new_fresh(recipient, message, false, lane); - self.input_tx + if self + .input_tx .as_ref() .expect("start method was not called before!") - .unbounded_send(input_msg) - .unwrap(); + .send(input_msg) + .await + .is_err() + { + panic!(); + } self } diff --git a/common/client-connections/src/lib.rs b/common/client-connections/src/lib.rs index 265e21f9fd..d1782e8d90 100644 --- a/common/client-connections/src/lib.rs +++ b/common/client-connections/src/lib.rs @@ -24,7 +24,7 @@ pub type ClosedConnectionReceiver = mpsc::UnboundedReceiver; // The `OutQueueControl` publishes the backlog per lane, primarily so that upstream can slow down // if needed. -#[derive(Clone)] +#[derive(Clone, Debug)] pub struct LaneQueueLengths(std::sync::Arc>); impl LaneQueueLengths { @@ -52,6 +52,16 @@ impl LaneQueueLengths { Err(err) => log::warn!("Failed to set lane queue length: {err}"), } } + + pub fn get(&self, lane: &TransmissionLane) -> Option { + match self.0.lock() { + Ok(inner) => inner.get(lane), + Err(err) => { + log::warn!("Failed to get lane queue length: {err}"); + None + } + } + } } impl Default for LaneQueueLengths { @@ -68,6 +78,24 @@ impl std::ops::Deref for LaneQueueLengths { } } +#[derive(Debug)] pub struct LaneQueueLengthsInner { - map: HashMap, + pub map: HashMap, +} + +impl LaneQueueLengthsInner { + pub fn get(&self, lane: &TransmissionLane) -> Option { + self.map.get(lane).copied() + } + + pub fn values(&self) -> impl Iterator { + self.map.values() + } + + pub fn modify(&mut self, lane: &TransmissionLane, f: F) + where + F: FnOnce(&mut usize), + { + self.map.entry(*lane).and_modify(f); + } } diff --git a/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs b/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs index e116ba1dc5..ef0a09336d 100644 --- a/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs +++ b/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs @@ -5,17 +5,20 @@ use super::MixProxySender; use super::SHUTDOWN_TIMEOUT; use crate::available_reader::AvailableReader; use bytes::Bytes; +use client_connections::LaneQueueLengths; +use client_connections::TransmissionLane; use futures::FutureExt; use futures::StreamExt; use log::*; use ordered_buffer::OrderedMessageSender; use socks5_requests::ConnectionId; +use std::time::Duration; use std::{io, sync::Arc}; use task::ShutdownListener; use tokio::select; use tokio::{net::tcp::OwnedReadHalf, sync::Notify, time::sleep}; -fn send_empty_close( +async fn send_empty_close( connection_id: ConnectionId, message_sender: &mut OrderedMessageSender, mix_sender: &MixProxySender, @@ -24,12 +27,17 @@ fn send_empty_close( F: Fn(ConnectionId, Vec, bool) -> S, { let ordered_msg = message_sender.wrap_message(Vec::new()).into_bytes(); - mix_sender - .unbounded_send(adapter_fn(connection_id, ordered_msg, true)) - .unwrap(); + if mix_sender + .send(adapter_fn(connection_id, ordered_msg, true)) + .await + .is_err() + { + panic!(); + } } -fn deal_with_data( +#[allow(clippy::too_many_arguments)] +async fn deal_with_data( read_data: Option>, local_destination_address: &str, remote_source_address: &str, @@ -37,6 +45,7 @@ fn deal_with_data( message_sender: &mut OrderedMessageSender, mix_sender: &MixProxySender, adapter_fn: F, + lane_queue_lengths: Option, ) -> bool where F: Fn(ConnectionId, Vec, bool) -> S, @@ -67,9 +76,26 @@ where "pushing data down the input sender: size: {}", ordered_msg.len() ); - mix_sender - .unbounded_send(adapter_fn(connection_id, ordered_msg, is_finished)) - .unwrap(); + + // If we are closing the socket, wait until the data has passed `OutQueueControl` and the lane + // is empty, otherwise just wait until we are reasonably close to finish sending as a way to + // throttle the incoming data. + if let Some(lane_queue_lengths) = lane_queue_lengths { + if is_finished { + wait_until_lane_empty(lane_queue_lengths, connection_id).await; + } else { + // We allow a bit of slack when this is not the last msg + wait_until_lane_almost_empty(lane_queue_lengths, connection_id).await; + } + } + + if mix_sender + .send(adapter_fn(connection_id, ordered_msg, is_finished)) + .await + .is_err() + { + panic!(); + } if is_finished { // technically we already informed it when we sent the message to mixnet above @@ -79,6 +105,41 @@ where is_finished } +async fn wait_until_lane_empty(lane_queue_lengths: LaneQueueLengths, connection_id: u64) { + wait_for_lane( + lane_queue_lengths, + connection_id, + 0, + Duration::from_millis(500), + ) + .await +} + +async fn wait_until_lane_almost_empty(lane_queue_lengths: LaneQueueLengths, connection_id: u64) { + wait_for_lane( + lane_queue_lengths, + connection_id, + 10, + Duration::from_millis(100), + ) + .await +} + +async fn wait_for_lane( + lane_queue_lengths: LaneQueueLengths, + connection_id: u64, + queue_length_threshold: usize, + sleep_duration: Duration, +) { + while let Some(queue) = lane_queue_lengths.get(&TransmissionLane::ConnectionId(connection_id)) { + if queue > queue_length_threshold { + sleep(sleep_duration).await; + } else { + break; + } + } +} + #[allow(clippy::too_many_arguments)] pub(super) async fn run_inbound( mut reader: OwnedReadHalf, @@ -88,6 +149,7 @@ pub(super) async fn run_inbound( mix_sender: MixProxySender, adapter_fn: F, shutdown_notify: Arc, + lane_queue_lengths: Option, mut shutdown_listener: ShutdownListener, ) -> OwnedReadHalf where @@ -102,15 +164,27 @@ where loop { select! { read_data = &mut available_reader.next() => { - if deal_with_data(read_data, &local_destination_address, &remote_source_address, connection_id, &mut message_sender, &mix_sender, &adapter_fn) { + if deal_with_data( + read_data, + &local_destination_address, + &remote_source_address, + connection_id, + &mut message_sender, + &mix_sender, + &adapter_fn, + lane_queue_lengths.clone() + ).await { break } } _ = &mut shutdown_future => { - debug!("closing inbound proxy after outbound was closed {:?} ago", SHUTDOWN_TIMEOUT); + debug!( + "closing inbound proxy after outbound was closed {:?} ago", + SHUTDOWN_TIMEOUT + ); // 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); + send_empty_close(connection_id, &mut message_sender, &mix_sender, &adapter_fn).await; break; } _ = shutdown_listener.recv() => { diff --git a/common/socks5/proxy-helpers/src/proxy_runner/mod.rs b/common/socks5/proxy-helpers/src/proxy_runner/mod.rs index 1d5b493ced..573f09f884 100644 --- a/common/socks5/proxy-helpers/src/proxy_runner/mod.rs +++ b/common/socks5/proxy-helpers/src/proxy_runner/mod.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use crate::connection_controller::ConnectionReceiver; -use futures::channel::mpsc; +use client_connections::LaneQueueLengths; use socks5_requests::ConnectionId; use std::{sync::Arc, time::Duration}; use task::ShutdownListener; @@ -29,7 +29,8 @@ impl From<(Vec, bool)> for ProxyMessage { } } -pub type MixProxySender = mpsc::UnboundedSender; +pub type MixProxySender = tokio::sync::mpsc::Sender; +pub type MixProxyReader = tokio::sync::mpsc::Receiver; // TODO: when we finally get to implementing graceful shutdown, // on Drop this guy should tell the remote that it's closed now @@ -45,6 +46,7 @@ pub struct ProxyRunner { local_destination_address: String, remote_source_address: String, connection_id: ConnectionId, + lane_queue_lengths: Option, // Listens to shutdown commands from higher up shutdown_listener: ShutdownListener, @@ -54,6 +56,7 @@ impl ProxyRunner where S: Send + 'static, { + #[allow(clippy::too_many_arguments)] pub fn new( socket: TcpStream, local_destination_address: String, // addresses are provided for better logging @@ -61,6 +64,7 @@ where mix_receiver: ConnectionReceiver, mix_sender: MixProxySender, connection_id: ConnectionId, + lane_queue_lengths: Option, shutdown_listener: ShutdownListener, ) -> Self { ProxyRunner { @@ -70,6 +74,7 @@ where local_destination_address, remote_source_address, connection_id, + lane_queue_lengths, shutdown_listener, } } @@ -78,7 +83,7 @@ where // request/response as required by entity running particular side of the proxy. pub async fn run(mut self, adapter_fn: F) -> Self where - F: Fn(ConnectionId, Vec, bool) -> S + Send + 'static, + F: Fn(ConnectionId, Vec, bool) -> S + Send + Sync + 'static, { let (read_half, write_half) = self.socket.take().unwrap().into_split(); let shutdown_notify = Arc::new(Notify::new()); @@ -92,6 +97,7 @@ where self.mix_sender.clone(), adapter_fn, Arc::clone(&shutdown_notify), + self.lane_queue_lengths.clone(), self.shutdown_listener.clone(), ); diff --git a/service-providers/network-requester/src/connection.rs b/service-providers/network-requester/src/connection.rs index 5c5fea04b4..6627fdbf4b 100644 --- a/service-providers/network-requester/src/connection.rs +++ b/service-providers/network-requester/src/connection.rs @@ -1,10 +1,9 @@ // Copyright 2020 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use futures::channel::mpsc; use nymsphinx::addressing::clients::Recipient; use proxy_helpers::connection_controller::ConnectionReceiver; -use proxy_helpers::proxy_runner::ProxyRunner; +use proxy_helpers::proxy_runner::{MixProxySender, ProxyRunner}; use socks5_requests::{ConnectionId, Message as Socks5Message, RemoteAddress, Response}; use std::io; use task::ShutdownListener; @@ -40,7 +39,7 @@ impl Connection { pub(crate) async fn run_proxy( &mut self, mix_receiver: ConnectionReceiver, - mix_sender: mpsc::UnboundedSender<(Socks5Message, Recipient)>, + mix_sender: MixProxySender<(Socks5Message, Recipient)>, shutdown: ShutdownListener, ) { let stream = self.conn.take().unwrap(); @@ -54,6 +53,7 @@ impl Connection { mix_receiver, mix_sender, connection_id, + None, shutdown, ) .run(move |conn_id, read_data, socket_closed| { diff --git a/service-providers/network-requester/src/core.rs b/service-providers/network-requester/src/core.rs index c9333d6bf6..88d66d159d 100644 --- a/service-providers/network-requester/src/core.rs +++ b/service-providers/network-requester/src/core.rs @@ -15,6 +15,7 @@ use log::*; use nymsphinx::addressing::clients::Recipient; use nymsphinx::receiver::ReconstructedMessage; use proxy_helpers::connection_controller::{Controller, ControllerCommand, ControllerSender}; +use proxy_helpers::proxy_runner::{MixProxyReader, MixProxySender}; use socks5_requests::{ ConnectionId, Message as Socks5Message, NetworkRequesterResponse, Request, Response, }; @@ -67,14 +68,14 @@ impl ServiceProvider { /// via the `websocket_writer`. async fn mixnet_response_listener( mut websocket_writer: SplitSink, - mut mix_reader: mpsc::UnboundedReceiver<(Socks5Message, Recipient)>, + mut mix_reader: MixProxyReader<(Socks5Message, Recipient)>, stats_collector: Option, mut closed_connection_rx: ClosedConnectionReceiver, ) { loop { tokio::select! { // TODO: wire SURBs in here once they're available - socks5_msg = mix_reader.next() => { + socks5_msg = mix_reader.recv() => { if let Some((msg, return_address)) = socks5_msg { if let Some(stats_collector) = stats_collector.as_ref() { if let Some(remote_addr) = stats_collector @@ -153,7 +154,7 @@ impl ServiceProvider { remote_addr: String, return_address: Recipient, controller_sender: ControllerSender, - mix_input_sender: mpsc::UnboundedSender<(Socks5Message, Recipient)>, + mix_input_sender: MixProxySender<(Socks5Message, Recipient)>, shutdown: ShutdownListener, ) { let mut conn = match Connection::new(conn_id, remote_addr.clone(), return_address).await { @@ -166,12 +167,16 @@ impl ServiceProvider { ); // inform the remote that the connection is closed before it even was established - mix_input_sender - .unbounded_send(( + if mix_input_sender + .send(( Socks5Message::Response(Response::new(conn_id, Vec::new(), true)), return_address, )) - .unwrap(); + .await + .is_err() + { + panic!(); + } return; } @@ -207,10 +212,10 @@ impl ServiceProvider { ); } - fn handle_proxy_connect( + async fn handle_proxy_connect( &mut self, controller_sender: &mut ControllerSender, - mix_input_sender: &mpsc::UnboundedSender<(Socks5Message, Recipient)>, + mix_input_sender: &MixProxySender<(Socks5Message, Recipient)>, conn_id: ConnectionId, remote_addr: String, return_address: Recipient, @@ -219,14 +224,18 @@ impl ServiceProvider { if !self.open_proxy && !self.outbound_request_filter.check(&remote_addr) { let log_msg = format!("Domain {:?} failed filter check", remote_addr); log::info!("{}", log_msg); - mix_input_sender - .unbounded_send(( + if mix_input_sender + .send(( Socks5Message::NetworkRequesterResponse(NetworkRequesterResponse::new( conn_id, log_msg, )), return_address, )) - .unwrap(); + .await + .is_err() + { + panic!(); + } return; } @@ -263,7 +272,7 @@ impl ServiceProvider { &mut self, raw_request: &[u8], controller_sender: &mut ControllerSender, - mix_input_sender: &mpsc::UnboundedSender<(Socks5Message, Recipient)>, + mix_input_sender: &MixProxySender<(Socks5Message, Recipient)>, stats_collector: Option, shutdown: ShutdownListener, ) { @@ -292,6 +301,7 @@ impl ServiceProvider { req.return_address, shutdown, ) + .await } Request::Send(conn_id, data, closed) => { @@ -326,7 +336,7 @@ impl ServiceProvider { // channels responsible for managing messages that are to be sent to the mix network. The receiver is // going to be used by `mixnet_response_listener` let (mix_input_sender, mix_input_receiver) = - mpsc::unbounded::<(Socks5Message, Recipient)>(); + tokio::sync::mpsc::channel::<(Socks5Message, Recipient)>(1); // Used to notify tasks to shutdown. Not all tasks fully supports this (yet). let shutdown = task::ShutdownNotifier::default(); diff --git a/service-providers/network-requester/src/statistics/collector.rs b/service-providers/network-requester/src/statistics/collector.rs index 96e7292d82..72fc276860 100644 --- a/service-providers/network-requester/src/statistics/collector.rs +++ b/service-providers/network-requester/src/statistics/collector.rs @@ -2,8 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 use async_trait::async_trait; -use futures::channel::mpsc; use log::*; +use proxy_helpers::proxy_runner::MixProxySender; use rand::RngCore; use serde::Deserialize; use sqlx::types::chrono::{DateTime, Utc}; @@ -77,13 +77,13 @@ pub struct ServiceStatisticsCollector { pub(crate) response_stats_data: Arc>, pub(crate) connected_services: Arc>>, stats_provider_addr: Recipient, - mix_input_sender: mpsc::UnboundedSender<(Socks5Message, Recipient)>, + mix_input_sender: MixProxySender<(Socks5Message, Recipient)>, } impl ServiceStatisticsCollector { pub async fn new( stats_provider_addr: Option, - mix_input_sender: mpsc::UnboundedSender<(Socks5Message, Recipient)>, + mix_input_sender: MixProxySender<(Socks5Message, Recipient)>, ) -> Result { let client = reqwest::Client::builder() .timeout(Duration::from_secs(3)) @@ -175,20 +175,30 @@ impl StatisticsCollector for ServiceStatisticsCollector { ), self.stats_provider_addr, ); - self.mix_input_sender - .unbounded_send(( + if self + .mix_input_sender + .send(( Socks5Message::Request(connect_req), self.stats_provider_addr, )) - .unwrap(); + .await + .is_err() + { + panic!(); + } trace!("Sending data to statistics service"); let mut message_sender = OrderedMessageSender::new(); let ordered_msg = message_sender.wrap_message(msg).into_bytes(); let send_req = Request::new_send(conn_id, ordered_msg, true); - self.mix_input_sender - .unbounded_send((Socks5Message::Request(send_req), self.stats_provider_addr)) - .unwrap(); + if self + .mix_input_sender + .send((Socks5Message::Request(send_req), self.stats_provider_addr)) + .await + .is_err() + { + panic!(); + } Ok(()) }