diff --git a/common/socks5/proxy-helpers/src/available_reader.rs b/common/socks5/proxy-helpers/src/available_reader.rs index 70d968f968..018486e68b 100644 --- a/common/socks5/proxy-helpers/src/available_reader.rs +++ b/common/socks5/proxy-helpers/src/available_reader.rs @@ -1,8 +1,9 @@ -// Copyright 2021 - Nym Technologies SA +// Copyright 2021-2023 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 use bytes::{BufMut, Bytes, BytesMut}; use futures::Stream; +use log::error; use std::cell::RefCell; use std::future::Future; use std::io; @@ -15,6 +16,7 @@ use tokio_util::io::poll_read_buf; const MAX_READ_AMOUNT: usize = 500 * 1000; // 0.5MB const GRACE_DURATION: Duration = Duration::from_millis(1); +const READ_TIMEOUT: Duration = Duration::from_millis(10); 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 @@ -22,6 +24,7 @@ pub struct AvailableReader<'a, R: AsyncRead + Unpin> { buf: RefCell, inner: RefCell<&'a mut R>, grace_period: Option>>, + read_deadline: Option>>, } impl<'a, R> AvailableReader<'a, R> @@ -31,10 +34,58 @@ where const BUF_INCREMENT: usize = 4096; pub fn new(reader: &'a mut R) -> Self { + // pub fn new(reader: &'a mut R, buffer_size: usize) -> Self { AvailableReader { buf: RefCell::new(BytesMut::with_capacity(Self::BUF_INCREMENT)), inner: RefCell::new(reader), - grace_period: Some(Box::pin(sleep(GRACE_DURATION))), + grace_period: None, + read_deadline: None, + } + } + + fn return_buf(mut self: Pin<&mut Self>) -> Poll::Item>> { + // reset timeouts so the poll wouldn't be accidentally called + self.grace_period = None; + self.read_deadline = None; + + let buf = self.buf.borrow_mut().split(); + // there's no data in the buffer, it means the underlying source is done + if buf.is_empty() { + Poll::Ready(None) + } else { + Poll::Ready(Some(Ok(buf.freeze()))) + } + } + + fn poll_read_deadline(&mut self, cx: &mut Context<'_>) -> Poll<()> { + match self.read_deadline.as_mut() { + Some(read_deadline) => Pin::new(read_deadline).poll(cx), + None => { + // this is the first time we're polling this reader - set the deadline + self.read_deadline = Some(Box::pin(sleep(READ_TIMEOUT))); + Poll::Pending + } + } + } + + fn poll_grace_period(&mut self, cx: &mut Context<'_>) -> Poll<()> { + match self.grace_period.as_mut() { + Some(grace_period) => Pin::new(grace_period).poll(cx), + None => { + // this is the first time we're polling this reader - set the grace period + self.grace_period = Some(Box::pin(sleep(GRACE_DURATION))); + Poll::Pending + } + } + } + + fn reset_grace_period(&mut self) { + if let Some(grace_period) = self.grace_period.as_mut() { + let now = Instant::now(); + grace_period.as_mut().reset(now + GRACE_DURATION); + } else { + // this branch should be impossible! + error!("reached an impossible branch - attempted to reset a non-existent grace period timeout") } } } @@ -55,50 +106,46 @@ impl<'a, R: AsyncRead + Unpin> Stream for AvailableReader<'a, R> { self.buf.borrow_mut().deref_mut(), ); + let deadline_poll_res = self.poll_read_deadline(cx); + let grace_period_poll_res = self.poll_grace_period(cx); + match poll_res { Poll::Pending => { - // there's nothing for us here, just return whatever we have (assuming we read anything!) + // there's nothing for us here, just return whatever we have + // (assuming we read anything and we SHOULD HAVE read something) if self.buf.borrow().is_empty() { + // this case shouldn't be possible - if something woke the future we MUST HAVE HAD some data + // (it's because we never set the sleeps until the first call, meaning some data was already there) + error!("AvailableReader got woken with not data to return!"); + Poll::Pending + } else if grace_period_poll_res.is_pending() && deadline_poll_res.is_pending() { + // if neither timeout was reached yet, wait a bit more 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()))) + // otherwise return what we managed to read (and reset timeouts) + self.return_buf() } } 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); - } + self.reset_grace_period(); - // 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) - } + // if we read 0 bytes, it means the underlying source is done, so return what we have + // (note, this might get called twice if we already had something in the buffer) + self.return_buf() } else { // tell the waker we should be polled again! cx.waker().wake_by_ref(); - // if we reached our maximum amount - return it + // if we reached our maximum amount or we've been trying to read the data for too long + // return what we have 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()))); + if read_bytes_len >= MAX_READ_AMOUNT || deadline_poll_res.is_ready() { + self.return_buf() + } else { + Poll::Pending } - Poll::Pending } } } @@ -154,7 +201,6 @@ mod tests { 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);