fix: retry close_notify in poll_shutdown when the send buffer is full

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<io::Result<()>>;
+
+    /// 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<R>(&self, f: impl FnOnce() -> io::Result<R>) -> io::Result<R>;
+}
+
+impl AsyncWriteReady for tokio::net::TcpStream {
+    fn poll_write_ready(&self, cx: &mut task::Context<'_>) -> task::Poll<io::Result<()>> {
+        tokio::net::TcpStream::poll_write_ready(self, cx)
+    }
+
+    fn try_write_io<R>(&self, f: impl FnOnce() -> io::Result<R>) -> io::Result<R> {
+        self.try_io(Interest::WRITABLE, f)
+    }
+}

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<IO> AsyncWriteReady for CorkStream<IO>
+where
+    IO: AsyncWriteReady,
+{
+    fn poll_write_ready(&self, cx: &mut task::Context<'_>) -> task::Poll<io::Result<()>> {
+        self.io.poll_write_ready(cx)
+    }
+
+    fn try_write_io<R>(&self, f: impl FnOnce() -> io::Result<R>) -> io::Result<R> {
+        self.io.try_write_io(f)
+    }
+}
+
 impl<IO> AsyncWrite for CorkStream<IO>
 where
     IO: AsyncWrite,

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<IO> AsyncWrite for KtlsStream<IO>
 where
-    IO: AsRawFd + AsyncWrite,
+    IO: AsRawFd + AsyncWrite + AsyncWriteReady,
 {
     fn poll_write(
         self: Pin<&mut Self>,
@@ -323,12 +323,22 @@ where
     ) -> task::Poll<io::Result<()>> {
         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();
+                }
             }
         }
 

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};

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<IO> AsyncWriteReady for SpyStream<IO>
+where
+    IO: AsyncWriteReady,
+{
+    fn poll_write_ready(&self, cx: &mut task::Context<'_>) -> task::Poll<io::Result<()>> {
+        self.0.poll_write_ready(cx)
+    }
+
+    fn try_write_io<R>(&self, f: impl FnOnce() -> io::Result<R>) -> io::Result<R> {
+        self.0.try_write_io(f)
+    }
+}
+
 impl<IO> AsyncWrite for SpyStream<IO>
 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();
+}