Switch to recvmsg

Cargo.toml

@@ -22,10 +22,15 @@ pin-project-lite = "0.2.9"
 tokio = { version = "1.28.2", features = ["net", "macros", "io-util"] }
 futures = "0.3.28"
 ktls-sys = "1.0.0"
+nix = { version = "0.26.1", features = ["socket", "uio", "net"], default-features = false }
 
 [dev-dependencies]
 const-random = "0.1.15"
 rcgen = "0.10.0"
 socket2 = "0.5.3"
 tokio = { version = "1.28.2", features = ["full"] }
 tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
+
+[patch.crates-io]
+nix = { git = "https://github.com/fasterthanlime/nix", rev = "004d31c" } # on branch 'sol_tls'
+

src/async_read_ready.rs

@@ -0,0 +1,12 @@
+use std::{io, task};
+
+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<io::Result<()>>;
+}
+
+impl AsyncReadReady for tokio::net::TcpStream {
+    fn poll_read_ready(&self, cx: &mut task::Context<'_>) -> task::Poll<io::Result<()>> {
+        tokio::net::TcpStream::poll_read_ready(self, cx)
+    }
+}

src/cork_stream.rs

@@ -1,12 +1,10 @@
-use std::{
-    io,
-    pin::Pin,
-    task::{Context, Poll},
-};
+use std::{io, pin::Pin, task};
 
 use rustls::internal::msgs::codec::Codec;
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 
+use crate::AsyncReadReady;
+
 enum State {
     ReadHeader { header_buf: [u8; 5], offset: usize },
     ReadPayload { msg_size: usize, offset: usize },
@@ -59,9 +57,9 @@ where
     #[inline]
     fn poll_read(
         self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
+        cx: &mut task::Context<'_>,
         buf: &mut ReadBuf<'_>,
-    ) -> Poll<io::Result<()>> {
+    ) -> task::Poll<io::Result<()>> {
         let this = unsafe { self.get_unchecked_mut() };
         let mut io = unsafe { Pin::new_unchecked(&mut this.io) };
 
@@ -75,7 +73,7 @@ where
                             "corked, returning empty read (but waking to prevent stalls)"
                         );
                         cx.waker().wake_by_ref();
-                        return Poll::Ready(Ok(()));
+                        return task::Poll::Ready(Ok(()));
                     }
 
                     let left = header_buf.len() - *offset;
@@ -97,7 +95,7 @@ where
                             buf.put_slice(&header_buf[..*offset]);
                             *state = State::Passthrough;
 
-                            return Poll::Ready(Ok(()));
+                            return task::Poll::Ready(Ok(()));
                         }
                         tracing::trace!("read {} bytes off of header", rest.filled().len());
                     }
@@ -127,7 +125,7 @@ where
                             }
                         }
 
-                        return Poll::Ready(Ok(()));
+                        return task::Poll::Ready(Ok(()));
                     } else {
                         // keep trying
                     }
@@ -156,7 +154,7 @@ where
                     let new_filled = buf.filled().len() + just_read;
                     buf.set_filled(new_filled);
 
-                    return Poll::Ready(Ok(()));
+                    return task::Poll::Ready(Ok(()));
                 }
                 State::Passthrough => {
                     // we encountered EOF while reading, or saw an invalid header and we're just
@@ -168,28 +166,40 @@ where
     }
 }
 
+impl<IO> AsyncReadReady for CorkStream<IO>
+where
+    IO: AsyncReadReady,
+{
+    fn poll_read_ready(&self, cx: &mut task::Context<'_>) -> task::Poll<io::Result<()>> {
+        self.io.poll_read_ready(cx)
+    }
+}
+
 impl<IO> AsyncWrite for CorkStream<IO>
 where
     IO: AsyncWrite,
 {
     #[inline]
     fn poll_write(
         self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
+        cx: &mut task::Context<'_>,
         buf: &[u8],
-    ) -> Poll<io::Result<usize>> {
+    ) -> task::Poll<io::Result<usize>> {
         let io = unsafe { self.map_unchecked_mut(|s| &mut s.io) };
         io.poll_write(cx, buf)
     }
 
     #[inline]
-    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<io::Result<()>> {
         let io = unsafe { self.map_unchecked_mut(|s| &mut s.io) };
         io.poll_flush(cx)
     }
 
     #[inline]
-    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        cx: &mut task::Context<'_>,
+    ) -> task::Poll<io::Result<()>> {
         let io = unsafe { self.map_unchecked_mut(|s| &mut s.io) };
         io.poll_shutdown(cx)
     }

src/ktls_stream.rs

@@ -1,7 +1,18 @@
-use std::{io, os::unix::prelude::AsRawFd, pin::Pin, task};
-
+use std::{
+    io::{self, IoSliceMut},
+    os::unix::prelude::AsRawFd,
+    pin::Pin,
+    task,
+};
+
+use nix::{
+    cmsg_space,
+    sys::socket::{ControlMessageOwned, MsgFlags, SockaddrIn},
+};
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 
+use crate::AsyncReadReady;
+
 // A wrapper around `IO` that sends a `close_notify` when shut down or dropped.
 pin_project_lite::pin_project! {
     pub struct KtlsStream<IO>
@@ -54,14 +65,14 @@ where
 
 impl<IO> AsyncRead for KtlsStream<IO>
 where
-    IO: AsRawFd + AsyncRead,
+    IO: AsRawFd + AsyncRead + AsyncReadReady,
 {
     fn poll_read(
         self: Pin<&mut Self>,
         cx: &mut task::Context<'_>,
         buf: &mut ReadBuf<'_>,
     ) -> task::Poll<io::Result<()>> {
-        tracing::trace!(remaining = %buf.remaining(), "KtlsStream::poll_read");
+        tracing::trace!(buf.remaining = %buf.remaining(), "KtlsStream::poll_read");
 
         let this = self.project();
 
@@ -79,11 +90,57 @@ where
             }
             cx.waker().wake_by_ref();
 
+            tracing::trace!("KtlsStream::poll_read, returning after drain");
             return task::Poll::Ready(Ok(()));
         }
 
-        tracing::trace!("KtlsStream::poll_read, forwarding to inner IO");
-        this.inner.poll_read(cx, buf)
+        let fd = this.inner.as_raw_fd();
+
+        let mut cmsgspace = cmsg_space!(nix::sys::time::TimeVal);
+        let mut iov = [IoSliceMut::new(buf.initialize_unfilled())];
+        let flags = MsgFlags::empty();
+
+        let r = nix::sys::socket::recvmsg::<SockaddrIn>(fd, &mut iov, Some(&mut cmsgspace), flags);
+        let r = match r {
+            Ok(r) => r,
+            Err(nix::errno::Errno::EAGAIN) => {
+                // this time don't `wake_by_ref` on purpose, but clear readiness by calling
+                // poll_read, knowing it'll fail
+                tracing::trace!("KtlsStream::poll_read, recvmsg gave us EAGAIN/EWOULDBLOCK");
+                match this.inner.poll_read(cx, buf) {
+                    task::Poll::Ready(_) => {
+                        unreachable!("one of ktls's core assumptions about async I/O didn't hold")
+                    }
+                    task::Poll::Pending => return task::Poll::Pending,
+                }
+            }
+            Err(e) => {
+                tracing::trace!(?e, "recvmsg failed");
+                return Err(e.into()).into();
+            }
+        };
+        let cmsg = r.cmsgs().next().unwrap();
+        tracing::trace!("cmsg = {cmsg:#?}");
+        let message_type = match cmsg {
+            ControlMessageOwned::TlsGetRecordType(t) => t,
+            _ => panic!("unexpected cmsg type: {cmsg:#?}"),
+        };
+        match message_type {
+            23 => {
+                tracing::trace!(%r.bytes, "KtlsStream::poll_read, returning Ok");
+                let read_bytes = r.bytes;
+                buf.advance(read_bytes);
+                task::Poll::Ready(Ok(()))
+            }
+            _ => {
+                tracing::trace!("received message_type {message_type:#?}");
+
+                // FIXME: hacky, maybe there's a way to avoid a loop here if
+                // the socket isn't actually ready to be read
+                cx.waker().wake_by_ref();
+                task::Poll::Pending
+            }
+        }
     }
 }
 

src/lib.rs

@@ -14,6 +14,9 @@ use tokio::{
 mod ffi;
 use crate::ffi::CryptoInfo;
 
+mod async_read_ready;
+pub use async_read_ready::AsyncReadReady;
+
 mod ktls_stream;
 pub use ktls_stream::KtlsStream;
 
@@ -203,7 +206,7 @@ pub async fn config_ktls_server<IO>(
     mut stream: tokio_rustls::server::TlsStream<CorkStream<IO>>,
 ) -> Result<KtlsStream<IO>, Error>
 where
-    IO: AsRawFd + AsyncRead + AsyncWrite + Unpin,
+    IO: AsRawFd + AsyncRead + AsyncReadReady + AsyncWrite + Unpin,
 {
     stream.get_mut().0.corked = true;
     let drained = drain(&mut stream).await.map_err(Error::DrainError)?;