fix: handshake does not fully flush writes (#112)

conradludgate · web-flow · commit dd5f27d78c7e · 2025-06-06T21:54:50.000+08:00
* write test to verify that flush does not get resumed

* attempt to persist the need_flush state during handshake
diff --git a/src/client.rs b/src/client.rs
@@ -20,6 +20,7 @@ pub struct TlsStream<IO> {
     pub(crate) io: IO,
     pub(crate) session: ClientConnection,
     pub(crate) state: TlsState,
+    pub(crate) need_flush: bool,
 
     #[cfg(feature = "early-data")]
     pub(crate) early_waker: Option<Waker>,
@@ -72,8 +73,13 @@ impl<IO> IoSession for TlsStream<IO> {
     }
 
     #[inline]
-    fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session) {
-        (&mut self.state, &mut self.io, &mut self.session)
+    fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session, &mut bool) {
+        (
+            &mut self.state,
+            &mut self.io,
+            &mut self.session,
+            &mut self.need_flush,
+        )
     }
 
     #[inline]
@@ -174,21 +180,27 @@ where
         buf: &[u8],
     ) -> Poll<io::Result<usize>> {
         let this = self.get_mut();
-        let mut stream =
-            Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
+        let mut stream = Stream::new(&mut this.io, &mut this.session)
+            .set_eof(!this.state.readable())
+            .set_need_flush(this.need_flush);
 
         #[cfg(feature = "early-data")]
         {
             let bufs = [io::IoSlice::new(buf)];
-            let written = ready!(poll_handle_early_data(
+            let written = poll_handle_early_data(
                 &mut this.state,
                 &mut stream,
                 &mut this.early_waker,
                 cx,
-                &bufs
-            ))?;
-            if written != 0 {
-                return Poll::Ready(Ok(written));
+                &bufs,
+            )?;
+            match written {
+                Poll::Ready(0) => {}
+                Poll::Ready(written) => return Poll::Ready(Ok(written)),
+                Poll::Pending => {
+                    this.need_flush = stream.need_flush;
+                    return Poll::Pending;
+                }
             }
         }
 
@@ -203,20 +215,26 @@ where
         bufs: &[io::IoSlice<'_>],
     ) -> Poll<io::Result<usize>> {
         let this = self.get_mut();
-        let mut stream =
-            Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
+        let mut stream = Stream::new(&mut this.io, &mut this.session)
+            .set_eof(!this.state.readable())
+            .set_need_flush(this.need_flush);
 
         #[cfg(feature = "early-data")]
         {
-            let written = ready!(poll_handle_early_data(
+            let written = poll_handle_early_data(
                 &mut this.state,
                 &mut stream,
                 &mut this.early_waker,
                 cx,
-                bufs
-            ))?;
-            if written != 0 {
-                return Poll::Ready(Ok(written));
+                bufs,
+            )?;
+            match written {
+                Poll::Ready(0) => {}
+                Poll::Ready(written) => return Poll::Ready(Ok(written)),
+                Poll::Pending => {
+                    this.need_flush = stream.need_flush;
+                    return Poll::Pending;
+                }
             }
         }
 
@@ -230,17 +248,24 @@ where
 
     fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
         let this = self.get_mut();
-        let mut stream =
-            Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
+        let mut stream = Stream::new(&mut this.io, &mut this.session)
+            .set_eof(!this.state.readable())
+            .set_need_flush(this.need_flush);
 
         #[cfg(feature = "early-data")]
-        ready!(poll_handle_early_data(
-            &mut this.state,
-            &mut stream,
-            &mut this.early_waker,
-            cx,
-            &[]
-        ))?;
+        {
+            let written = poll_handle_early_data(
+                &mut this.state,
+                &mut stream,
+                &mut this.early_waker,
+                cx,
+                &[],
+            )?;
+            if written.is_pending() {
+                this.need_flush = stream.need_flush;
+                return Poll::Pending;
+            }
+        }
 
         stream.as_mut_pin().poll_flush(cx)
     }
diff --git a/src/common/handshake.rs b/src/common/handshake.rs
@@ -15,7 +15,7 @@ pub(crate) trait IoSession {
     type Session;
 
     fn skip_handshake(&self) -> bool;
-    fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session);
+    fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session, &mut bool);
     fn into_io(self) -> Self::Io;
 }
 
@@ -67,15 +67,18 @@ where
         };
 
         if !stream.skip_handshake() {
-            let (state, io, session) = stream.get_mut();
-            let mut tls_stream = Stream::new(io, session).set_eof(!state.readable());
+            let (state, io, session, need_flush) = stream.get_mut();
+            let mut tls_stream = Stream::new(io, session)
+                .set_eof(!state.readable())
+                .set_need_flush(*need_flush);
 
             macro_rules! try_poll {
                 ( $e:expr ) => {
                     match $e {
-                        Poll::Ready(Ok(_)) => (),
+                        Poll::Ready(Ok(x)) => x,
                         Poll::Ready(Err(err)) => return Poll::Ready(Err((err, stream.into_io()))),
                         Poll::Pending => {
+                            *need_flush = tls_stream.need_flush;
                             *this = MidHandshake::Handshaking(stream);
                             return Poll::Pending;
                         }
diff --git a/src/common/mod.rs b/src/common/mod.rs
@@ -63,6 +63,7 @@ pub(crate) struct Stream<'a, IO, C> {
     pub(crate) io: &'a mut IO,
     pub(crate) session: &'a mut C,
     pub(crate) eof: bool,
+    pub(crate) need_flush: bool,
 }
 
 impl<'a, IO: AsyncRead + AsyncWrite + Unpin, C, SD> Stream<'a, IO, C>
@@ -77,6 +78,8 @@ where
             // The state so far is only used to detect EOF, so either Stream
             // or EarlyData state should both be all right.
             eof: false,
+            // Whether a previous flush returned pending, or a write occured without a flush.
+            need_flush: false,
         }
     }
 
@@ -85,6 +88,11 @@ where
         self
     }
 
+    pub(crate) fn set_need_flush(mut self, need_flush: bool) -> Self {
+        self.need_flush = need_flush;
+        self
+    }
+
     pub(crate) fn as_mut_pin(&mut self) -> Pin<&mut Self> {
         Pin::new(self)
     }
@@ -126,14 +134,13 @@ where
         loop {
             let mut write_would_block = false;
             let mut read_would_block = false;
-            let mut need_flush = false;
 
             while self.session.wants_write() {
                 match self.write_io(cx) {
                     Poll::Ready(Ok(0)) => return Poll::Ready(Err(io::ErrorKind::WriteZero.into())),
                     Poll::Ready(Ok(n)) => {
                         wrlen += n;
-                        need_flush = true;
+                        self.need_flush = true;
                     }
                     Poll::Pending => {
                         write_would_block = true;
@@ -143,9 +150,9 @@ where
                 }
             }
 
-            if need_flush {
+            if self.need_flush {
                 match Pin::new(&mut self.io).poll_flush(cx) {
-                    Poll::Ready(Ok(())) => (),
+                    Poll::Ready(Ok(())) => self.need_flush = false,
                     Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
                     Poll::Pending => write_would_block = true,
                 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -149,6 +149,8 @@ impl TlsConnector {
                 TlsState::Stream
             },
 
+            need_flush: false,
+
             #[cfg(feature = "early-data")]
             early_waker: None,
 
@@ -193,6 +195,7 @@ impl TlsAcceptor {
             session,
             io: stream,
             state: TlsState::Stream,
+            need_flush: false,
         }))
     }
 
@@ -363,6 +366,7 @@ where
             session: conn,
             io: self.io,
             state: TlsState::Stream,
+            need_flush: false,
         }))
     }
 }
diff --git a/src/server.rs b/src/server.rs
@@ -18,6 +18,7 @@ pub struct TlsStream<IO> {
     pub(crate) io: IO,
     pub(crate) session: ServerConnection,
     pub(crate) state: TlsState,
+    pub(crate) need_flush: bool,
 }
 
 impl<IO> TlsStream<IO> {
@@ -47,8 +48,13 @@ impl<IO> IoSession for TlsStream<IO> {
     }
 
     #[inline]
-    fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session) {
-        (&mut self.state, &mut self.io, &mut self.session)
+    fn get_mut(&mut self) -> (&mut TlsState, &mut Self::Io, &mut Self::Session, &mut bool) {
+        (
+            &mut self.state,
+            &mut self.io,
+            &mut self.session,
+            &mut self.need_flush,
+        )
     }
 
     #[inline]
diff --git a/tests/early-data.rs b/tests/early-data.rs
@@ -10,7 +10,7 @@ use std::thread;
 use futures_util::{future::Future, ready};
 use rustls::pki_types::ServerName;
 use rustls::{self, ClientConfig, ServerConnection, Stream};
-use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt, ReadBuf};
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
 use tokio::net::TcpStream;
 use tokio_rustls::client::TlsStream;
 use tokio_rustls::TlsConnector;
@@ -35,14 +35,15 @@ impl<T: AsyncRead + Unpin> Future for Read1<T> {
     }
 }
 
-async fn send(
+async fn send<S: AsyncRead + AsyncWrite + Unpin>(
     config: Arc<ClientConfig>,
     addr: SocketAddr,
+    wrapper: impl Fn(TcpStream) -> S,
     data: &[u8],
     vectored: bool,
-) -> io::Result<(TlsStream<TcpStream>, Vec<u8>)> {
+) -> io::Result<(TlsStream<S>, Vec<u8>)> {
     let connector = TlsConnector::from(config).early_data(true);
-    let stream = TcpStream::connect(&addr).await?;
+    let stream = wrapper(TcpStream::connect(&addr).await?);
     let domain = ServerName::try_from("foobar.com").unwrap();
 
     let mut stream = connector.connect(domain, stream).await?;
@@ -58,15 +59,23 @@ async fn send(
 
 #[tokio::test]
 async fn test_0rtt() -> io::Result<()> {
-    test_0rtt_impl(false).await
+    test_0rtt_impl(|s| s, false).await
 }
 
 #[tokio::test]
 async fn test_0rtt_vectored() -> io::Result<()> {
-    test_0rtt_impl(true).await
+    test_0rtt_impl(|s| s, true).await
 }
 
-async fn test_0rtt_impl(vectored: bool) -> io::Result<()> {
+#[tokio::test]
+async fn test_0rtt_vectored_flush_pending() -> io::Result<()> {
+    test_0rtt_impl(utils::FlushWrapper::new, false).await
+}
+
+async fn test_0rtt_impl<S: AsyncRead + AsyncWrite + Unpin>(
+    wrapper: impl Fn(TcpStream) -> S,
+    vectored: bool,
+) -> io::Result<()> {
     let (mut server, mut client) = utils::make_configs();
     server.max_early_data_size = 8192;
     let server = Arc::new(server);
@@ -108,11 +117,11 @@ async fn test_0rtt_impl(vectored: bool) -> io::Result<()> {
     let client = Arc::new(client);
     let addr = SocketAddr::from(([127, 0, 0, 1], server_port));
 
-    let (io, buf) = send(client.clone(), addr, b"hello", vectored).await?;
+    let (io, buf) = send(client.clone(), addr, &wrapper, b"hello", vectored).await?;
     assert!(!io.get_ref().1.is_early_data_accepted());
     assert_eq!("LATE:hello", String::from_utf8_lossy(&buf));
 
-    let (io, buf) = send(client, addr, b"world!", vectored).await?;
+    let (io, buf) = send(client, addr, wrapper, b"world!", vectored).await?;
     assert!(io.get_ref().1.is_early_data_accepted());
     assert_eq!("EARLY:world!LATE:", String::from_utf8_lossy(&buf));
 
diff --git a/tests/test.rs b/tests/test.rs
diff --git a/tests/utils.rs b/tests/utils.rs

Original file line number	Diff line number	Diff line change
`@@ -63,6 +63,7 @@ pub(crate) struct Stream<'a, IO, C> {`
`63`	`63`	`pub(crate) io: &'a mut IO,`
`64`	`64`	`pub(crate) session: &'a mut C,`
`65`	`65`	`pub(crate) eof: bool,`
	`66`	`+ pub(crate) need_flush: bool,`
`66`	`67`	`}`
`67`	`68`
`68`	`69`	`impl<'a, IO: AsyncRead + AsyncWrite + Unpin, C, SD> Stream<'a, IO, C>`
`@@ -77,6 +78,8 @@ where`
`77`	`78`	`// The state so far is only used to detect EOF, so either Stream`
`78`	`79`	`// or EarlyData state should both be all right.`
`79`	`80`	`eof: false,`
	`81`	`+ // Whether a previous flush returned pending, or a write occured without a flush.`
	`82`	`+ need_flush: false,`
`80`	`83`	`}`
`81`	`84`	`}`
`82`	`85`
`@@ -85,6 +88,11 @@ where`
`85`	`88`	`self`
`86`	`89`	`}`
`87`	`90`
	`91`	`+ pub(crate) fn set_need_flush(mut self, need_flush: bool) -> Self {`
	`92`	`+ self.need_flush = need_flush;`
	`93`	`+ self`
	`94`	`+ }`
	`95`	`+`
`88`	`96`	`pub(crate) fn as_mut_pin(&mut self) -> Pin<&mut Self> {`
`89`	`97`	`Pin::new(self)`
`90`	`98`	`}`
`@@ -126,14 +134,13 @@ where`
`126`	`134`	`loop {`
`127`	`135`	`let mut write_would_block = false;`
`128`	`136`	`let mut read_would_block = false;`
`129`		`- let mut need_flush = false;`
`130`	`137`
`131`	`138`	`while self.session.wants_write() {`
`132`	`139`	`match self.write_io(cx) {`
`133`	`140`	`Poll::Ready(Ok(0)) => return Poll::Ready(Err(io::ErrorKind::WriteZero.into())),`
`134`	`141`	`Poll::Ready(Ok(n)) => {`
`135`	`142`	`wrlen += n;`
`136`		`- need_flush = true;`
	`143`	`+ self.need_flush = true;`
`137`	`144`	`}`
`138`	`145`	`Poll::Pending => {`
`139`	`146`	`write_would_block = true;`
`@@ -143,9 +150,9 @@ where`
`143`	`150`	`}`
`144`	`151`	`}`
`145`	`152`
`146`		`- if need_flush {`
	`153`	`+ if self.need_flush {`
`147`	`154`	`match Pin::new(&mut self.io).poll_flush(cx) {`
`148`		`- Poll::Ready(Ok(())) => (),`
	`155`	`+ Poll::Ready(Ok(())) => self.need_flush = false,`
`149`	`156`	`Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),`
`150`	`157`	`Poll::Pending => write_would_block = true,`
`151`	`158`	`}`