Use recvmsg to handle/ignore control messages when poll_read fails

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,16 +65,16 @@ 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();
+        let mut this = self.project();
 
         if let Some((drain_index, drained)) = this.drained.as_mut() {
             let drained = &drained[*drain_index..];
@@ -79,11 +90,95 @@ 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 read_res = this.inner.as_mut().poll_read(cx, buf);
+        if let task::Poll::Ready(Err(e)) = &read_res {
+            // 5 is a generic "input/output error", it happens when
+            // using poll_read on a kTLS socket that just received
+            // a control message
+            if let Some(5) = e.raw_os_error() {
+                // could be a control message, let's check
+                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) => {
+                        unreachable!("expected a control message, got EAGAIN")
+                    }
+                    Err(e) => {
+                        // ok I guess it really failed then
+                        tracing::trace!(?e, "recvmsg failed");
+                        return Err(e.into()).into();
+                    }
+                };
+                let cmsg = r
+                    .cmsgs()
+                    .next()
+                    .expect("we should've received exactly one control message");
+                match cmsg {
+                    // cf. RFC 5246, Section 6.2.1
+                    // https://datatracker.ietf.org/doc/html/rfc5246#section-6.2.1
+                    ControlMessageOwned::TlsGetRecordType(t) => {
+                        match t {
+                            // change_cipher_spec
+                            20 => {
+                                panic!(
+                                    "received TLS change_cipher_spec, this isn't supported by ktls"
+                                )
+                            }
+                            // alert
+                            21 => {
+                                panic!("received TLS alert, this isn't supported by ktls")
+                            }
+                            // handshake
+                            22 => {
+                                // TODO: this is where we receive TLS 1.3 resumption tickets,
+                                // should those be stored anywhere? I'm not even sure what
+                                // format they have at this point
+                                tracing::trace!(
+                                    "ignoring handshake message (probably a resumption ticket)"
+                                );
+                            }
+                            // application data
+                            23 => {
+                                unreachable!("received TLS application in recvmsg, this is supposed to happen in the poll_read codepath")
+                            }
+                            _ => {
+                                // just ignore the message type then
+                                tracing::trace!("received message_type {t:#?}");
+                            }
+                        }
+                    }
+                    _ => panic!("unexpected cmsg type: {cmsg:#?}"),
+                };
+
+                // FIXME: this is hacky, but can we do better?
+                // after we handled (..ignored) the control message, we don't
+                // know whether the scoket is still ready to be read or not.
+                //
+                // we could try looping (tricky code structure), but we can't,
+                // for example, just call `poll_read`, which might fail not
+                // not with EAGAIN/EWOULDBLOCK, but because _another_ control
+                // message is available.
+                cx.waker().wake_by_ref();
+                return task::Poll::Pending;
+            }
+        }
+
+        read_res
     }
 }
 

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)?;