From 77ae1108dca24039588ae13db320852d95751c5e Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Fri, 9 Oct 2026 19:59:03 +0300 Subject: [PATCH] Flush buffered adaptive relay writes when the source stalls --- src/proxy/relay/adaptive_copy.rs | 13 +- src/proxy/relay/adaptive_copy/copy.rs | 45 ++++ src/proxy/relay/adaptive_copy/tests.rs | 235 +++++++++++++++++ src/proxy/relay/adaptive_copy/tests/flow.rs | 275 ++++++++++++++++++++ 4 files changed, 565 insertions(+), 3 deletions(-) create mode 100644 src/proxy/relay/adaptive_copy/copy.rs create mode 100644 src/proxy/relay/adaptive_copy/tests.rs create mode 100644 src/proxy/relay/adaptive_copy/tests/flow.rs diff --git a/src/proxy/relay/adaptive_copy.rs b/src/proxy/relay/adaptive_copy.rs index ae6ca2a..72c6785 100644 --- a/src/proxy/relay/adaptive_copy.rs +++ b/src/proxy/relay/adaptive_copy.rs @@ -5,7 +5,7 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::task::{Context, Poll}; use std::time::Duration; -use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf, copy_buf}; +use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf}; use tokio::time::Instant; use tokio_util::sync::CancellationToken; use tracing::{debug, warn}; @@ -25,8 +25,15 @@ use super::WATCHDOG_INTERVAL; use super::io::{SharedCounters, StatsIo, is_quota_io_error}; use super::watchdog_delta; +// Buffered copy progress and downstream pressure sampling remain local to this relay. +mod copy; mod write_pressure; +// Deterministic transport and relay lifecycle regressions. +#[cfg(test)] +mod tests; + +use self::copy::copy_with_flush; use self::write_pressure::WritePressureIo; struct AdaptiveBufferState { @@ -227,12 +234,12 @@ where let relay_outcome = { let copy = async { let c2s = async { - let copied = copy_buf(&mut c2s_reader, &mut server_writer).await?; + let copied = copy_with_flush(&mut c2s_reader, &mut server_writer).await?; server_writer.shutdown().await?; Ok::(copied) }; let s2c = async { - let copied = copy_buf(&mut s2c_reader, &mut client_writer).await?; + let copied = copy_with_flush(&mut s2c_reader, &mut client_writer).await?; client_writer.shutdown().await?; Ok::(copied) }; diff --git a/src/proxy/relay/adaptive_copy/copy.rs b/src/proxy/relay/adaptive_copy/copy.rs new file mode 100644 index 0000000..ead0afe --- /dev/null +++ b/src/proxy/relay/adaptive_copy/copy.rs @@ -0,0 +1,45 @@ +use std::future::poll_fn; +use std::io; +use std::pin::Pin; +use std::task::{Poll, ready}; + +use tokio::io::{AsyncBufRead, AsyncWrite}; + +/// Copies from the existing reader buffer, flushing accepted bytes when reads stall. +pub(super) async fn copy_with_flush(reader: &mut R, writer: &mut W) -> io::Result +where + R: AsyncBufRead + Unpin + ?Sized, + W: AsyncWrite + Unpin + ?Sized, +{ + let mut copied = 0; + let mut need_flush = false; + poll_fn(|cx| { + loop { + let buffer = match Pin::new(&mut *reader).poll_fill_buf(cx) { + Poll::Ready(result) => result?, + Poll::Pending => { + // A writer can accept plaintext while retaining ciphertext. Its + // progress must not depend on another source read or eventual EOF. + if need_flush { + ready!(Pin::new(&mut *writer).poll_flush(cx))?; + need_flush = false; + } + return Poll::Pending; + } + }; + if buffer.is_empty() { + ready!(Pin::new(&mut *writer).poll_flush(cx))?; + return Poll::Ready(Ok(copied)); + } + + let written = ready!(Pin::new(&mut *writer).poll_write(cx, buffer))?; + if written == 0 { + return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); + } + Pin::new(&mut *reader).consume(written); + copied += written as u64; + need_flush = true; + } + }) + .await +} diff --git a/src/proxy/relay/adaptive_copy/tests.rs b/src/proxy/relay/adaptive_copy/tests.rs new file mode 100644 index 0000000..00cd0a2 --- /dev/null +++ b/src/proxy/relay/adaptive_copy/tests.rs @@ -0,0 +1,235 @@ +use super::*; +use crate::crypto::aes::AesCtr; +use crate::protocol::constants::MAX_TLS_CIPHERTEXT_SIZE; +use crate::stream::{CryptoWriter, FakeTlsWriter}; +use std::future::Future; +use std::sync::Mutex; +use std::task::{Wake, Waker}; + +// Copy-loop error, backpressure, and wakeup invariants share these transport fixtures. +mod flow; + +struct BurstReader { + bytes: Vec, + pos: usize, + eof: bool, +} + +impl AsyncRead for BurstReader { + fn poll_read( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + output: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + if this.pos == this.bytes.len() && !this.eof { + return Poll::Pending; + } + let n = output.remaining().min(this.bytes.len() - this.pos); + output.put_slice(&this.bytes[this.pos..this.pos + n]); + this.pos += n; + Poll::Ready(Ok(())) + } +} + +struct SinkState { + bytes: Vec, + allowance: usize, + max_write: usize, + waker: Option, + flushes: usize, +} + +impl SinkState { + fn new(allowance: usize, max_write: usize) -> Arc> { + Arc::new(Mutex::new(Self { + bytes: Vec::new(), + allowance, + max_write, + waker: None, + flushes: 0, + })) + } +} + +struct Sink(Arc>); + +impl AsyncWrite for Sink { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bytes: &[u8], + ) -> Poll> { + let mut state = self.0.lock().unwrap(); + if state.allowance == 0 { + state.waker = Some(cx.waker().clone()); + return Poll::Pending; + } + let n = bytes.len().min(state.allowance).min(state.max_write); + state.bytes.extend_from_slice(&bytes[..n]); + state.allowance -= n; + Poll::Ready(Ok(n)) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + self.0.lock().unwrap().flushes += 1; + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} + +#[derive(Default)] +struct WakeCount(AtomicUsize); + +impl Wake for WakeCount { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +fn plaintext(records: &[u8]) -> Vec { + let mut ciphertext = Vec::new(); + let mut offset = 0; + while offset < records.len() { + assert!(records.len() - offset >= 5, "incomplete TLS header"); + assert_eq!(&records[offset..offset + 3], &[0x17, 3, 3]); + let n = u16::from_be_bytes([records[offset + 3], records[offset + 4]]) as usize; + assert!(n <= MAX_TLS_CIPHERTEXT_SIZE); + assert!(records.len() - offset - 5 >= n, "incomplete TLS payload"); + ciphertext.extend_from_slice(&records[offset + 5..offset + 5 + n]); + offset += 5 + n; + } + AesCtr::new(&[7; 32], 9).apply(&mut ciphertext); + ciphertext +} + +#[tokio::test(start_paused = true)] +async fn adaptive_relay_flushes_faketls_tail_without_another_source_event() { + let payload: Vec<_> = (0..65536).map(|i| (i % 251) as u8).collect(); + let sink = SinkState::new(usize::MAX, usize::MAX); + let writer = CryptoWriter::new( + FakeTlsWriter::new(Sink(Arc::clone(&sink))), + AesCtr::new(&[7; 32], 9), + 65536, + ); + let stats = Arc::new(Stats::new()); + let user = "adaptive-tail"; + let cancel = CancellationToken::new(); + let mut relay = Box::pin(relay_direct_adaptive( + tokio::io::empty(), + writer, + BurstReader { + bytes: payload.clone(), + pos: 0, + eof: false, + }, + tokio::io::sink(), + 4096, + 65536, + 1, + user, + Arc::clone(&stats), + stats.current_user_quota_handle(user), + Some(1024 * 1024), + None, + Duration::from_secs(60), + cancel.clone(), + DirectBufferBudget::new(1024 * 1024), + )); + let mut cx = Context::from_waker(Waker::noop()); + assert!(relay.as_mut().poll(&mut cx).is_pending()); + tokio::time::advance(Duration::from_millis(1)).await; + for _ in 0..3 { + assert!(relay.as_mut().poll(&mut cx).is_pending()); + } + assert_eq!(stats.get_user_total_octets(user), payload.len() as u64); + let delivered = plaintext(&sink.lock().unwrap().bytes); + assert_eq!( + delivered.len(), + payload.len(), + "accepted ciphertext must not await another read" + ); + assert_eq!(delivered, payload); + assert_eq!(stats.get_user_total_octets(user), payload.len() as u64); + assert_eq!(stats.get_user_quota_used(user), payload.len() as u64); + assert_eq!( + sink.lock().unwrap().flushes, + 1, + "idle polls must not repeat a completed flush" + ); + cancel.cancel(); + assert!(matches!( + relay.as_mut().poll(&mut cx), + Poll::Ready(Err(ProxyError::UserDisabled { .. })) + )); +} + +#[tokio::test(start_paused = true)] +async fn adaptive_relay_keeps_quota_and_blocked_cancellation_limits() { + for blocked in [false, true] { + let payload = vec![29; 65536]; + let sink = SinkState::new(if blocked { 0 } else { usize::MAX }, usize::MAX); + let writer = CryptoWriter::new( + FakeTlsWriter::new(Sink(Arc::clone(&sink))), + AesCtr::new(&[7; 32], 9), + 65536, + ); + let stats = Arc::new(Stats::new()); + let user = "adaptive-limits"; + let cancel = CancellationToken::new(); + let budget = DirectBufferBudget::new(1024 * 1024); + let limit = if blocked { 1024 * 1024 } else { 32768 }; + let mut relay = Box::pin(relay_direct_adaptive( + tokio::io::empty(), + writer, + BurstReader { + bytes: payload.clone(), + pos: 0, + eof: false, + }, + tokio::io::sink(), + 4096, + 65536, + 1, + user, + Arc::clone(&stats), + stats.current_user_quota_handle(user), + Some(limit), + None, + Duration::from_secs(60), + cancel.clone(), + Arc::clone(&budget), + )); + let mut cx = Context::from_waker(Waker::noop()); + let mut result = relay.as_mut().poll(&mut cx); + if result.is_pending() { + tokio::time::advance(Duration::from_millis(1)).await; + result = relay.as_mut().poll(&mut cx); + } + if blocked { + assert!(result.is_pending()); + assert_eq!(stats.get_user_total_octets(user), payload.len() as u64); + assert!(sink.lock().unwrap().bytes.is_empty()); + cancel.cancel(); + assert!(matches!( + relay.as_mut().poll(&mut cx), + Poll::Ready(Err(ProxyError::UserDisabled { .. })) + )); + } else { + assert!(matches!( + result, + Poll::Ready(Err(ProxyError::DataQuotaExceeded { .. })) + )); + assert_eq!(stats.get_user_quota_used(user), limit); + assert_eq!(stats.get_user_total_octets(user), limit); + assert_eq!( + plaintext(&sink.lock().unwrap().bytes), + payload[..limit as usize] + ); + } + assert_eq!(budget.snapshot().reserved_bytes, 0); + } +} diff --git a/src/proxy/relay/adaptive_copy/tests/flow.rs b/src/proxy/relay/adaptive_copy/tests/flow.rs new file mode 100644 index 0000000..a85c378 --- /dev/null +++ b/src/proxy/relay/adaptive_copy/tests/flow.rs @@ -0,0 +1,275 @@ +use super::*; + +#[test] +fn socket_wakeup_drains_both_encrypted_layers_while_source_stays_pending() { + let payload: Vec<_> = (0..65536).map(|i| (i % 251) as u8).collect(); + let mut reader = AdaptiveBufReader::new( + BurstReader { + bytes: payload.clone(), + pos: 0, + eof: false, + }, + AdaptiveBufferState::new(payload.len()), + ); + let sink = SinkState::new(13, 7); + let mut writer = CryptoWriter::new( + FakeTlsWriter::new(Sink(Arc::clone(&sink))), + AesCtr::new(&[7; 32], 9), + 65536, + ); + let wake = Arc::new(WakeCount::default()); + let waker = Waker::from(Arc::clone(&wake)); + let mut cx = Context::from_waker(&waker); + let mut copy = Box::pin(copy_with_flush(&mut reader, &mut writer)); + for _ in 0..3 { + assert!(copy.as_mut().poll(&mut cx).is_pending()); + assert_eq!(sink.lock().unwrap().bytes.len(), 13); + } + assert_eq!( + wake.0.load(Ordering::Relaxed), + 0, + "blocked copy must not spin or self-wake" + ); + let wake_writer = { + let mut state = sink.lock().unwrap(); + state.allowance = usize::MAX; + state + .waker + .take() + .expect("blocked sink must register readiness") + }; + wake_writer.wake(); + assert_eq!(wake.0.load(Ordering::Relaxed), 1); + assert!(copy.as_mut().poll(&mut cx).is_pending()); + assert_eq!(plaintext(&sink.lock().unwrap().bytes), payload); + drop(copy); + assert!(!writer.has_pending()); + assert!(!writer.get_ref().has_pending()); +} + +#[test] +fn writer_pending_preserves_unconsumed_reader_bytes_and_reservation() { + let state = AdaptiveBufferState::new(64); + let payload = vec![42; 64]; + let mut reader = AdaptiveBufReader::new( + BurstReader { + bytes: payload.clone(), + pos: 0, + eof: true, + }, + Arc::clone(&state), + ); + let sink = SinkState::new(0, 3); + let mut writer = Sink(Arc::clone(&sink)); + let mut cx = Context::from_waker(Waker::noop()); + let mut copy = Box::pin(copy_with_flush(&mut reader, &mut writer)); + assert!(copy.as_mut().poll(&mut cx).is_pending()); + drop(copy); + assert_eq!(reader.pos, 0); + assert_eq!(reader.cap, payload.len()); + assert!(sink.lock().unwrap().bytes.is_empty()); + state.desired_bytes.store(8, Ordering::Release); + sink.lock().unwrap().allowance = usize::MAX; + let mut copy = Box::pin(copy_with_flush(&mut reader, &mut writer)); + assert!(matches!(copy.as_mut().poll(&mut cx), Poll::Ready(Ok(64)))); + drop(copy); + assert_eq!(sink.lock().unwrap().bytes, payload); + assert_eq!(state.actual_bytes.load(Ordering::Acquire), 8); +} + +#[test] +fn eof_flushes_complete_ciphertext_and_reports_plaintext_count() { + for len in [ + 0, + 1, + MAX_TLS_CIPHERTEXT_SIZE, + MAX_TLS_CIPHERTEXT_SIZE + 1, + 65536, + ] { + let payload = vec![17; len]; + let mut reader = AdaptiveBufReader::new( + BurstReader { + bytes: payload.clone(), + pos: 0, + eof: true, + }, + AdaptiveBufferState::new(65536), + ); + let sink = SinkState::new(usize::MAX, 113); + let mut writer = CryptoWriter::new( + FakeTlsWriter::new(Sink(Arc::clone(&sink))), + AesCtr::new(&[7; 32], 9), + 65536, + ); + let mut cx = Context::from_waker(Waker::noop()); + let mut copy = Box::pin(copy_with_flush(&mut reader, &mut writer)); + assert!(matches!(copy.as_mut().poll(&mut cx), Poll::Ready(Ok(n)) if n == len as u64)); + assert_eq!(plaintext(&sink.lock().unwrap().bytes), payload); + } +} + +struct FlushGate { + accepted: usize, + polls: Arc, + blocked: Arc, + error: Option, +} + +impl AsyncWrite for FlushGate { + fn poll_write( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + bytes: &[u8], + ) -> Poll> { + self.accepted += bytes.len(); + Poll::Ready(Ok(bytes.len())) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + self.polls.fetch_add(1, Ordering::Relaxed); + if self.blocked.load(Ordering::Relaxed) { + return Poll::Pending; + } + Poll::Ready(self.error.map_or(Ok(()), |kind| Err(kind.into()))) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + panic!("copy must leave shutdown to its caller") + } +} + +#[test] +fn pending_flush_is_retried_and_successful_idle_flush_is_not_repeated() { + let mut reader = AdaptiveBufReader::new( + BurstReader { + bytes: vec![3; 128], + pos: 0, + eof: false, + }, + AdaptiveBufferState::new(256), + ); + let polls = Arc::new(AtomicUsize::new(0)); + let blocked = Arc::new(AtomicBool::new(true)); + let mut writer = FlushGate { + accepted: 0, + polls: Arc::clone(&polls), + blocked: Arc::clone(&blocked), + error: None, + }; + let mut cx = Context::from_waker(Waker::noop()); + let mut copy = Box::pin(copy_with_flush(&mut reader, &mut writer)); + assert!(copy.as_mut().poll(&mut cx).is_pending()); + assert_eq!(polls.load(Ordering::Relaxed), 1); + blocked.store(false, Ordering::Relaxed); + assert!(copy.as_mut().poll(&mut cx).is_pending()); + assert_eq!(polls.load(Ordering::Relaxed), 2); + for _ in 0..8 { + assert!(copy.as_mut().poll(&mut cx).is_pending()); + } + assert_eq!(polls.load(Ordering::Relaxed), 2); + drop(copy); + assert_eq!(writer.accepted, 128); +} + +#[test] +fn source_pending_before_any_write_does_not_flush() { + let mut reader = AdaptiveBufReader::new( + BurstReader { + bytes: Vec::new(), + pos: 0, + eof: false, + }, + AdaptiveBufferState::new(64), + ); + let sink = SinkState::new(usize::MAX, usize::MAX); + let mut writer = Sink(Arc::clone(&sink)); + let mut cx = Context::from_waker(Waker::noop()); + let mut copy = Box::pin(copy_with_flush(&mut reader, &mut writer)); + for _ in 0..8 { + assert!(copy.as_mut().poll(&mut cx).is_pending()); + } + assert_eq!(sink.lock().unwrap().flushes, 0); +} + +#[test] +fn flush_errors_propagate_on_source_pending_and_eof() { + for eof in [false, true] { + let mut reader = AdaptiveBufReader::new( + BurstReader { + bytes: vec![3; 128], + pos: 0, + eof, + }, + AdaptiveBufferState::new(256), + ); + let mut writer = FlushGate { + accepted: 0, + polls: Arc::new(AtomicUsize::new(0)), + blocked: Arc::new(AtomicBool::new(false)), + error: Some(io::ErrorKind::ConnectionReset), + }; + let mut cx = Context::from_waker(Waker::noop()); + let mut copy = Box::pin(copy_with_flush(&mut reader, &mut writer)); + assert!( + matches!(copy.as_mut().poll(&mut cx), Poll::Ready(Err(e)) if e.kind() == io::ErrorKind::ConnectionReset) + ); + } +} + +struct ErrorIo; + +impl AsyncRead for ErrorIo { + fn poll_read( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + _buf: &mut ReadBuf<'_>, + ) -> Poll> { + Poll::Ready(Err(io::ErrorKind::UnexpectedEof.into())) + } +} + +impl AsyncWrite for ErrorIo { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + _buf: &[u8], + ) -> Poll> { + Poll::Ready(Err(io::ErrorKind::BrokenPipe.into())) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + panic!("copy must not flush after a read/write error") + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + panic!("copy must leave shutdown to its caller") + } +} + +#[test] +fn read_write_and_write_zero_errors_preserve_their_categories() { + let mut cx = Context::from_waker(Waker::noop()); + let mut reader = AdaptiveBufReader::new(ErrorIo, AdaptiveBufferState::new(8)); + let mut writer = ErrorIo; + let mut copy = Box::pin(copy_with_flush(&mut reader, &mut writer)); + assert!( + matches!(copy.as_mut().poll(&mut cx), Poll::Ready(Err(e)) if e.kind() == io::ErrorKind::UnexpectedEof) + ); + drop(copy); + + let mut reader = AdaptiveBufReader::new(&b"abc"[..], AdaptiveBufferState::new(8)); + let mut copy = Box::pin(copy_with_flush(&mut reader, &mut writer)); + assert!( + matches!(copy.as_mut().poll(&mut cx), Poll::Ready(Err(e)) if e.kind() == io::ErrorKind::BrokenPipe) + ); + drop(copy); + assert_eq!(reader.pos, 0); + + let mut writer = Sink(SinkState::new(usize::MAX, 0)); + let mut copy = Box::pin(copy_with_flush(&mut reader, &mut writer)); + assert!( + matches!(copy.as_mut().poll(&mut cx), Poll::Ready(Err(e)) if e.kind() == io::ErrorKind::WriteZero) + ); + drop(copy); + assert_eq!(reader.pos, 0); +}