Bound the time spent in a single TlsStream::poll_read

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,

src/common/mod.rs

@@ -2,13 +2,21 @@ 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};
 
 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")]

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,

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");