From dede3d7409cf3a92e1ccdf87cb1f59749bb36ed9 Mon Sep 17 00:00:00 2001 From: Remi Dettai Date: Thu, 8 Oct 2026 17:04:04 +0200 Subject: [PATCH] Bound the time spent in a single TlsStream::poll_read MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Alternative to #209: stop batching records once 100 µs have elapsed instead of after a byte count. Co-Authored-By: Claude Opus 5.5 --- src/client.rs | 6 ++++-- src/common/mod.rs | 8 ++++++++ src/server.rs | 9 +++++++-- tests/test.rs | 46 ++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 65 insertions(+), 4 deletions(-) diff --git a/src/client.rs b/src/client.rs index 2bd36592..5745fea7 100644 --- a/src/client.rs +++ b/src/client.rs @@ -9,12 +9,13 @@ use std::sync::Arc; #[cfg(feature = "early-data")] use std::task::Waker; use std::task::{Context, Poll}; +use std::time::Instant; use rustls::pki_types::ServerName; use rustls::{ClientConfig, ClientConnection}; use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; -use crate::common::{IoSession, MidHandshake, Stream, TlsState}; +use crate::common::{IoSession, MAX_READ_DURATION_PER_POLL, MidHandshake, Stream, TlsState}; /// A wrapper around a `rustls::ClientConfig`, providing an async `connect` method. #[derive(Clone)] @@ -341,6 +342,7 @@ where if let Some(err) = self.error.take() { return Poll::Ready(Err(err)); }; + let start = Instant::now(); let data = ready!(self.as_mut().poll_fill_buf(cx))?; let len = data.len().min(buf.remaining()); if len == 0 { @@ -349,7 +351,7 @@ where buf.put_slice(&data[..len]); self.as_mut().consume(len); - while buf.remaining() > 0 { + while buf.remaining() > 0 && start.elapsed() < MAX_READ_DURATION_PER_POLL { let data = match self.as_mut().poll_fill_buf(cx) { Poll::Ready(Ok([])) => break, Poll::Ready(Ok(data)) => data, diff --git a/src/common/mod.rs b/src/common/mod.rs index ede69d9b..908f1cfd 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -2,6 +2,7 @@ use std::io::{self, BufRead as _, IoSlice, Read, Write}; use std::ops::{Deref, DerefMut}; use std::pin::Pin; use std::task::{Context, Poll}; +use std::time::Duration; use rustls::{ConnectionCommon, SideData}; use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; @@ -9,6 +10,13 @@ use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; mod handshake; pub(crate) use handshake::{IoSession, MidHandshake}; +/// Time after which `TlsStream::poll_read` stops batching records. +/// +/// `poll_read` drains as many buffered records as fit in the caller's buffer to save +/// wake-ups, but every record costs a decryption. Returning once this much time has been +/// spent bounds how long a single poll holds the executor thread, whatever the buffer size. +pub(crate) const MAX_READ_DURATION_PER_POLL: Duration = Duration::from_micros(100); + #[derive(Debug)] pub(crate) enum TlsState { #[cfg(feature = "early-data")] diff --git a/src/server.rs b/src/server.rs index c24c4890..bcae9bf0 100644 --- a/src/server.rs +++ b/src/server.rs @@ -7,12 +7,16 @@ use std::os::windows::io::{AsRawSocket, RawSocket}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use std::time::Instant; use rustls::server::AcceptedAlert; use rustls::{ServerConfig, ServerConnection}; use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf}; -use crate::common::{IoSession, MidHandshake, Stream, SyncReadAdapter, SyncWriteAdapter, TlsState}; +use crate::common::{ + IoSession, MAX_READ_DURATION_PER_POLL, MidHandshake, Stream, SyncReadAdapter, SyncWriteAdapter, + TlsState, +}; /// A wrapper around a `rustls::ServerConfig`, providing an async `accept` method. #[derive(Clone)] @@ -426,6 +430,7 @@ where if let Some(err) = self.error.take() { return Poll::Ready(Err(err)); }; + let start = Instant::now(); let data = ready!(self.as_mut().poll_fill_buf(cx))?; let len = data.len().min(buf.remaining()); if len == 0 { @@ -434,7 +439,7 @@ where buf.put_slice(&data[..len]); self.as_mut().consume(len); - while buf.remaining() > 0 { + while buf.remaining() > 0 && start.elapsed() < MAX_READ_DURATION_PER_POLL { let data = match self.as_mut().poll_fill_buf(cx) { Poll::Ready(Ok([])) => break, Poll::Ready(Ok(data)) => data, diff --git a/tests/test.rs b/tests/test.rs index 1f1f3b3a..93dd8447 100644 --- a/tests/test.rs +++ b/tests/test.rs @@ -442,5 +442,51 @@ async fn propagate_connection_aborted() { server.abort(); } +/// Writes `len` bytes into `writer`, then performs a single `poll_read` on `reader` with a +/// buffer large enough to hold everything, and returns how many bytes that one call returned. +async fn single_poll_read_len( + mut writer: impl AsyncWrite + Unpin, + mut reader: impl AsyncRead + Unpin, + len: usize, +) -> usize { + writer.write_all(&vec![0x42; len]).await.unwrap(); + writer.flush().await.unwrap(); + + let mut storage = vec![0u8; len]; + let mut buf = ReadBuf::new(&mut storage); + future::poll_fn(|cx| Pin::new(&mut reader).poll_read(cx, &mut buf)) + .await + .unwrap(); + buf.filled().len() +} + +#[tokio::test] +async fn poll_read_is_bounded() { + // Decrypting 64 records takes well over the per-poll time budget, so a single + // `poll_read` must stop before draining all of them. + const PAYLOAD_LEN: usize = 1024 * 1024; + // The transport must hold the whole encrypted payload, so that every record is + // available to the first `poll_read`. + let (client_io, server_io) = tokio::io::duplex(2 * PAYLOAD_LEN); + + let (server_config, client_config) = utils::make_configs(); + let (client, server) = tokio::join!( + TlsConnector::from(Arc::new(client_config)).connect( + ServerName::try_from(utils::TEST_SERVER_DOMAIN).unwrap(), + client_io, + ), + TlsAcceptor::from(Arc::new(server_config)).accept(server_io), + ); + let (mut client, mut server) = (client.unwrap(), server.unwrap()); + + // Server writes, client reads. + let n = single_poll_read_len(&mut server, &mut client, PAYLOAD_LEN).await; + assert!(n < PAYLOAD_LEN, "expected a partial read, got {n} bytes"); + + // Client writes, server reads. + let n = single_poll_read_len(&mut client, &mut server, PAYLOAD_LEN).await; + assert!(n < PAYLOAD_LEN, "expected a partial read, got {n} bytes"); +} + // Include `utils` module include!("utils.rs");