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

@@ -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)

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<InputMessage>;
-pub type InputMessageReceiver = mpsc::UnboundedReceiver<InputMessage>;
+pub type InputMessageSender = tokio::sync::mpsc::Sender<InputMessage>;
+pub type InputMessageReceiver = tokio::sync::mpsc::Receiver<InputMessage>;
 
 #[derive(Debug)]
 pub enum InputMessage {

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;
         }
     }

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"))]

clients/client-core/src/client/real_messages_control/mod.rs

@@ -119,7 +119,7 @@ impl RealMessagesController<OsRng> {
     ) -> 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(

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<RealMessage>, TransmissionLane)>;
-type BatchRealMessageReceiver = mpsc::UnboundedReceiver<(Vec<RealMessage>, TransmissionLane)>;
+    tokio::sync::mpsc::Sender<(Vec<RealMessage>, TransmissionLane)>;
+type BatchRealMessageReceiver = tokio::sync::mpsc::Receiver<(Vec<RealMessage>, 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),

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<u8>, with_reply_surb: bool) {
+    pub async fn send_message(
+        &mut self,
+        recipient: Recipient,
+        message: Vec<u8>,
+        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<u8>) {
+    pub async fn send_reply(&mut self, reply_surb: ReplySurb, message: Vec<u8>) {
         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::<InputMessage>();
+        let (input_sender, input_receiver) = tokio::sync::mpsc::channel::<InputMessage>(1);
 
         // channels responsible for controlling ack messages
         let (ack_sender, ack_receiver) = mpsc::unbounded();

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<u8>,
@@ -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<u8>) -> Option<ServerResponse> {
+    async fn handle_reply(
+        &mut self,
+        reply_surb: ReplySurb,
+        message: Vec<u8>,
+    ) -> Option<ServerResponse> {
         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<ServerResponse> {
+    async fn handle_request(&mut self, request: ClientRequest) -> Option<ServerResponse> {
         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<WsMessage> {
+    async fn handle_text_message(&mut self, msg: String) -> Option<WsMessage> {
         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<WsMessage> {
+    async fn handle_binary_message(&mut self, msg: &[u8]) -> Option<WsMessage> {
         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<WsMessage> {
+    async fn handle_ws_request(&mut self, raw_request: WsMessage) -> Option<WsMessage> {
         // 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.",

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::<InputMessage>();
+        let (input_sender, input_receiver) = tokio::sync::mpsc::channel::<InputMessage>(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(),
         );
 

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| {

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(),
                     );
 

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"] }

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::<InputMessage>();
+        let (input_sender, input_receiver) = tokio::sync::mpsc::channel::<InputMessage>(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
     }

common/client-connections/src/lib.rs

@@ -24,7 +24,7 @@ pub type ClosedConnectionReceiver = mpsc::UnboundedReceiver<ConnectionId>;
 
 // 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<std::sync::Mutex<LaneQueueLengthsInner>>);
 
 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<usize> {
+        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<TransmissionLane, usize>,
+    pub map: HashMap<TransmissionLane, usize>,
+}
+
+impl LaneQueueLengthsInner {
+    pub fn get(&self, lane: &TransmissionLane) -> Option<usize> {
+        self.map.get(lane).copied()
+    }
+
+    pub fn values(&self) -> impl Iterator<Item = &usize> {
+        self.map.values()
+    }
+
+    pub fn modify<F>(&mut self, lane: &TransmissionLane, f: F)
+    where
+        F: FnOnce(&mut usize),
+    {
+        self.map.entry(*lane).and_modify(f);
+    }
 }

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<F, S>(
+async fn send_empty_close<F, S>(
     connection_id: ConnectionId,
     message_sender: &mut OrderedMessageSender,
     mix_sender: &MixProxySender<S>,
@@ -24,12 +27,17 @@ fn send_empty_close<F, S>(
     F: Fn(ConnectionId, Vec<u8>, 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<F, S>(
+#[allow(clippy::too_many_arguments)]
+async fn deal_with_data<F, S>(
     read_data: Option<io::Result<Bytes>>,
     local_destination_address: &str,
     remote_source_address: &str,
@@ -37,6 +45,7 @@ fn deal_with_data<F, S>(
     message_sender: &mut OrderedMessageSender,
     mix_sender: &MixProxySender<S>,
     adapter_fn: F,
+    lane_queue_lengths: Option<LaneQueueLengths>,
 ) -> bool
 where
     F: Fn(ConnectionId, Vec<u8>, 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<F, S>(
     mut reader: OwnedReadHalf,
@@ -88,6 +149,7 @@ pub(super) async fn run_inbound<F, S>(
     mix_sender: MixProxySender<S>,
     adapter_fn: F,
     shutdown_notify: Arc<Notify>,
+    lane_queue_lengths: Option<LaneQueueLengths>,
     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() => {

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<u8>, bool)> for ProxyMessage {
     }
 }
 
-pub type MixProxySender<S> = mpsc::UnboundedSender<S>;
+pub type MixProxySender<S> = tokio::sync::mpsc::Sender<S>;
+pub type MixProxyReader<S> = tokio::sync::mpsc::Receiver<S>;
 
 // 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<S> {
     local_destination_address: String,
     remote_source_address: String,
     connection_id: ConnectionId,
+    lane_queue_lengths: Option<LaneQueueLengths>,
 
     // Listens to shutdown commands from higher up
     shutdown_listener: ShutdownListener,
@@ -54,6 +56,7 @@ impl<S> ProxyRunner<S>
 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<S>,
         connection_id: ConnectionId,
+        lane_queue_lengths: Option<LaneQueueLengths>,
         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<F>(mut self, adapter_fn: F) -> Self
     where
-        F: Fn(ConnectionId, Vec<u8>, bool) -> S + Send + 'static,
+        F: Fn(ConnectionId, Vec<u8>, 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(),
         );
 

service-providers/network-requester/src/connection.rs

@@ -1,10 +1,9 @@
 // Copyright 2020 - Nym Technologies SA <contact@nymtech.net>
 // 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| {

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<TSWebsocketStream, Message>,
-        mut mix_reader: mpsc::UnboundedReceiver<(Socks5Message, Recipient)>,
+        mut mix_reader: MixProxyReader<(Socks5Message, Recipient)>,
         stats_collector: Option<ServiceStatisticsCollector>,
         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<ServiceStatisticsCollector>,
         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();

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<RwLock<StatsData>>,
     pub(crate) connected_services: Arc<RwLock<HashMap<ConnectionId, RemoteAddress>>>,
     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<Recipient>,
-        mix_input_sender: mpsc::UnboundedSender<(Socks5Message, Recipient)>,
+        mix_input_sender: MixProxySender<(Socks5Message, Recipient)>,
     ) -> Result<Self, StatsError> {
         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(())
     }