From 42fa3bc0bdc8f14c364d47709ba3da811776b57b Mon Sep 17 00:00:00 2001 From: Arni Dagur Date: Wed, 2 Sep 2026 17:47:57 +0100 Subject: [PATCH 1/2] refactor: rename `async_read_ready.rs` to `async_ready.rs` The module is about to gain an `AsyncWriteReady` trait next to `AsyncReadReady`, so give it a name that covers both. Renaming it on its own keeps the change visible to Git's rename detection. --- ktls/src/{async_read_ready.rs => async_ready.rs} | 0 ktls/src/lib.rs | 4 ++-- 2 files changed, 2 insertions(+), 2 deletions(-) rename ktls/src/{async_read_ready.rs => async_ready.rs} (100%) diff --git a/ktls/src/async_read_ready.rs b/ktls/src/async_ready.rs similarity index 100% rename from ktls/src/async_read_ready.rs rename to ktls/src/async_ready.rs diff --git a/ktls/src/lib.rs b/ktls/src/lib.rs index 7e31d86..277ea47 100644 --- a/ktls/src/lib.rs +++ b/ktls/src/lib.rs @@ -3,7 +3,7 @@ compile_error!("This crate needs wither the 'ring' or 'aws_lc_rs' feature enable #[cfg(all(feature = "ring", feature = "aws_lc_rs"))] compile_error!("The 'ring' and 'aws_lc_rs' features are mutually exclusive"); -mod async_read_ready; +mod async_ready; mod cork_stream; mod ffi; mod ktls_stream; @@ -24,7 +24,7 @@ use smallvec::SmallVec; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite}; use tokio::net::{TcpListener, TcpStream}; -pub use crate::async_read_ready::AsyncReadReady; +pub use crate::async_ready::AsyncReadReady; pub use crate::cork_stream::CorkStream; pub use crate::ffi::CryptoInfo; use crate::ffi::{setup_tls_info, setup_ulp, KtlsCompatibilityError}; From 198807bd6bbe0f072ae5388dde83772ba76921c9 Mon Sep 17 00:00:00 2001 From: Arni Dagur Date: Fri, 7 Aug 2026 22:44:59 +0100 Subject: [PATCH 2/2] fix: retry close_notify in poll_shutdown when the send buffer is full `poll_shutdown` marked the write side closed _before_ enforcing a successful `close_notify`. With a full send buffer this resulted in two bugs: 1. Shutdown failed with `WouldBlock`, which is an error type that should not escape a poll-based API. 2. Subsequent retries of the shutdown skipped the `close_notify` entirely, since `write_closed` was already set. We fix this by only marking the write side closed once the alert is sent (or has failed fatally), and retrying on `WouldBlock`. The retry uses a new `AsyncWriteReady` trait, which mirrors the preexisting `AsyncReadReady`. It exposes tokio's `poll_write_ready` and `try_io`. The latter `try_write_io` clears write-readiness when `send_close_notify` returns `WouldBlock`, so the task parks until the socket becomes writable instead of [busy-polling][1]. [1]: https://github.com/rustls/ktls/blob/5e3c7d6ceadbb1ae98d06908d559490723899aed/ktls/src/ktls_stream.rs#L268-L277 This PR introduces a minor breaking change, since the `AsyncWrite` impl for `KtlsStream` now requires `IO: AsyncWriteReady`. This change also lays the foundation for additional work, including in relation to properly implementing `KeyUpdate`. --- ktls/src/async_ready.rs | 30 ++++++++- ktls/src/cork_stream.rs | 15 ++++- ktls/src/ktls_stream.rs | 27 ++++++--- ktls/src/lib.rs | 2 +- ktls/tests/integration_test.rs | 107 ++++++++++++++++++++++++++++++++- 5 files changed, 166 insertions(+), 15 deletions(-) diff --git a/ktls/src/async_ready.rs b/ktls/src/async_ready.rs index b8d9a21..7255f41 100644 --- a/ktls/src/async_ready.rs +++ b/ktls/src/async_ready.rs @@ -1,12 +1,36 @@ -use std::{io, task}; +use std::io; +use std::task::{Context, Poll}; + +use tokio::io::Interest; pub trait AsyncReadReady { /// cf. https://docs.rs/tokio/latest/tokio/net/struct.TcpStream.html#method.poll_read_ready - fn poll_read_ready(&self, cx: &mut task::Context<'_>) -> task::Poll>; + fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll>; } impl AsyncReadReady for tokio::net::TcpStream { - fn poll_read_ready(&self, cx: &mut task::Context<'_>) -> task::Poll> { + fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll> { tokio::net::TcpStream::poll_read_ready(self, cx) } } + +pub trait AsyncWriteReady { + /// cf. https://docs.rs/tokio/latest/tokio/net/struct.TcpStream.html#method.poll_write_ready + fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll>; + + /// Perform a write to the socket using a user-provided I/O operation + /// + /// If the operation returns `WouldBlock`, the socket's write-readiness is cleared. + /// cf. https://docs.rs/tokio/latest/tokio/net/struct.TcpStream.html#method.try_io + fn try_write_io(&self, f: impl FnOnce() -> io::Result) -> io::Result; +} + +impl AsyncWriteReady for tokio::net::TcpStream { + fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll> { + tokio::net::TcpStream::poll_write_ready(self, cx) + } + + fn try_write_io(&self, f: impl FnOnce() -> io::Result) -> io::Result { + self.try_io(Interest::WRITABLE, f) + } +} diff --git a/ktls/src/cork_stream.rs b/ktls/src/cork_stream.rs index 4f01e7d..fe4b32c 100644 --- a/ktls/src/cork_stream.rs +++ b/ktls/src/cork_stream.rs @@ -4,7 +4,7 @@ use std::{io, task}; use rustls::internal::msgs::codec::Codec; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; -use crate::AsyncReadReady; +use crate::{AsyncReadReady, AsyncWriteReady}; enum State { ReadHeader { header_buf: [u8; 5], offset: usize }, @@ -176,6 +176,19 @@ where } } +impl AsyncWriteReady for CorkStream +where + IO: AsyncWriteReady, +{ + fn poll_write_ready(&self, cx: &mut task::Context<'_>) -> task::Poll> { + self.io.poll_write_ready(cx) + } + + fn try_write_io(&self, f: impl FnOnce() -> io::Result) -> io::Result { + self.io.try_write_io(f) + } +} + impl AsyncWrite for CorkStream where IO: AsyncWrite, diff --git a/ktls/src/ktls_stream.rs b/ktls/src/ktls_stream.rs index 8f77bbb..388421a 100644 --- a/ktls/src/ktls_stream.rs +++ b/ktls/src/ktls_stream.rs @@ -8,7 +8,8 @@ use nix::sys::socket::{recvmsg, ControlMessageOwned, MsgFlags, SockaddrIn, TlsGe use num_enum::FromPrimitive; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; -use crate::AsyncReadReady; +use crate::ffi::send_close_notify; +use crate::{AsyncReadReady, AsyncWriteReady}; // A wrapper around `IO` that sends a `close_notify` when shut down or dropped. pin_project_lite::pin_project! { @@ -205,9 +206,7 @@ where tracing::trace!(?level, ?description, "got TLS alert"); *this.read_closed = true; *this.write_closed = true; - if let Err(e) = - crate::ffi::send_close_notify(this.inner.as_raw_fd()) - { + if let Err(e) = send_close_notify(this.inner.as_raw_fd()) { // This can fail in case of a full send buffer (EAGAIN), // or a dead socket. Ignore the error, as replying with // close_notify is best-effort. @@ -299,7 +298,7 @@ where impl AsyncWrite for KtlsStream where - IO: AsRawFd + AsyncWrite, + IO: AsRawFd + AsyncWrite + AsyncWriteReady, { fn poll_write( self: Pin<&mut Self>, @@ -323,12 +322,22 @@ where ) -> task::Poll> { let this = self.project(); - if !*this.write_closed { + while !*this.write_closed { // they didn't hang up on us, we're nicely being asked to shut down, // let's send a close_notify (and not wait for them to send it back) - *this.write_closed = true; - if let Err(e) = crate::ffi::send_close_notify(this.inner.as_raw_fd()) { - return Err(e).into(); + task::ready!(this.inner.poll_write_ready(cx))?; + + let fd = this.inner.as_raw_fd(); + let res = this + .inner + .try_write_io(|| send_close_notify(fd)); + match res { + Ok(()) => *this.write_closed = true, + Err(e) if e.kind() == io::ErrorKind::WouldBlock => continue, + Err(e) => { + *this.write_closed = true; + return Err(e).into(); + } } } diff --git a/ktls/src/lib.rs b/ktls/src/lib.rs index 277ea47..6dded98 100644 --- a/ktls/src/lib.rs +++ b/ktls/src/lib.rs @@ -24,7 +24,7 @@ use smallvec::SmallVec; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite}; use tokio::net::{TcpListener, TcpStream}; -pub use crate::async_ready::AsyncReadReady; +pub use crate::async_ready::{AsyncReadReady, AsyncWriteReady}; pub use crate::cork_stream::CorkStream; pub use crate::ffi::CryptoInfo; use crate::ffi::{setup_tls_info, setup_ulp, KtlsCompatibilityError}; diff --git a/ktls/tests/integration_test.rs b/ktls/tests/integration_test.rs index 066ecdb..8b7d0af 100644 --- a/ktls/tests/integration_test.rs +++ b/ktls/tests/integration_test.rs @@ -3,7 +3,9 @@ use std::sync::Arc; use std::time::Duration; use std::{io, task}; -use ktls::{AsyncReadReady, CorkStream, KtlsCipherSuite, KtlsCipherType, KtlsVersion}; +use ktls::{ + AsyncReadReady, AsyncWriteReady, CorkStream, KtlsCipherSuite, KtlsCipherType, KtlsVersion, +}; use lazy_static::lazy_static; use rcgen::generate_simple_self_signed; use rustls::client::Resumption; @@ -563,6 +565,19 @@ where } } +impl AsyncWriteReady for SpyStream +where + IO: AsyncWriteReady, +{ + fn poll_write_ready(&self, cx: &mut task::Context<'_>) -> task::Poll> { + self.0.poll_write_ready(cx) + } + + fn try_write_io(&self, f: impl FnOnce() -> io::Result) -> io::Result { + self.0.try_write_io(f) + } +} + impl AsyncWrite for SpyStream where IO: AsyncWrite, @@ -862,3 +877,93 @@ async fn ktls_server_rustls_client( }; tokio::join!(server, client) } + +#[tokio::test] +async fn shutdown_retries_close_notify_when_send_buffer_full() { + let cipher_suite = KtlsCipherSuite { + version: KtlsVersion::TLS13, + typ: KtlsCipherType::AesGcm128, + }; + + let ckey = generate_simple_self_signed(vec!["localhost".to_string()]).unwrap(); + + let mut server_config = + ServerConfig::builder_with_provider(single_suite_provider(cipher_suite)) + .with_protocol_versions(&[cipher_suite + .version + .as_supported_version()]) + .unwrap() + .with_no_client_auth() + .with_single_cert( + vec![ckey.cert.der().clone()], + rustls::pki_types::PrivatePkcs8KeyDer::from(ckey.key_pair.serialize_der()).into(), + ) + .unwrap(); + server_config.enable_secret_extraction = true; + + let acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(server_config)); + let ln = TcpListener::bind("[::]:0") + .await + .unwrap(); + let addr = ln.local_addr().unwrap(); + + let mut root_store = RootCertStore::empty(); + root_store + .add(ckey.cert.der().clone()) + .unwrap(); + let client_config = ClientConfig::builder() + .with_root_certificates(root_store) + .with_no_client_auth(); + let tls_connector = TlsConnector::from(Arc::new(client_config)); + + let (drain_tx, drain_rx) = tokio::sync::oneshot::channel::<()>(); + + let jh = tokio::spawn(async move { + let stream = TcpStream::connect(addr).await.unwrap(); + let mut stream = tls_connector + .connect("localhost".try_into().unwrap(), stream) + .await + .unwrap(); + + // 3. Drain everything + drain_rx.await.unwrap(); + let mut sink = vec![0u8; 65536]; + loop { + match stream.read(&mut sink).await.unwrap() { + // EOF signals the `close_notify` was delivered + 0 => break, + _ => continue, + } + } + }); + + let (stream, _) = ln.accept().await.unwrap(); + socket2::SockRef::from(&stream) + .set_send_buffer_size(4096) + .unwrap(); + let stream = CorkStream::new(stream); + let stream = acceptor.accept(stream).await.unwrap(); + let mut stream = ktls::config_ktls_server(stream) + .await + .unwrap(); + + // 1. Fill the send buffer (the client is not reading yet). + let chunk = vec![0u8; 65536]; + while let Ok(res) = tokio::time::timeout(Duration::from_millis(250), stream.write(&chunk)).await + { + res.unwrap(); + } + + // 2. With no room for the alert, shutdown must stay pending, not fail. + let res = tokio::time::timeout(Duration::from_millis(250), stream.shutdown()).await; + assert!( + res.is_err(), + "shutdown must stay pending while the buffer is full, got {res:?}" + ); + + // 4. Signal client to start draining. The retried shutdown now completes. + drain_tx.send(()).unwrap(); + stream.shutdown().await.unwrap(); + + jh.await.unwrap(); +}