From e862e48240d1070bfd471aa7e4ae0970931ea9f9 Mon Sep 17 00:00:00 2001 From: Arni Dagur Date: Fri, 7 Aug 2026 22:44:59 +0100 Subject: [PATCH] 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_write_ready.rs | 24 ++++++++ ktls/src/cork_stream.rs | 15 ++++- ktls/src/ktls_stream.rs | 22 +++++-- ktls/src/lib.rs | 2 + ktls/tests/integration_test.rs | 107 ++++++++++++++++++++++++++++++++- 5 files changed, 162 insertions(+), 8 deletions(-) create mode 100644 ktls/src/async_write_ready.rs diff --git a/ktls/src/async_write_ready.rs b/ktls/src/async_write_ready.rs new file mode 100644 index 0000000..46ec23e --- /dev/null +++ b/ktls/src/async_write_ready.rs @@ -0,0 +1,24 @@ +use std::{io, task}; + +use tokio::io::Interest; + +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 task::Context<'_>) -> task::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 task::Context<'_>) -> task::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..4cdf610 100644 --- a/ktls/src/ktls_stream.rs +++ b/ktls/src/ktls_stream.rs @@ -8,7 +8,7 @@ use nix::sys::socket::{recvmsg, ControlMessageOwned, MsgFlags, SockaddrIn, TlsGe use num_enum::FromPrimitive; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; -use crate::AsyncReadReady; +use crate::{AsyncReadReady, AsyncWriteReady}; // A wrapper around `IO` that sends a `close_notify` when shut down or dropped. pin_project_lite::pin_project! { @@ -299,7 +299,7 @@ where impl AsyncWrite for KtlsStream where - IO: AsRawFd + AsyncWrite, + IO: AsRawFd + AsyncWrite + AsyncWriteReady, { fn poll_write( self: Pin<&mut Self>, @@ -323,12 +323,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(|| crate::ffi::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 7e31d86..289728d 100644 --- a/ktls/src/lib.rs +++ b/ktls/src/lib.rs @@ -4,6 +4,7 @@ compile_error!("This crate needs wither the 'ring' or 'aws_lc_rs' feature enable compile_error!("The 'ring' and 'aws_lc_rs' features are mutually exclusive"); mod async_read_ready; +mod async_write_ready; mod cork_stream; mod ffi; mod ktls_stream; @@ -25,6 +26,7 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite}; use tokio::net::{TcpListener, TcpStream}; pub use crate::async_read_ready::AsyncReadReady; +pub use crate::async_write_ready::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(); +}