From f3c1ff02e247f160c0441bedff4cb2ffd256438d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= Date: Wed, 23 Nov 2022 12:03:58 +0100 Subject: [PATCH] Network-requester: throttle inbound connections (#1789) * Return and handle ClientRequest::LaneQueueLenghts * Pass lane queue lengths to inbound future * Remove unused self reference * Request lane queue lengths periodically for all open connections * Add timeouts * Rename to ConnectionCommandSender and Receiver * Rename to client_connection_tx/rx * Fix wasm build * Replace bool with enum --- .../src/client/real_messages_control/mod.rs | 6 +- .../real_traffic_stream.rs | 24 ++-- clients/native/src/client/mod.rs | 23 ++-- clients/native/src/websocket/handler.rs | 43 +++++-- .../native/websocket-requests/src/requests.rs | 27 +++++ clients/socks5/src/client/mod.rs | 16 +-- clients/socks5/src/socks/server.rs | 13 ++- clients/webassembly/src/client/mod.rs | 10 +- common/client-connections/src/lib.rs | 20 +++- .../src/connection_controller.rs | 49 ++++++-- .../proxy-helpers/src/proxy_runner/inbound.rs | 34 ++++-- .../network-requester/src/connection.rs | 4 +- .../network-requester/src/core.rs | 107 ++++++++++++++---- 13 files changed, 282 insertions(+), 94 deletions(-) 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 d8b9dd34a6..ffeeea61e3 100644 --- a/clients/client-core/src/client/real_messages_control/mod.rs +++ b/clients/client-core/src/client/real_messages_control/mod.rs @@ -16,7 +16,7 @@ use crate::{ }, spawn_future, }; -use client_connections::{ClosedConnectionReceiver, LaneQueueLengths}; +use client_connections::{ConnectionCommandReceiver, LaneQueueLengths}; use futures::channel::mpsc; use gateway_client::AcknowledgementReceiver; use log::*; @@ -115,7 +115,7 @@ impl RealMessagesController { topology_access: TopologyAccessor, #[cfg(feature = "reply-surb")] reply_key_storage: ReplyKeyStorage, lane_queue_lengths: LaneQueueLengths, - closed_connection_rx: ClosedConnectionReceiver, + client_connection_rx: ConnectionCommandReceiver, ) -> Self { let rng = OsRng; @@ -166,7 +166,7 @@ impl RealMessagesController { config.self_recipient, topology_access, lane_queue_lengths, - closed_connection_rx, + client_connection_rx, ); RealMessagesController { 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 d6f34a5460..68b078e53c 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 @@ -5,7 +5,7 @@ use crate::client::mix_traffic::BatchMixMessageSender; use crate::client::real_messages_control::acknowledgement_control::SentPacketNotificationSender; use crate::client::topology_control::TopologyAccessor; use client_connections::{ - ClosedConnectionReceiver, ConnectionId, LaneQueueLengths, TransmissionLane, + ConnectionCommand, ConnectionCommandReceiver, ConnectionId, LaneQueueLengths, TransmissionLane, }; use futures::task::{Context, Poll}; use futures::{Future, Stream, StreamExt}; @@ -134,7 +134,7 @@ where /// Incoming channel for being notified of closed connections, so that we can close lanes /// corresponding to connections. To avoid sending traffic unnecessary - closed_connection_rx: ClosedConnectionReceiver, + client_connection_rx: ConnectionCommandReceiver, /// Report queue lengths so that upstream can backoff sending data, and keep connections open. lane_queue_lengths: LaneQueueLengths, @@ -182,7 +182,7 @@ where our_full_destination: Recipient, topology_access: TopologyAccessor, lane_queue_lengths: LaneQueueLengths, - closed_connection_rx: ClosedConnectionReceiver, + client_connection_rx: ConnectionCommandReceiver, ) -> Self { OutQueueControl { config, @@ -196,7 +196,7 @@ where rng, topology_access, transmission_buffer: Default::default(), - closed_connection_rx, + client_connection_rx, lane_queue_lengths, } } @@ -262,7 +262,7 @@ where self.sent_notify(fragment_id); } - // In addition to closing connections on receiving messages throught closed_connection_rx, + // In addition to closing connections on receiving messages throught client_connection_rx, // also close connections when sufficiently stale. self.transmission_buffer.prune_stale_connections(); @@ -335,8 +335,11 @@ where // Start by checking if we have any incoming messages about closed connections // NOTE: this feels a bit iffy, the `OutQueueControl` is getting ripe for a rewrite to // something simpler. - if let Poll::Ready(Some(id)) = Pin::new(&mut self.closed_connection_rx).poll_next(cx) { - self.on_close_connection(id); + if let Poll::Ready(Some(id)) = Pin::new(&mut self.client_connection_rx).poll_next(cx) { + match id { + ConnectionCommand::Close(id) => self.on_close_connection(id), + ConnectionCommand::ActiveConnections(_) => panic!(), + } } if let Some(ref mut next_delay) = &mut self.next_delay { @@ -411,8 +414,11 @@ where fn poll_immediate(&mut self, cx: &mut Context<'_>) -> Poll> { // Start by checking if we have any incoming messages about closed connections - if let Poll::Ready(Some(id)) = Pin::new(&mut self.closed_connection_rx).poll_next(cx) { - self.on_close_connection(id); + if let Poll::Ready(Some(id)) = Pin::new(&mut self.client_connection_rx).poll_next(cx) { + match id { + ConnectionCommand::Close(id) => self.on_close_connection(id), + ConnectionCommand::ActiveConnections(_) => panic!(), + } } match Pin::new(&mut self.real_receiver).poll_recv(cx) { diff --git a/clients/native/src/client/mod.rs b/clients/native/src/client/mod.rs index 357293d17e..5e16cf0a0f 100644 --- a/clients/native/src/client/mod.rs +++ b/clients/native/src/client/mod.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use client_connections::{ - ClosedConnectionReceiver, ClosedConnectionSender, LaneQueueLengths, TransmissionLane, + ConnectionCommandReceiver, ConnectionCommandSender, LaneQueueLengths, TransmissionLane, }; use client_core::client::cover_traffic_stream::LoopCoverTrafficStream; use client_core::client::inbound_messages::{ @@ -123,7 +123,7 @@ impl NymClient { input_receiver: InputMessageReceiver, mix_sender: BatchMixMessageSender, lane_queue_lengths: LaneQueueLengths, - closed_connection_rx: ClosedConnectionReceiver, + client_connection_rx: ConnectionCommandReceiver, shutdown: ShutdownListener, ) { let mut controller_config = real_messages_control::Config::new( @@ -154,7 +154,7 @@ impl NymClient { topology_accessor, reply_key_storage, lane_queue_lengths, - closed_connection_rx, + client_connection_rx, ) .start_with_shutdown(shutdown); } @@ -299,15 +299,17 @@ impl NymClient { &self, buffer_requester: ReceivedBufferRequestSender, msg_input: InputMessageSender, - closed_connection_tx: ClosedConnectionSender, + shared_lane_queue_lengths: LaneQueueLengths, + client_connection_tx: ConnectionCommandSender, ) { info!("Starting websocket listener..."); let websocket_handler = websocket::Handler::new( msg_input, - closed_connection_tx, + client_connection_tx, buffer_requester, &self.as_mix_recipient(), + shared_lane_queue_lengths, ); websocket::Listener::new(self.config.get_listening_port()).start(websocket_handler); @@ -449,11 +451,11 @@ impl NymClient { // Channels that the websocket listener can use to signal downstream to the real traffic // controller that connections are closed. - let (closed_connection_tx, closed_connection_rx) = mpsc::unbounded(); + let (client_connection_tx, client_connection_rx) = mpsc::unbounded(); // Shared queue length data. Published by the `OutQueueController` in the client, and used // primarily to throttle incoming connections (e.g socks5 for attached network-requesters) - let shared_lane_queue_length = LaneQueueLengths::new(); + let shared_lane_queue_lengths = LaneQueueLengths::new(); self.start_real_traffic_controller( shared_topology_accessor.clone(), @@ -461,8 +463,8 @@ impl NymClient { ack_receiver, input_receiver, sphinx_message_sender.clone(), - shared_lane_queue_length, - closed_connection_rx, + shared_lane_queue_lengths.clone(), + client_connection_rx, shutdown.subscribe(), ); @@ -482,7 +484,8 @@ impl NymClient { SocketType::WebSocket => self.start_websocket_listener( received_buffer_request_sender, input_sender, - closed_connection_tx, + shared_lane_queue_lengths, + client_connection_tx, ), SocketType::None => { // if we did not start the socket, it means we're running (supposedly) in the native mode diff --git a/clients/native/src/websocket/handler.rs b/clients/native/src/websocket/handler.rs index 07b196354d..53e6a4d491 100644 --- a/clients/native/src/websocket/handler.rs +++ b/clients/native/src/websocket/handler.rs @@ -1,7 +1,9 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use client_connections::{ClosedConnectionSender, TransmissionLane}; +use client_connections::{ + ConnectionCommand, ConnectionCommandSender, LaneQueueLengths, TransmissionLane, +}; use client_core::client::{ inbound_messages::{InputMessage, InputMessageSender}, received_buffer::{ @@ -35,11 +37,12 @@ impl Default for ReceivedResponseType { pub(crate) struct Handler { msg_input: InputMessageSender, - closed_connection_tx: ClosedConnectionSender, + client_connection_tx: ConnectionCommandSender, buffer_requester: ReceivedBufferRequestSender, self_full_address: Recipient, socket: Option>, received_response_type: ReceivedResponseType, + lane_queue_lengths: LaneQueueLengths, } // clone is used to use handler on a new connection, which initially is `None` @@ -47,11 +50,12 @@ impl Clone for Handler { fn clone(&self) -> Self { Handler { msg_input: self.msg_input.clone(), - closed_connection_tx: self.closed_connection_tx.clone(), + client_connection_tx: self.client_connection_tx.clone(), buffer_requester: self.buffer_requester.clone(), self_full_address: self.self_full_address, socket: None, received_response_type: Default::default(), + lane_queue_lengths: self.lane_queue_lengths.clone(), } } } @@ -67,17 +71,19 @@ impl Drop for Handler { impl Handler { pub(crate) fn new( msg_input: InputMessageSender, - closed_connection_tx: ClosedConnectionSender, + client_connection_tx: ConnectionCommandSender, buffer_requester: ReceivedBufferRequestSender, self_full_address: &Recipient, + lane_queue_lengths: LaneQueueLengths, ) -> Self { Handler { msg_input, - closed_connection_tx, + client_connection_tx, buffer_requester, self_full_address: *self_full_address, socket: None, received_response_type: Default::default(), + lane_queue_lengths, } } @@ -95,6 +101,15 @@ impl Handler { panic!(); } + // on receiving a send, we reply back the current lane queue length for that connection id. + // Note that this does _NOT_ take into account the packets that have been received but not + // yet reach `OutQueueControl`, so it might be a tad low. + if let Ok(lane_queue_lengths) = self.lane_queue_lengths.lock() { + let queue_length = lane_queue_lengths.get(&lane).unwrap_or(0); + return Some(ServerResponse::LaneQueueLength(connection_id, queue_length)); + } + + log::warn!("Failed to get the lane queue length lock, not responding back with the current queue length"); None } @@ -120,12 +135,25 @@ impl Handler { } fn handle_closed_connection(&self, connection_id: u64) -> Option { - self.closed_connection_tx - .unbounded_send(connection_id) + self.client_connection_tx + .unbounded_send(ConnectionCommand::Close(connection_id)) .unwrap(); None } + fn handle_get_lane_queue_length(&self, connection_id: u64) -> Option { + let Ok(lane_queue_lengths) = self.lane_queue_lengths.lock() else { + log::warn!( + "Failed to get the lane queue length lock, not responding back with the current queue length" + ); + return None; + }; + + let lane = TransmissionLane::ConnectionId(connection_id); + let queue_length = lane_queue_lengths.get(&lane).unwrap_or(0); + Some(ServerResponse::LaneQueueLength(connection_id, queue_length)) + } + async fn handle_request(&mut self, request: ClientRequest) -> Option { match request { ClientRequest::Send { @@ -143,6 +171,7 @@ impl Handler { } => self.handle_reply(reply_surb, message).await, ClientRequest::SelfAddress => Some(self.handle_self_address()), ClientRequest::ClosedConnection(id) => self.handle_closed_connection(id), + ClientRequest::GetLaneQueueLength(id) => self.handle_get_lane_queue_length(id), } } diff --git a/clients/native/websocket-requests/src/requests.rs b/clients/native/websocket-requests/src/requests.rs index f681b4264c..9b5af6e70d 100644 --- a/clients/native/websocket-requests/src/requests.rs +++ b/clients/native/websocket-requests/src/requests.rs @@ -23,6 +23,9 @@ pub const SELF_ADDRESS_REQUEST_TAG: u8 = 0x02; /// Value tag representing [`ClosedConnection`] variant of the [`ClientRequest`] pub const CLOSED_CONNECTION_REQUEST_TAG: u8 = 0x03; +/// Value tag representing [`GetLaneQueueLength`] variant of the [`ClientRequest`] +pub const GET_LANE_QUEUE_LENGHT_TAG: u8 = 0x04; + #[allow(non_snake_case)] #[derive(Debug)] pub enum ClientRequest { @@ -39,6 +42,7 @@ pub enum ClientRequest { }, SelfAddress, ClosedConnection(u64), + GetLaneQueueLength(u64), } // we could have been parsing it directly TryFrom, but we want to retain @@ -241,6 +245,26 @@ impl ClientRequest { ClientRequest::ClosedConnection(connection_id) } + // GET_LANE_QUEUE_LENGHT_TAG + fn serialize_get_lane_queue_lengths(connection_id: u64) -> Vec { + let conn_id_bytes = connection_id.to_be_bytes(); + std::iter::once(GET_LANE_QUEUE_LENGHT_TAG) + .chain(conn_id_bytes.iter().copied()) + .collect() + } + + // GET_LANE_QUEUE_LENGHT_TAG + fn deserialize_get_lane_queue_length(b: &[u8]) -> Self { + // this MUST match because it was called by 'deserialize' + debug_assert_eq!(b[0], GET_LANE_QUEUE_LENGHT_TAG); + + let mut connection_id_bytes = [0u8; size_of::()]; + connection_id_bytes.copy_from_slice(&b[1..=size_of::()]); + let connection_id = u64::from_be_bytes(connection_id_bytes); + + ClientRequest::GetLaneQueueLength(connection_id) + } + pub fn serialize(self) -> Vec { match self { ClientRequest::Send { @@ -258,6 +282,8 @@ impl ClientRequest { ClientRequest::SelfAddress => Self::serialize_self_address(), ClientRequest::ClosedConnection(id) => Self::serialize_closed_connection(id), + + ClientRequest::GetLaneQueueLength(id) => Self::serialize_get_lane_queue_lengths(id), } } @@ -288,6 +314,7 @@ impl ClientRequest { REPLY_REQUEST_TAG => Self::deserialize_reply(b), SELF_ADDRESS_REQUEST_TAG => Ok(Self::deserialize_self_address(b)), CLOSED_CONNECTION_REQUEST_TAG => Ok(Self::deserialize_closed_connection(b)), + GET_LANE_QUEUE_LENGHT_TAG => Ok(Self::deserialize_get_lane_queue_length(b)), n => Err(error::Error::new( ErrorKind::UnknownRequest, format!("type {n}"), diff --git a/clients/socks5/src/client/mod.rs b/clients/socks5/src/client/mod.rs index ef07360e2b..c6b734a749 100644 --- a/clients/socks5/src/client/mod.rs +++ b/clients/socks5/src/client/mod.rs @@ -9,7 +9,7 @@ use crate::socks::{ authentication::{AuthenticationMethods, Authenticator, User}, server::SphinxSocksServer, }; -use client_connections::{ClosedConnectionReceiver, ClosedConnectionSender, LaneQueueLengths}; +use client_connections::{ConnectionCommandReceiver, ConnectionCommandSender, LaneQueueLengths}; use client_core::client::cover_traffic_stream::LoopCoverTrafficStream; use client_core::client::inbound_messages::{ InputMessage, InputMessageReceiver, InputMessageSender, @@ -120,7 +120,7 @@ impl NymClient { ack_receiver: AcknowledgementReceiver, input_receiver: InputMessageReceiver, mix_sender: BatchMixMessageSender, - closed_connection_rx: ClosedConnectionReceiver, + client_connection_rx: ConnectionCommandReceiver, lane_queue_lengths: LaneQueueLengths, shutdown: ShutdownListener, ) { @@ -152,7 +152,7 @@ impl NymClient { topology_accessor, reply_key_storage, lane_queue_lengths, - closed_connection_rx, + client_connection_rx, ) .start_with_shutdown(shutdown); } @@ -297,7 +297,7 @@ impl NymClient { &self, buffer_requester: ReceivedBufferRequestSender, msg_input: InputMessageSender, - closed_connection_tx: ClosedConnectionSender, + client_connection_tx: ConnectionCommandSender, lane_queue_lengths: LaneQueueLengths, shutdown: ShutdownListener, ) { @@ -316,7 +316,7 @@ impl NymClient { ); tokio::spawn(async move { sphinx_socks - .serve(msg_input, buffer_requester, closed_connection_tx) + .serve(msg_input, buffer_requester, client_connection_tx) .await }); } @@ -426,7 +426,7 @@ impl NymClient { // Channel for announcing closed (socks5) connections by the controller. // This will be forwarded to `OutQueueControl` - let (closed_connection_tx, closed_connection_rx) = mpsc::unbounded(); + let (client_connection_tx, client_connection_rx) = mpsc::unbounded(); // Shared queue length data. Published by the `OutQueueController` in the client, and used // primarily to throttle incoming connections @@ -438,7 +438,7 @@ impl NymClient { ack_receiver, input_receiver, sphinx_message_sender.clone(), - closed_connection_rx, + client_connection_rx, shared_lane_queue_lengths.clone(), shutdown.subscribe(), ); @@ -458,7 +458,7 @@ impl NymClient { self.start_socks5_listener( received_buffer_request_sender, input_sender, - closed_connection_tx, + client_connection_tx, shared_lane_queue_lengths, shutdown.subscribe(), ); diff --git a/clients/socks5/src/socks/server.rs b/clients/socks5/src/socks/server.rs index af5a6f9fed..f61d6052a0 100644 --- a/clients/socks5/src/socks/server.rs +++ b/clients/socks5/src/socks/server.rs @@ -4,13 +4,13 @@ use super::{ mixnet_responses::MixnetResponseListener, types::{ResponseCode, SocksProxyError}, }; -use client_connections::{ClosedConnectionSender, LaneQueueLengths}; +use client_connections::{ConnectionCommandSender, LaneQueueLengths}; use client_core::client::{ inbound_messages::InputMessageSender, received_buffer::ReceivedBufferRequestSender, }; use log::*; use nymsphinx::addressing::clients::Recipient; -use proxy_helpers::connection_controller::Controller; +use proxy_helpers::connection_controller::{BroadcastActiveConnections, Controller}; use std::net::SocketAddr; use task::ShutdownListener; use tokio::net::TcpListener; @@ -55,14 +55,17 @@ impl SphinxSocksServer { &mut self, input_sender: InputMessageSender, buffer_requester: ReceivedBufferRequestSender, - closed_connection_tx: ClosedConnectionSender, + client_connection_tx: ConnectionCommandSender, ) -> Result<(), SocksProxyError> { let listener = TcpListener::bind(self.listening_address).await.unwrap(); info!("Serving Connections..."); // controller for managing all active connections - let (mut active_streams_controller, controller_sender) = - Controller::new(closed_connection_tx, self.shutdown.clone()); + let (mut active_streams_controller, controller_sender) = Controller::new( + client_connection_tx, + BroadcastActiveConnections::Off, + self.shutdown.clone(), + ); tokio::spawn(async move { active_streams_controller.run().await; }); diff --git a/clients/webassembly/src/client/mod.rs b/clients/webassembly/src/client/mod.rs index 654989644e..608674a169 100644 --- a/clients/webassembly/src/client/mod.rs +++ b/clients/webassembly/src/client/mod.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use self::config::Config; -use client_connections::{ClosedConnectionReceiver, LaneQueueLengths, TransmissionLane}; +use client_connections::{ConnectionCommandReceiver, LaneQueueLengths, TransmissionLane}; use client_core::client::{ cover_traffic_stream::LoopCoverTrafficStream, inbound_messages::{InputMessage, InputMessageReceiver, InputMessageSender}, @@ -128,7 +128,7 @@ impl NymClient { ack_receiver: AcknowledgementReceiver, input_receiver: InputMessageReceiver, mix_sender: BatchMixMessageSender, - closed_connection_rx: ClosedConnectionReceiver, + client_connection_rx: ConnectionCommandReceiver, lane_queue_lengths: LaneQueueLengths, ) { let mut controller_config = real_messages_control::Config::new( @@ -155,7 +155,7 @@ impl NymClient { mix_sender, topology_accessor, lane_queue_lengths, - closed_connection_rx, + client_connection_rx, ) .start(); } @@ -334,7 +334,7 @@ impl NymClient { // Channel that the real traffix controller can listed to for closing connections. // Currently unused in the wasm client. - let (_closed_connection_tx, closed_connection_rx) = mpsc::unbounded(); + let (_client_connection_tx, client_connection_rx) = mpsc::unbounded(); // the components are started in very specific order. Unless you know what you are doing, // do not change that. @@ -364,7 +364,7 @@ impl NymClient { ack_receiver, input_receiver, sphinx_message_sender.clone(), - closed_connection_rx, + client_connection_rx, shared_lane_queue_lengths, ); diff --git a/common/client-connections/src/lib.rs b/common/client-connections/src/lib.rs index d1782e8d90..515938ed54 100644 --- a/common/client-connections/src/lib.rs +++ b/common/client-connections/src/lib.rs @@ -16,11 +16,21 @@ pub enum TransmissionLane { ConnectionId(ConnectionId), } -/// Announce connections that are closed, for whoever is interested. -/// One usecase is that the network-requester and socks5-client wants to know about this, so that -/// they can forward this to the `OutQueueControl` (via `ClientRequest` for the network-requester) -pub type ClosedConnectionSender = mpsc::UnboundedSender; -pub type ClosedConnectionReceiver = mpsc::UnboundedReceiver; +/// Used by the connection controller to report current state for client connections. +pub type ConnectionCommandSender = mpsc::UnboundedSender; +pub type ConnectionCommandReceiver = mpsc::UnboundedReceiver; + +pub enum ConnectionCommand { + // Announce that at a connection was closed. E.g the `OutQueueControl` uses this to discard + // transmission lanes. + Close(ConnectionId), + + // In the network requester for example, we usually want to broadcast active connections + // regularly, so we know what connections we need to request lane queue lengths for from the + // client. + // In the socks5-client, this is not needed since have direct access to the lane queue lengths. + ActiveConnections(Vec), +} // The `OutQueueControl` publishes the backlog per lane, primarily so that upstream can slow down // if needed. diff --git a/common/socks5/proxy-helpers/src/connection_controller.rs b/common/socks5/proxy-helpers/src/connection_controller.rs index d78915d2a5..12fdba7dbb 100644 --- a/common/socks5/proxy-helpers/src/connection_controller.rs +++ b/common/socks5/proxy-helpers/src/connection_controller.rs @@ -1,14 +1,18 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use client_connections::ClosedConnectionSender; +use client_connections::{ConnectionCommand, ConnectionCommandSender}; use futures::channel::mpsc; use futures::StreamExt; use log::*; use ordered_buffer::{OrderedMessage, OrderedMessageBuffer, ReadContiguousData}; use socks5_requests::ConnectionId; -use std::collections::{HashMap, HashSet}; +use std::{ + collections::{HashMap, HashSet}, + time::Duration, +}; use task::ShutdownListener; +use tokio::time; /// A generic message produced after reading from a socket/connection. It includes data that was /// actually read alongside boolean indicating whether the connection got closed so that @@ -64,6 +68,12 @@ impl ActiveConnection { } } +#[derive(PartialEq, Eq)] +pub enum BroadcastActiveConnections { + On, + Off, +} + /// Controller represents a way of managing multiple open connections that are used for socks5 /// proxy. pub struct Controller { @@ -75,7 +85,11 @@ pub struct Controller { recently_closed: HashSet, // Broadcast closed connections - closed_connection_tx: ClosedConnectionSender, + client_connection_tx: ConnectionCommandSender, + + // The controller can broadcast active connections. This is useful in the network-requester + // where its used to query the client for lane queue lengths + broadcast_connections: BroadcastActiveConnections, // TODO: this can potentially be abused to ddos and kill provider. Not sure at this point // how to handle it more gracefully @@ -89,7 +103,8 @@ pub struct Controller { impl Controller { pub fn new( - closed_connection_tx: ClosedConnectionSender, + client_connection_tx: ConnectionCommandSender, + broadcast_connections: BroadcastActiveConnections, shutdown: ShutdownListener, ) -> (Self, ControllerSender) { let (sender, receiver) = mpsc::unbounded(); @@ -98,7 +113,8 @@ impl Controller { active_connections: HashMap::new(), receiver, recently_closed: HashSet::new(), - closed_connection_tx, + client_connection_tx, + broadcast_connections, pending_messages: HashMap::new(), shutdown, }, @@ -137,7 +153,10 @@ impl Controller { self.recently_closed.insert(conn_id); // Announce closed connections, currently used by the `OutQueueControl`. - if let Err(err) = self.closed_connection_tx.unbounded_send(conn_id) { + if let Err(err) = self + .client_connection_tx + .unbounded_send(ConnectionCommand::Close(conn_id)) + { if self.shutdown.is_shutdown_poll() { log::debug!("Failed to send: {}", err); } else { @@ -146,6 +165,15 @@ impl Controller { } } + fn broadcast_active_connections(&mut self) { + // What about the recently closed ones? Hopefully we can ignore them ... + let conn_ids = self.active_connections.keys().copied().collect(); + + self.client_connection_tx + .unbounded_send(ConnectionCommand::ActiveConnections(conn_ids)) + .unwrap(); + } + fn send_to_connection(&mut self, conn_id: ConnectionId, payload: Vec, is_closed: bool) { if let Some(active_connection) = self.active_connections.get_mut(&conn_id) { if !payload.is_empty() { @@ -202,6 +230,8 @@ impl Controller { } pub async fn run(&mut self) { + let mut interval = time::interval(Duration::from_millis(500)); + loop { tokio::select! { command = self.receiver.next() => match command { @@ -216,7 +246,12 @@ impl Controller { log::trace!("SOCKS5 Controller: Stopping since channel closed"); break; } - } + }, + _ = interval.tick() => { + if self.broadcast_connections == BroadcastActiveConnections::On { + self.broadcast_active_connections(); + } + }, } } assert!(self.shutdown.is_shutdown_poll()); diff --git a/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs b/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs index ef0a09336d..bb5cf98f37 100644 --- a/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs +++ b/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs @@ -106,23 +106,37 @@ where } 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), + if tokio::time::timeout( + Duration::from_secs(2 * 60), + wait_for_lane( + lane_queue_lengths, + connection_id, + 0, + Duration::from_millis(500), + ), ) .await + .is_err() + { + log::warn!("Wait until lane empty timed out"); + } } 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), + if tokio::time::timeout( + Duration::from_secs(2 * 60), + wait_for_lane( + lane_queue_lengths, + connection_id, + 10, + Duration::from_millis(100), + ), ) .await + .is_err() + { + log::debug!("Wait until lane almost empty timed out"); + } } async fn wait_for_lane( diff --git a/service-providers/network-requester/src/connection.rs b/service-providers/network-requester/src/connection.rs index 6627fdbf4b..95bdac06a3 100644 --- a/service-providers/network-requester/src/connection.rs +++ b/service-providers/network-requester/src/connection.rs @@ -1,6 +1,7 @@ // Copyright 2020 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 +use client_connections::LaneQueueLengths; use nymsphinx::addressing::clients::Recipient; use proxy_helpers::connection_controller::ConnectionReceiver; use proxy_helpers::proxy_runner::{MixProxySender, ProxyRunner}; @@ -40,6 +41,7 @@ impl Connection { &mut self, mix_receiver: ConnectionReceiver, mix_sender: MixProxySender<(Socks5Message, Recipient)>, + lane_queue_lengths: LaneQueueLengths, shutdown: ShutdownListener, ) { let stream = self.conn.take().unwrap(); @@ -53,7 +55,7 @@ impl Connection { mix_receiver, mix_sender, connection_id, - None, + Some(lane_queue_lengths), 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 88d66d159d..6bea582581 100644 --- a/service-providers/network-requester/src/core.rs +++ b/service-providers/network-requester/src/core.rs @@ -7,14 +7,18 @@ use crate::error::NetworkRequesterError; use crate::statistics::ServiceStatisticsCollector; use crate::websocket; use crate::websocket::TSWebsocketStream; -use client_connections::ClosedConnectionReceiver; +use client_connections::{ + ConnectionCommand, ConnectionCommandReceiver, LaneQueueLengths, TransmissionLane, +}; use futures::channel::mpsc; use futures::stream::{SplitSink, SplitStream}; use futures::{SinkExt, StreamExt}; use log::*; use nymsphinx::addressing::clients::Recipient; use nymsphinx::receiver::ReconstructedMessage; -use proxy_helpers::connection_controller::{Controller, ControllerCommand, ControllerSender}; +use proxy_helpers::connection_controller::{ + BroadcastActiveConnections, Controller, ControllerCommand, ControllerSender, +}; use proxy_helpers::proxy_runner::{MixProxyReader, MixProxySender}; use socks5_requests::{ ConnectionId, Message as Socks5Message, NetworkRequesterResponse, Request, Response, @@ -70,7 +74,7 @@ impl ServiceProvider { mut websocket_writer: SplitSink, mut mix_reader: MixProxyReader<(Socks5Message, Recipient)>, stats_collector: Option, - mut closed_connection_rx: ClosedConnectionReceiver, + mut client_connection_rx: ConnectionCommandReceiver, ) { loop { tokio::select! { @@ -108,17 +112,49 @@ impl ServiceProvider { break; } }, - Some(id) = closed_connection_rx.next() => { - let msg = ClientRequest::ClosedConnection(id); - let ws_msg = Message::Binary(msg.serialize()); - websocket_writer.send(ws_msg).await.unwrap(); - } + Some(command) = client_connection_rx.next() => { + match command { + ConnectionCommand::Close(id) => { + let msg = ClientRequest::ClosedConnection(id); + let ws_msg = Message::Binary(msg.serialize()); + websocket_writer.send(ws_msg).await.unwrap(); + } + ConnectionCommand::ActiveConnections(ids) => { + // We can optimize this by sending a single request, but this is + // usually in the low single digits, max a few tens, so we leave that + // for a rainy day. + // Also that means fiddling with the currently manual + // serialize/deserialize we do with ClientRequests ... bleh + for id in ids { + log::trace!("Requesting lane queue length for: {}", id); + let msg = ClientRequest::GetLaneQueueLength(id); + let ws_msg = Message::Binary(msg.serialize()); + websocket_writer.send(ws_msg).await.unwrap(); + } + } + } + }, } } } + fn handle_lane_queue_length_response( + lane_queue_lengths: &LaneQueueLengths, + lane: u64, + queue_length: usize, + ) { + log::trace!("Received LaneQueueLength lane: {lane}, queue_length: {queue_length}"); + if let Ok(mut lane_queue_lengths) = lane_queue_lengths.lock() { + let lane = TransmissionLane::ConnectionId(lane); + lane_queue_lengths.map.insert(lane, queue_length); + } else { + log::warn!("Unable to lock lane queue lengths, skipping updating received lane length") + } + } + async fn read_websocket_message( websocket_reader: &mut SplitStream, + lane_queue_lengths: LaneQueueLengths, ) -> Option { while let Some(msg) = websocket_reader.next().await { let data = msg @@ -139,6 +175,14 @@ impl ServiceProvider { let received = match deserialized_message { ServerResponse::Received(received) => received, + ServerResponse::LaneQueueLength(lane, queue_length) => { + Self::handle_lane_queue_length_response( + &lane_queue_lengths, + lane, + queue_length, + ); + continue; + } ServerResponse::Error(err) => { panic!("received error from native client! - {}", err) } @@ -155,6 +199,7 @@ impl ServiceProvider { return_address: Recipient, controller_sender: ControllerSender, mix_input_sender: MixProxySender<(Socks5Message, Recipient)>, + lane_queue_lengths: LaneQueueLengths, shutdown: ShutdownListener, ) { let mut conn = match Connection::new(conn_id, remote_addr.clone(), return_address).await { @@ -196,7 +241,7 @@ impl ServiceProvider { ); // run the proxy on the connection - conn.run_proxy(mix_receiver, mix_input_sender, shutdown) + conn.run_proxy(mix_receiver, mix_input_sender, lane_queue_lengths, shutdown) .await; // proxy is done - remove the access channel from the controller @@ -212,10 +257,12 @@ impl ServiceProvider { ); } + #[allow(clippy::too_many_arguments)] async fn handle_proxy_connect( &mut self, controller_sender: &mut ControllerSender, mix_input_sender: &MixProxySender<(Socks5Message, Recipient)>, + lane_queue_lengths: LaneQueueLengths, conn_id: ConnectionId, remote_addr: String, return_address: Recipient, @@ -250,6 +297,7 @@ impl ServiceProvider { return_address, controller_sender_clone, mix_input_sender_clone, + lane_queue_lengths, shutdown, ) .await @@ -257,7 +305,6 @@ impl ServiceProvider { } fn handle_proxy_send( - &self, controller_sender: &mut ControllerSender, conn_id: ConnectionId, data: Vec, @@ -273,6 +320,7 @@ impl ServiceProvider { raw_request: &[u8], controller_sender: &mut ControllerSender, mix_input_sender: &MixProxySender<(Socks5Message, Recipient)>, + lane_queue_lengths: LaneQueueLengths, stats_collector: Option, shutdown: ShutdownListener, ) { @@ -296,6 +344,7 @@ impl ServiceProvider { self.handle_proxy_connect( controller_sender, mix_input_sender, + lane_queue_lengths, req.conn_id, req.remote_addr, req.return_address, @@ -319,7 +368,7 @@ impl ServiceProvider { .processed(remote_addr, data.len() as u32); } } - self.handle_proxy_send(controller_sender, conn_id, data, closed) + Self::handle_proxy_send(controller_sender, conn_id, data, closed) } }, Socks5Message::Response(_) | Socks5Message::NetworkRequesterResponse(_) => {} @@ -341,16 +390,23 @@ impl ServiceProvider { // Used to notify tasks to shutdown. Not all tasks fully supports this (yet). let shutdown = task::ShutdownNotifier::default(); - // Channel for announcing closed (socks5) connections by the controller. - // The `mixnet_response_listener` will forward this info to the client using a - // `ClientRequest`. - let (closed_connection_tx, closed_connection_rx) = mpsc::unbounded(); + // Channel for announcing client connection state by the controller. + // The `mixnet_response_listener` will use this to either report closed connection to the + // client or request lane queue lengths. + let (client_connection_tx, client_connection_rx) = mpsc::unbounded(); + + // Shared queue length data. Published by the `OutQueueController` in the client, and used + // primarily to throttle incoming connections + let shared_lane_queue_lengths = LaneQueueLengths::new(); // Controller for managing all active connections. // We provide it with a ShutdownListener since it requires it, even though for the network // requester shutdown signalling is not yet fully implemented. - let (mut active_connections_controller, mut controller_sender) = - Controller::new(closed_connection_tx, shutdown.subscribe()); + let (mut active_connections_controller, mut controller_sender) = Controller::new( + client_connection_tx, + BroadcastActiveConnections::On, + shutdown.subscribe(), + ); tokio::spawn(async move { active_connections_controller.run().await; @@ -378,7 +434,7 @@ impl ServiceProvider { websocket_writer, mix_input_receiver, stats_collector_clone, - closed_connection_rx, + client_connection_rx, ) .await; }); @@ -386,12 +442,14 @@ impl ServiceProvider { println!("\nAll systems go. Press CTRL-C to stop the server."); // for each incoming message from the websocket... (which in 99.99% cases is going to be a mix message) loop { - let received = match Self::read_websocket_message(&mut websocket_reader).await { - Some(msg) => msg, - None => { - error!("The websocket stream has finished!"); - return Ok(()); - } + let Some(received) = Self::read_websocket_message( + &mut websocket_reader, + shared_lane_queue_lengths.clone() + ) + .await + else { + log::error!("The websocket stream has finished!"); + return Ok(()); }; let raw_message = received.message; @@ -401,6 +459,7 @@ impl ServiceProvider { &raw_message, &mut controller_sender, &mix_input_sender, + shared_lane_queue_lengths.clone(), stats_collector.clone(), shutdown.subscribe(), )