Flush buffered adaptive relay writes when the source stalls

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::<u64, io::Error>(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::<u64, io::Error>(copied)
             };

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<R, W>(reader: &mut R, writer: &mut W) -> io::Result<u64>
+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
+}

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<u8>,
+    pos: usize,
+    eof: bool,
+}
+
+impl AsyncRead for BurstReader {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+        output: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        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<u8>,
+    allowance: usize,
+    max_write: usize,
+    waker: Option<Waker>,
+    flushes: usize,
+}
+
+impl SinkState {
+    fn new(allowance: usize, max_write: usize) -> Arc<Mutex<Self>> {
+        Arc::new(Mutex::new(Self {
+            bytes: Vec::new(),
+            allowance,
+            max_write,
+            waker: None,
+            flushes: 0,
+        }))
+    }
+}
+
+struct Sink(Arc<Mutex<SinkState>>);
+
+impl AsyncWrite for Sink {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bytes: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        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<io::Result<()>> {
+        self.0.lock().unwrap().flushes += 1;
+        Poll::Ready(Ok(()))
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Poll::Ready(Ok(()))
+    }
+}
+
+#[derive(Default)]
+struct WakeCount(AtomicUsize);
+
+impl Wake for WakeCount {
+    fn wake(self: Arc<Self>) {
+        self.0.fetch_add(1, Ordering::Relaxed);
+    }
+}
+
+fn plaintext(records: &[u8]) -> Vec<u8> {
+    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);
+    }
+}

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<AtomicUsize>,
+    blocked: Arc<AtomicBool>,
+    error: Option<io::ErrorKind>,
+}
+
+impl AsyncWrite for FlushGate {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+        bytes: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        self.accepted += bytes.len();
+        Poll::Ready(Ok(bytes.len()))
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        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<io::Result<()>> {
+        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<io::Result<()>> {
+        Poll::Ready(Err(io::ErrorKind::UnexpectedEof.into()))
+    }
+}
+
+impl AsyncWrite for ErrorIo {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+        _buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        panic!("copy must not flush after a read/write error")
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        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);
+}