GRIN/nym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Network-requester: throttle inbound connections (#1789)

Browse Source 
* 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
This commit is contained in:
Jon Häggblad
2022-11-23 12:03:58 +01:00
committed by GitHub
parent f4fb0d6d6c
commit f3c1ff02e2
13 changed files with 282 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
clients/client-core/src/client/real_messages_control/mod.rs
+3 -3
View File
@@ -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<OsRng> {
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<OsRng> {
config.self_recipient,
topology_access,
lane_queue_lengths,
closed_connection_rx,
client_connection_rx,
);
RealMessagesController {
clients/client-core/src/client/real_messages_control/real_traffic_stream.rs
+15 -9
View File
@@ -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<Option<StreamMessage>> {
// 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) {
clients/native/src/client/mod.rs
+13 -10
View File
@@ -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
clients/native/src/websocket/handler.rs
+36 -7
View File
@@ -1,7 +1,9 @@
// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
// 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<WebSocketStream<TcpStream>>,
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<ServerResponse> {
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<ServerResponse> {
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<ServerResponse> {
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),
}
}
clients/native/websocket-requests/src/requests.rs
+27
View File
@@ -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<WsMessage>, 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<u8> {
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::<u64>()];
connection_id_bytes.copy_from_slice(&b[1..=size_of::<u64>()]);
let connection_id = u64::from_be_bytes(connection_id_bytes);
ClientRequest::GetLaneQueueLength(connection_id)
}
pub fn serialize(self) -> Vec<u8> {
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}"),
clients/socks5/src/client/mod.rs
+8 -8
View File
@@ -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(),
);
clients/socks5/src/socks/server.rs
+8 -5
View File
@@ -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;
});
clients/webassembly/src/client/mod.rs
+5 -5
View File
@@ -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,
);
common/client-connections/src/lib.rs
+15 -5
View File
@@ -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<ConnectionId>;
pub type ClosedConnectionReceiver = mpsc::UnboundedReceiver<ConnectionId>;
/// Used by the connection controller to report current state for client connections.
pub type ConnectionCommandSender = mpsc::UnboundedSender<ConnectionCommand>;
pub type ConnectionCommandReceiver = mpsc::UnboundedReceiver<ConnectionCommand>;
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<ConnectionId>),
}
// The `OutQueueControl` publishes the backlog per lane, primarily so that upstream can slow down
// if needed.
common/socks5/proxy-helpers/src/connection_controller.rs
+42 -7
View File
@@ -1,14 +1,18 @@
// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
// 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<ConnectionId>,
// 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<u8>, 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());
common/socks5/proxy-helpers/src/proxy_runner/inbound.rs
+24 -10
View File
@@ -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(
service-providers/network-requester/src/connection.rs
+3 -1
View File
@@ -1,6 +1,7 @@
// Copyright 2020 - Nym Technologies SA <contact@nymtech.net>
// 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| {
service-providers/network-requester/src/core.rs
+83 -24
View File
@@ -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<TSWebsocketStream, Message>,
mut mix_reader: MixProxyReader<(Socks5Message, Recipient)>,
stats_collector: Option<ServiceStatisticsCollector>,
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<TSWebsocketStream>,
lane_queue_lengths: LaneQueueLengths,
) -> Option<ReconstructedMessage> {
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<u8>,
@@ -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<ServiceStatisticsCollector>,
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(),
)