Forward vectored writes (#45)

* Migrate early-data test to rustls

* Replace `match` with `if let` on `TlsState::EarlyData`

* Extract client early data handling

* Forward vectored writes

Author: paolobarbolini

src/client.rs

@@ -4,6 +4,8 @@ use std::os::unix::io::{AsRawFd, RawFd};
 #[cfg(windows)]
 use std::os::windows::io::{AsRawSocket, RawSocket};
 use std::pin::Pin;
+#[cfg(feature = "early-data")]
+use std::task::Waker;
 use std::task::{Context, Poll};
 
 use rustls::ClientConnection;
@@ -20,7 +22,7 @@ pub struct TlsStream<IO> {
     pub(crate) state: TlsState,
 
     #[cfg(feature = "early-data")]
-    pub(crate) early_waker: Option<std::task::Waker>,
+    pub(crate) early_waker: Option<Waker>,
 }
 
 impl<IO> TlsStream<IO> {
@@ -152,78 +154,70 @@ where
         let mut stream =
             Stream::new(&mut this.io, &mut this.session).set_eof(!this.state.readable());
 
-        #[allow(clippy::match_single_binding)]
-        match this.state {
-            #[cfg(feature = "early-data")]
-            TlsState::EarlyData(ref mut pos, ref mut data) => {
-                use std::io::Write;
-
-                // write early data
-                if let Some(mut early_data) = stream.session.early_data() {
-                    let len = match early_data.write(buf) {
-                        Ok(n) => n,
-                        Err(err) => return Poll::Ready(Err(err)),
-                    };
-                    if len != 0 {
-                        data.extend_from_slice(&buf[..len]);
-                        return Poll::Ready(Ok(len));
-                    }
-                }
-
-                // complete handshake
-                while stream.session.is_handshaking() {
-                    ready!(stream.handshake(cx))?;
-                }
-
-                // write early data (fallback)
-                if !stream.session.is_early_data_accepted() {
-                    while *pos < data.len() {
-                        let len = ready!(stream.as_mut_pin().poll_write(cx, &data[*pos..]))?;
-                        *pos += len;
-                    }
-                }
-
-                // end
-                this.state = TlsState::Stream;
-
-                if let Some(waker) = this.early_waker.take() {
-                    waker.wake();
-                }
-
-                stream.as_mut_pin().poll_write(cx, buf)
+        #[cfg(feature = "early-data")]
+        {
+            let bufs = [io::IoSlice::new(buf)];
+            let written = ready!(poll_handle_early_data(
+                &mut this.state,
+                &mut stream,
+                &mut this.early_waker,
+                cx,
+                &bufs
+            ))?;
+            if written != 0 {
+                return Poll::Ready(Ok(written));
             }
-            _ => stream.as_mut_pin().poll_write(cx, buf),
         }
+
+        stream.as_mut_pin().poll_write(cx, buf)
     }
 
-    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+    /// Note: that it does not guarantee the final data to be sent.
+    /// To be cautious, you must manually call `flush`.
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        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());
 
         #[cfg(feature = "early-data")]
         {
-            if let TlsState::EarlyData(ref mut pos, ref mut data) = this.state {
-                // complete handshake
-                while stream.session.is_handshaking() {
-                    ready!(stream.handshake(cx))?;
-                }
+            let written = ready!(poll_handle_early_data(
+                &mut this.state,
+                &mut stream,
+                &mut this.early_waker,
+                cx,
+                bufs
+            ))?;
+            if written != 0 {
+                return Poll::Ready(Ok(written));
+            }
+        }
 
-                // write early data (fallback)
-                if !stream.session.is_early_data_accepted() {
-                    while *pos < data.len() {
-                        let len = ready!(stream.as_mut_pin().poll_write(cx, &data[*pos..]))?;
-                        *pos += len;
-                    }
-                }
+        stream.as_mut_pin().poll_write_vectored(cx, bufs)
+    }
 
-                this.state = TlsState::Stream;
+    #[inline]
+    fn is_write_vectored(&self) -> bool {
+        true
+    }
 
-                if let Some(waker) = this.early_waker.take() {
-                    waker.wake();
-                }
-            }
-        }
+    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());
+
+        #[cfg(feature = "early-data")]
+        ready!(poll_handle_early_data(
+            &mut this.state,
+            &mut stream,
+            &mut this.early_waker,
+            cx,
+            &[]
+        ))?;
 
         stream.as_mut_pin().poll_flush(cx)
     }
@@ -248,3 +242,69 @@ where
         stream.as_mut_pin().poll_shutdown(cx)
     }
 }
+
+#[cfg(feature = "early-data")]
+fn poll_handle_early_data<IO>(
+    state: &mut TlsState,
+    stream: &mut Stream<IO, ClientConnection>,
+    early_waker: &mut Option<Waker>,
+    cx: &mut Context<'_>,
+    bufs: &[io::IoSlice<'_>],
+) -> Poll<io::Result<usize>>
+where
+    IO: AsyncRead + AsyncWrite + Unpin,
+{
+    if let TlsState::EarlyData(pos, data) = state {
+        use std::io::Write;
+
+        // write early data
+        if let Some(mut early_data) = stream.session.early_data() {
+            let mut written = 0;
+
+            for buf in bufs {
+                if buf.is_empty() {
+                    continue;
+                }
+
+                let len = match early_data.write(buf) {
+                    Ok(0) => break,
+                    Ok(n) => n,
+                    Err(err) => return Poll::Ready(Err(err)),
+                };
+
+                written += len;
+                data.extend_from_slice(&buf[..len]);
+
+                if len < buf.len() {
+                    break;
+                }
+            }
+
+            if written != 0 {
+                return Poll::Ready(Ok(written));
+            }
+        }
+
+        // complete handshake
+        while stream.session.is_handshaking() {
+            ready!(stream.handshake(cx))?;
+        }
+
+        // write early data (fallback)
+        if !stream.session.is_early_data_accepted() {
+            while *pos < data.len() {
+                let len = ready!(stream.as_mut_pin().poll_write(cx, &data[*pos..]))?;
+                *pos += len;
+            }
+        }
+
+        // end
+        *state = TlsState::Stream;
+
+        if let Some(waker) = early_waker.take() {
+            waker.wake();
+        }
+    }
+
+    Poll::Ready(Ok(0))
+}

src/common/mod.rs

@@ -289,6 +289,43 @@ where
         Poll::Ready(Ok(pos))
     }
 
+    fn poll_write_vectored(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[IoSlice<'_>],
+    ) -> Poll<io::Result<usize>> {
+        if bufs.iter().all(|buf| buf.is_empty()) {
+            return Poll::Ready(Ok(0));
+        }
+
+        loop {
+            let mut would_block = false;
+            let written = self.session.writer().write_vectored(bufs)?;
+
+            while self.session.wants_write() {
+                match self.write_io(cx) {
+                    Poll::Ready(Ok(0)) | Poll::Pending => {
+                        would_block = true;
+                        break;
+                    }
+                    Poll::Ready(Ok(_)) => (),
+                    Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
+                }
+            }
+
+            return match (written, would_block) {
+                (0, true) => Poll::Pending,
+                (0, false) => continue,
+                (n, _) => Poll::Ready(Ok(n)),
+            };
+        }
+    }
+
+    #[inline]
+    fn is_write_vectored(&self) -> bool {
+        true
+    }
+
     fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
         self.session.writer().flush()?;
         while self.session.wants_write() {

src/common/test_stream.rs

@@ -122,6 +122,15 @@ impl AsyncWrite for Expected {
 
 #[tokio::test]
 async fn stream_good() -> io::Result<()> {
+    stream_good_impl(false).await
+}
+
+#[tokio::test]
+async fn stream_good_vectored() -> io::Result<()> {
+    stream_good_impl(true).await
+}
+
+async fn stream_good_impl(vectored: bool) -> io::Result<()> {
     const FILE: &[u8] = include_bytes!("../../README.md");
 
     let (server, mut client) = make_pair();
@@ -139,7 +148,7 @@ async fn stream_good() -> io::Result<()> {
         dbg!(stream.read_to_end(&mut buf).await)?;
         assert_eq!(buf, FILE);
 
-        dbg!(stream.write_all(b"Hello World!").await)?;
+        dbg!(utils::write(&mut stream, b"Hello World!", vectored).await)?;
         stream.session.send_close_notify();
 
         dbg!(stream.shutdown().await)?;

src/lib.rs

@@ -564,6 +564,26 @@ where
         }
     }
 
+    #[inline]
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[io::IoSlice<'_>],
+    ) -> Poll<io::Result<usize>> {
+        match self.get_mut() {
+            TlsStream::Client(x) => Pin::new(x).poll_write_vectored(cx, bufs),
+            TlsStream::Server(x) => Pin::new(x).poll_write_vectored(cx, bufs),
+        }
+    }
+
+    #[inline]
+    fn is_write_vectored(&self) -> bool {
+        match self {
+            TlsStream::Client(x) => x.is_write_vectored(),
+            TlsStream::Server(x) => x.is_write_vectored(),
+        }
+    }
+
     #[inline]
     fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
         match self.get_mut() {

src/server.rs

@@ -113,6 +113,24 @@ where
         stream.as_mut_pin().poll_write(cx, buf)
     }
 
+    /// Note: that it does not guarantee the final data to be sent.
+    /// To be cautious, you must manually call `flush`.
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        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());
+        stream.as_mut_pin().poll_write_vectored(cx, bufs)
+    }
+
+    #[inline]
+    fn is_write_vectored(&self) -> bool {
+        true
+    }
+
     fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
         let this = self.get_mut();
         let mut stream =