Merge pull request #3222 from nymtech/feature/available_reader_changes

Feature/available reader changes
This commit is contained in:
Tommy Verrall
2023-04-24 15:48:08 +01:00
committed by GitHub
8 changed files with 64 additions and 193 deletions
+14 -5
View File
@@ -7,15 +7,10 @@ use crate::socks::{
authentication::{AuthenticationMethods, Authenticator, User},
server::SphinxSocksServer,
};
use futures::channel::mpsc;
use futures::StreamExt;
use log::*;
use nym_bandwidth_controller::BandwidthController;
#[cfg(target_os = "android")]
use nym_client_core::client::base_client::helpers::setup_empty_reply_surb_backend;
#[cfg(not(target_os = "android"))]
use nym_client_core::client::base_client::non_wasm_helpers;
use nym_client_core::client::base_client::{
BaseClientBuilder, ClientInput, ClientOutput, ClientState,
};
@@ -28,6 +23,12 @@ use nym_validator_client::nyxd::QueryNyxdClient;
use nym_validator_client::Client;
use std::error::Error;
#[cfg(target_os = "android")]
use nym_client_core::client::base_client::helpers::setup_empty_reply_surb_backend;
#[cfg(not(target_os = "android"))]
use nym_client_core::client::base_client::non_wasm_helpers;
use nym_client_core::config::DebugConfig;
pub mod config;
pub mod error;
pub mod socks;
@@ -102,6 +103,7 @@ impl NymClient {
pub fn start_socks5_listener(
socks5_config: &Socks5,
debug_config: DebugConfig,
client_input: ClientInput,
client_output: ClientOutput,
client_status: ClientState,
@@ -126,6 +128,11 @@ impl NymClient {
..
} = client_status;
let packet_size = debug_config
.traffic
.secondary_packet_size
.unwrap_or(debug_config.traffic.primary_packet_size);
let authenticator = Authenticator::new(auth_methods, allowed_users);
let mut sphinx_socks = SphinxSocksServer::new(
socks5_config.get_listening_port(),
@@ -134,6 +141,7 @@ impl NymClient {
self_address,
shared_lane_queue_lengths,
socks::client::Config::new(
packet_size,
socks5_config.get_provider_interface_version(),
socks5_config.get_socks5_protocol_version(),
socks5_config.get_send_anonymously(),
@@ -255,6 +263,7 @@ impl NymClient {
Self::start_socks5_listener(
self.config.get_socks5(),
*self.config.get_debug_settings(),
client_input,
client_output,
client_state,
@@ -17,6 +17,7 @@ use nym_socks5_requests::{
ConnectionId, RemoteAddress, Socks5ProtocolVersion, Socks5ProviderRequest, Socks5Request,
};
use nym_sphinx::addressing::clients::Recipient;
use nym_sphinx::params::PacketSize;
use nym_task::connections::{LaneQueueLengths, TransmissionLane};
use nym_task::TaskClient;
use pin_project::pin_project;
@@ -131,6 +132,7 @@ impl AsyncWrite for StreamState {
#[derive(Debug, Copy, Clone)]
pub(crate) struct Config {
biggest_packet_size: PacketSize,
provider_interface_version: ProviderInterfaceVersion,
socks5_protocol_version: Socks5ProtocolVersion,
use_surbs_for_responses: bool,
@@ -140,6 +142,7 @@ pub(crate) struct Config {
impl Config {
pub(crate) fn new(
biggest_packet_size: PacketSize,
provider_interface_version: ProviderInterfaceVersion,
socks5_protocol_version: Socks5ProtocolVersion,
use_surbs_for_responses: bool,
@@ -147,6 +150,7 @@ impl Config {
per_request_surbs: u32,
) -> Self {
Self {
biggest_packet_size,
provider_interface_version,
socks5_protocol_version,
use_surbs_for_responses,
@@ -410,6 +414,9 @@ impl SocksClient {
remote_proxy_target,
conn_receiver,
input_sender,
// FIXME: this does NOT include overhead due to acks or chunking
// (so actual true plaintext is smaller)
self.config.biggest_packet_size.plaintext_size(),
connection_id,
Some(self.lane_queue_lengths.clone()),
self.shutdown_listener.clone(),
@@ -1,208 +1,36 @@
// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
// Copyright 2021-2023 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
use bytes::{BufMut, Bytes, BytesMut};
use bytes::Bytes;
use futures::Stream;
use std::cell::RefCell;
use std::future::Future;
use std::io;
use std::ops::DerefMut;
use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::AsyncRead;
use tokio::time::{sleep, Duration, Instant, Sleep};
use tokio_util::io::poll_read_buf;
const MAX_READ_AMOUNT: usize = 500 * 1000; // 0.5MB
const GRACE_DURATION: Duration = Duration::from_millis(1);
// note, min_capacity doesn't mean we're going to always read at least this amount of data,
// it defines the smallest allowed (by yours truly) upper bound
const MIN_CAPACITY: usize = 16 * 1024;
const DEFAULT_CAPACITY: usize = 64 * 1024;
pub struct AvailableReader<'a, R: AsyncRead + Unpin> {
// TODO: come up with a way to avoid using RefCell (not sure if possible though due to having to
// mutably borrow both inner reader and buffer at the same time)
buf: RefCell<BytesMut>,
inner: RefCell<&'a mut R>,
grace_period: Option<Pin<Box<Sleep>>>,
pub struct AvailableReader<R> {
inner: tokio_util::io::ReaderStream<R>,
}
impl<'a, R> AvailableReader<'a, R>
where
R: AsyncRead + Unpin,
{
const BUF_INCREMENT: usize = 4096;
impl<R: AsyncRead> AvailableReader<R> {
pub fn new(reader: R, capacity: Option<usize>) -> Self {
let capacity = capacity.unwrap_or(DEFAULT_CAPACITY).max(MIN_CAPACITY);
pub fn new(reader: &'a mut R) -> Self {
AvailableReader {
buf: RefCell::new(BytesMut::with_capacity(Self::BUF_INCREMENT)),
inner: RefCell::new(reader),
grace_period: Some(Box::pin(sleep(GRACE_DURATION))),
inner: tokio_util::io::ReaderStream::with_capacity(reader, capacity),
}
}
}
impl<'a, R: AsyncRead + Unpin> Stream for AvailableReader<'a, R> {
impl<R: AsyncRead + Unpin> Stream for AvailableReader<R> {
type Item = io::Result<Bytes>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
// if we have no space in buffer left - expand it
if !self.buf.borrow().has_remaining_mut() {
self.buf.borrow_mut().reserve(Self::BUF_INCREMENT);
}
// note: poll_read_buf calls `buf.advance_mut(n)`
let poll_res = poll_read_buf(
Pin::new(self.inner.borrow_mut().deref_mut()),
cx,
self.buf.borrow_mut().deref_mut(),
);
match poll_res {
Poll::Pending => {
// there's nothing for us here, just return whatever we have (assuming we read anything!)
if self.buf.borrow().is_empty() {
Poll::Pending
} else {
// if exists - check grace period
if let Some(grace_period) = self.grace_period.as_mut() {
if Pin::new(grace_period).poll(cx).is_pending() {
return Poll::Pending;
}
}
let buf = self.buf.replace(BytesMut::new());
Poll::Ready(Some(Ok(buf.freeze())))
}
}
Poll::Ready(Err(err)) => Poll::Ready(Some(Err(err))),
Poll::Ready(Ok(n)) => {
// if exists - reset grace period
if let Some(grace_period) = self.grace_period.as_mut() {
let now = Instant::now();
grace_period.as_mut().reset(now + GRACE_DURATION);
}
// if we read a non-0 amount, we're not done yet!
if n == 0 {
let buf = self.buf.replace(BytesMut::new());
if !buf.is_empty() {
Poll::Ready(Some(Ok(buf.freeze())))
} else {
Poll::Ready(None)
}
} else {
// tell the waker we should be polled again!
cx.waker().wake_by_ref();
// if we reached our maximum amount - return it
let read_bytes_len = self.buf.borrow().len();
if read_bytes_len >= MAX_READ_AMOUNT {
let buf = self.buf.replace(BytesMut::new());
return Poll::Ready(Some(Ok(buf.freeze())));
}
Poll::Pending
}
}
}
Pin::new(&mut self.inner).poll_next(cx)
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures::{poll, StreamExt};
use std::io::Cursor;
use std::time::Duration;
use tokio::io::AsyncReadExt;
use tokio_test::assert_pending;
#[tokio::test]
async fn available_reader_reads_all_available_data_smaller_than_its_buf() {
let data = vec![42u8; 100];
let mut reader = Cursor::new(data.clone());
let mut available_reader = AvailableReader::new(&mut reader);
let read_data = available_reader.next().await.unwrap().unwrap();
assert_eq!(read_data, data);
assert!(available_reader.next().await.is_none());
}
#[tokio::test]
async fn available_reader_reads_all_available_data_bigger_than_its_buf() {
let data = vec![42u8; AvailableReader::<Cursor<Vec<u8>>>::BUF_INCREMENT + 100];
let mut reader = Cursor::new(data.clone());
let mut available_reader = AvailableReader::new(&mut reader);
let read_data = available_reader.next().await.unwrap().unwrap();
assert_eq!(read_data, data);
assert!(available_reader.next().await.is_none());
}
#[tokio::test]
async fn available_reader_will_not_wait_for_more_data_if_it_already_has_some() {
let first_data_chunk = vec![42u8; 100];
let second_data_chunk = vec![123u8; 100];
let mut reader_mock = tokio_test::io::Builder::new()
.read(&first_data_chunk)
.wait(Duration::from_millis(100)) // delay is irrelevant, what matters is that we don't get everything immediately
.read(&second_data_chunk)
.build();
let mut available_reader = AvailableReader::new(&mut reader_mock);
let read_data = available_reader.next().await.unwrap().unwrap();
assert_eq!(read_data, first_data_chunk);
assert_pending!(poll!(available_reader.next()));
// before dropping the mock, we need to empty it
let mut buf = vec![0u8; second_data_chunk.len()];
assert_eq!(reader_mock.read(&mut buf).await.unwrap(), 100);
}
#[tokio::test]
async fn available_reader_will_wait_for_more_data_if_it_doesnt_have_anything() {
let data = vec![42u8; 100];
let mut reader_mock = tokio_test::io::Builder::new()
.wait(Duration::from_millis(100))
.read(&data)
.build();
let mut available_reader = AvailableReader::new(&mut reader_mock);
let read_data = available_reader.next().await.unwrap().unwrap();
assert_eq!(read_data, data);
assert!(available_reader.next().await.is_none());
}
// perhaps the issue of tokio io builder will be resolved in tokio 0.3?
// #[tokio::test]
// async fn available_reader_will_wait_for_more_data_if_its_within_grace_period() {
// let first_data_chunk = vec![42u8; 100];
// let second_data_chunk = vec![123u8; 100];
//
// let combined_chunks: Vec<_> = first_data_chunk
// .iter()
// .cloned()
// .chain(second_data_chunk.iter().cloned())
// .collect();
//
// let mut reader_mock = tokio_test::io::Builder::new()
// .read(&first_data_chunk)
// .wait(Duration::from_millis(2))
// .read(&second_data_chunk)
// .build();
//
// let mut available_reader = AvailableReader {
// buf: RefCell::new(BytesMut::with_capacity(4096)),
// inner: RefCell::new(&mut reader_mock),
// grace_period: Some(delay_for(Duration::from_millis(5))),
// };
//
// let read_data = available_reader.next().await.unwrap().unwrap();
//
// assert_eq!(read_data, combined_chunks);
// assert!(available_reader.next().await.is_none())
// }
}
@@ -167,6 +167,7 @@ pub(super) async fn run_inbound<F, S>(
remote_source_address: String,
connection_id: ConnectionId,
mix_sender: MixProxySender<S>,
available_plaintext_per_mix_packet: usize,
adapter_fn: F,
shutdown_notify: Arc<Notify>,
lane_queue_lengths: Option<LaneQueueLengths>,
@@ -176,7 +177,9 @@ where
F: Fn(ConnectionId, Vec<u8>, bool) -> S + Send + 'static,
S: Debug,
{
let mut available_reader = AvailableReader::new(&mut reader);
// TODO: this multiplication by 4 is completely arbitrary here
let mut available_reader =
AvailableReader::new(&mut reader, Some(available_plaintext_per_mix_packet * 4));
let mut message_sender = OrderedMessageSender::new();
let shutdown_future = shutdown_notify.notified().then(|_| sleep(SHUTDOWN_TIMEOUT));
@@ -1,4 +1,4 @@
// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
// Copyright 2021-2023 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0
use crate::connection_controller::ConnectionReceiver;
@@ -49,6 +49,8 @@ pub struct ProxyRunner<S> {
connection_id: ConnectionId,
lane_queue_lengths: Option<LaneQueueLengths>,
available_plaintext_per_mix_packet: usize,
// Listens to shutdown commands from higher up
shutdown_listener: TaskClient,
}
@@ -64,6 +66,7 @@ where
remote_source_address: String,
mix_receiver: ConnectionReceiver,
mix_sender: MixProxySender<S>,
available_plaintext_per_mix_packet: usize,
connection_id: ConnectionId,
lane_queue_lengths: Option<LaneQueueLengths>,
shutdown_listener: TaskClient,
@@ -76,6 +79,7 @@ where
remote_source_address,
connection_id,
lane_queue_lengths,
available_plaintext_per_mix_packet,
shutdown_listener,
}
}
@@ -96,6 +100,7 @@ where
self.remote_source_address.clone(),
self.connection_id,
self.mix_sender.clone(),
self.available_plaintext_per_mix_packet,
adapter_fn,
Arc::clone(&shutdown_notify),
self.lane_queue_lengths.clone(),
+2
View File
@@ -506,6 +506,7 @@ where
.socks5_config
.clone()
.ok_or(Error::Socks5Config { set: false })?;
let debug_config = self.config.debug_config;
let (mut started_client, nym_address) = self.connect_to_mixnet_common().await?;
let (socks5_status_tx, mut socks5_status_rx) = mpsc::channel(128);
@@ -515,6 +516,7 @@ where
nym_socks5_client_core::NymClient::start_socks5_listener(
&socks5_config,
debug_config,
client_input,
client_output,
client_state.clone(),
@@ -29,6 +29,7 @@ use nym_socks5_requests::{
};
use nym_sphinx::addressing::clients::Recipient;
use nym_sphinx::anonymous_replies::requests::AnonymousSenderTag;
use nym_sphinx::params::PacketSize;
use nym_statistics_common::collector::StatisticsSender;
use nym_task::connections::LaneQueueLengths;
use nym_task::{TaskClient, TaskManager};
@@ -57,6 +58,8 @@ pub struct NRServiceProviderBuilder {
}
struct NRServiceProvider {
config: Config,
outbound_request_filter: OutboundRequestFilter,
open_proxy: bool,
mixnet_client: nym_sdk::mixnet::MixnetClient,
@@ -242,6 +245,7 @@ impl NRServiceProviderBuilder {
start_allowed_list_reloader(self.allowed_hosts, shutdown.subscribe()).await;
let service_provider = NRServiceProvider {
config: self.config,
outbound_request_filter: self.outbound_request_filter,
open_proxy: self.open_proxy,
mixnet_client,
@@ -329,6 +333,7 @@ impl NRServiceProvider {
connection_id: ConnectionId,
remote_addr: String,
return_address: reply::MixnetAddress,
biggest_packet_size: PacketSize,
controller_sender: ControllerSender,
mix_input_sender: MixProxySender<MixnetMessage>,
lane_queue_lengths: LaneQueueLengths,
@@ -385,6 +390,7 @@ impl NRServiceProvider {
// run the proxy on the connection
conn.run_proxy(
remote_version,
biggest_packet_size,
mix_receiver,
mix_input_sender,
lane_queue_lengths,
@@ -437,6 +443,11 @@ impl NRServiceProvider {
return;
}
let traffic_config = self.config.get_base().get_debug_config().traffic;
let packet_size = traffic_config
.secondary_packet_size
.unwrap_or(traffic_config.primary_packet_size);
let controller_sender_clone = self.controller_sender.clone();
let mix_input_sender_clone = self.mix_input_sender.clone();
let lane_queue_lengths_clone = self.mixnet_client.shared_lane_queue_lengths();
@@ -449,6 +460,7 @@ impl NRServiceProvider {
conn_id,
remote_addr,
return_address,
packet_size,
controller_sender_clone,
mix_input_sender_clone,
lane_queue_lengths_clone,
@@ -7,6 +7,7 @@ use nym_service_providers_common::interface::RequestVersion;
use nym_socks5_proxy_helpers::connection_controller::ConnectionReceiver;
use nym_socks5_proxy_helpers::proxy_runner::{MixProxySender, ProxyRunner};
use nym_socks5_requests::{ConnectionId, RemoteAddress, Socks5Request};
use nym_sphinx::params::PacketSize;
use nym_task::connections::LaneQueueLengths;
use nym_task::TaskClient;
use std::io;
@@ -42,6 +43,7 @@ impl Connection {
pub(crate) async fn run_proxy(
&mut self,
remote_version: RequestVersion<Socks5Request>,
biggest_packet_size: PacketSize,
mix_receiver: ConnectionReceiver,
mix_sender: MixProxySender<MixnetMessage>,
lane_queue_lengths: LaneQueueLengths,
@@ -57,6 +59,9 @@ impl Connection {
remote_source_address,
mix_receiver,
mix_sender,
// FIXME: this does NOT include overhead due to acks or chunking
// (so actual true plaintext is smaller)
biggest_packet_size.plaintext_size(),
connection_id,
Some(lane_queue_lengths),
shutdown,