rust: bssl-tls-tokio: Hyper I/O traits

Signed-off-by: Xiangfei Ding <xfding@google.com>
Change-Id: Ib35fe2e377ddbbed2837d480e44066906a6a6964
Reviewed-on: https://boringssl-review.googlesource.com/c/boringssl/+/100687
Reviewed-by: Adam Langley <agl@google.com>
diff --git a/rust/bssl-tls-tokio/Cargo.toml b/rust/bssl-tls-tokio/Cargo.toml
index 33b03bc..a6eaefd 100644
--- a/rust/bssl-tls-tokio/Cargo.toml
+++ b/rust/bssl-tls-tokio/Cargo.toml
@@ -17,7 +17,6 @@
 [dependencies.hyper]
 version = "1.0"
 optional = true
-features = ["client", "http1", "http2"]
 
 [dependencies.tower]
 version = "0.4"
@@ -30,9 +29,15 @@
 tokio = { version = "1.0", features = ["full"] }
 futures = "0.3"
 bssl-x509 = { path = "../bssl-x509" }
-hyper = { version = "1.0", features = ["client", "server", "http1", "http2"] }
-hyper-util = { version = "0.1", features = ["tokio"] }
+
+[dev-dependencies.hyper]
+version = "1.0"
+features = ["client", "server", "http1", "http2"]
+
+[dev-dependencies.hyper-util]
+version = "0.1"
+features = ["tokio"]
 
 [features]
-# `hyper` enables interop with `hyper` 1.0 and `tower`
+# `hyper` enables interop with `hyper` and `tower`
 hyper = ["dep:hyper", "dep:tower"]
diff --git a/rust/bssl-tls-tokio/src/hyper.rs b/rust/bssl-tls-tokio/src/hyper.rs
index 850d155..ec1add0 100644
--- a/rust/bssl-tls-tokio/src/hyper.rs
+++ b/rust/bssl-tls-tokio/src/hyper.rs
@@ -14,21 +14,52 @@
 
 //! Hyper support
 
-use crate::TlsConnector;
+use crate::{
+    TlsAcceptor,
+    TlsConnector,
+    TlsStream,
+    TokioTlsConnection,
+    translate_stdio_err, //
+};
 
 use std::{
     error::Error,
     fmt::Debug,
     future::Future,
+    io,
     pin::Pin,
     sync::Arc,
     task::{
         Context,
         Poll, //
-    }, //
+    },
 };
 
-use hyper::http;
+use bssl_tls::{
+    ReceiveBuffer,
+    connection::{
+        Client,
+        Server,
+        lifecycle::ShutdownStatus, //
+    },
+    errors::Error as TlsError,
+    io::{
+        AbstractReader,
+        AbstractSocket,
+        AbstractSocketResult,
+        AbstractWriter,
+        IoStatus,
+        NoAsyncContext, //
+    },
+};
+use hyper::{
+    http,
+    rt::{
+        Read,
+        ReadBufCursor,
+        Write, //
+    }, //
+};
 use tower::Service;
 
 /// A connector for `hyper` using `bssl-tls`.
@@ -59,11 +90,11 @@
 impl<Inner> Service<http::Uri> for HyperBsslConnector<Inner>
 where
     Inner: Service<http::Uri>,
-    Inner::Response: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + Sync + 'static,
+    Inner::Response: Read + Write + Unpin + Send + Sync + 'static,
     Inner::Future: Send + 'static,
     Inner::Error: Into<Box<dyn Error + Send + Sync>>,
 {
-    type Response = crate::TlsStream<bssl_tls::connection::Client, Inner::Response>;
+    type Response = TlsStream<Client, Inner::Response>;
     type Error = Box<dyn Error + Send + Sync>;
     type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
 
@@ -86,7 +117,243 @@
 
         Box::pin(async move {
             let stream = fut.await.map_err(Into::into)?;
-            Ok(connector.connect(&domain, stream).await?)
+            Ok(connector.hyper_connect(&domain, stream).await?)
         })
     }
 }
+
+impl TlsConnector {
+    /// Connect to the given domain using the provided stream implementing Hyper I/O traits.
+    pub async fn hyper_connect<S>(
+        &self,
+        domain: &str,
+        stream: S,
+    ) -> Result<TlsStream<Client, S>, TlsError>
+    where
+        S: Read + Write + Unpin + Send + 'static,
+    {
+        let mut conn = self.ctx.new_client_connection().build();
+        conn.in_handshake()
+            .expect("we have not started handshake")
+            .set_tlsext_host_name(domain)?
+            .set_host(domain)?;
+        conn.set_io(HyperIo(stream))?;
+        conn.async_handshake().await?;
+
+        Ok(TlsStream::new(TokioTlsConnection::new(conn)))
+    }
+}
+
+impl TlsAcceptor {
+    /// Accept a new connection using the provided stream implementing Hyper I/O traits.
+    pub async fn hyper_accept<S>(&self, stream: S) -> Result<TlsStream<Server, S>, TlsError>
+    where
+        S: Read + Write + Unpin + Send + 'static,
+    {
+        let mut conn = self.ctx.new_server_connection().build();
+        conn.set_io(HyperIo(stream))?;
+        conn.async_handshake().await?;
+
+        Ok(TlsStream::new(TokioTlsConnection::new(conn)))
+    }
+}
+
+/// IO object implementing [`hyper::rt::Read`] or [`hyper::rt::Write`] protocols.
+pub struct HyperIo<T>(pub T);
+
+fn hyper_async_read<T: Read>(
+    mut this: Pin<&mut T>,
+    ctx: &mut Context<'_>,
+    buffer: &mut [u8],
+) -> AbstractSocketResult {
+    let buffer_len = buffer.len();
+    let mut buf = hyper::rt::ReadBuf::new(buffer);
+    loop {
+        return match this.as_mut().poll_read(ctx, buf.unfilled()) {
+            Poll::Ready(Ok(())) => {
+                if buf.filled().is_empty() && buffer_len > 0 {
+                    AbstractSocketResult::EndOfStream
+                } else {
+                    AbstractSocketResult::Ok(buf.filled().len())
+                }
+            }
+            Poll::Pending => AbstractSocketResult::Retry,
+            Poll::Ready(Err(e)) => {
+                if e.kind() == io::ErrorKind::Interrupted {
+                    continue;
+                }
+                translate_stdio_err(e)
+            }
+        };
+    }
+}
+
+fn hyper_async_write<T: Write>(
+    mut this: Pin<&mut T>,
+    ctx: &mut Context<'_>,
+    buffer: &[u8],
+) -> AbstractSocketResult {
+    loop {
+        return match this.as_mut().poll_write(ctx, buffer) {
+            Poll::Ready(Ok(bytes)) => {
+                if buffer.is_empty() {
+                    AbstractSocketResult::Ok(0)
+                } else if bytes == 0 {
+                    AbstractSocketResult::EndOfStream
+                } else {
+                    AbstractSocketResult::Ok(bytes)
+                }
+            }
+            Poll::Pending => AbstractSocketResult::Retry,
+            Poll::Ready(Err(e)) => {
+                if e.kind() == io::ErrorKind::Interrupted {
+                    continue;
+                }
+                translate_stdio_err(e)
+            }
+        };
+    }
+}
+
+fn hyper_async_flush<T: Write>(
+    mut this: Pin<&mut T>,
+    ctx: &mut Context<'_>,
+) -> AbstractSocketResult {
+    loop {
+        return match this.as_mut().poll_flush(ctx) {
+            Poll::Ready(Ok(())) => AbstractSocketResult::Ok(0),
+            Poll::Pending => AbstractSocketResult::Retry,
+            Poll::Ready(Err(e)) => {
+                if e.kind() == io::ErrorKind::Interrupted {
+                    continue;
+                }
+                translate_stdio_err(e)
+            }
+        };
+    }
+}
+
+impl<T: Read + Send + Unpin> AbstractReader for HyperIo<T> {
+    fn read(
+        &mut self,
+        async_ctx: Option<&mut Context<'_>>,
+        buffer: &mut [u8],
+    ) -> AbstractSocketResult {
+        let Some(ctx) = async_ctx else {
+            return AbstractSocketResult::Err(Box::new(NoAsyncContext));
+        };
+        hyper_async_read(Pin::new(&mut self.0), ctx, buffer)
+    }
+}
+
+impl<T: Write + Send + Unpin> AbstractWriter for HyperIo<T> {
+    fn write(
+        &mut self,
+        async_ctx: Option<&mut Context<'_>>,
+        buffer: &[u8],
+    ) -> AbstractSocketResult {
+        let Some(ctx) = async_ctx else {
+            return AbstractSocketResult::Err(Box::new(NoAsyncContext));
+        };
+        hyper_async_write(Pin::new(&mut self.0), ctx, buffer)
+    }
+
+    fn flush(&mut self, async_ctx: Option<&mut Context<'_>>) -> AbstractSocketResult {
+        let Some(ctx) = async_ctx else {
+            return AbstractSocketResult::Err(Box::new(NoAsyncContext));
+        };
+        hyper_async_flush(Pin::new(&mut self.0), ctx)
+    }
+}
+
+impl<T: Read + Write + Send + Unpin> AbstractSocket for HyperIo<T> {}
+
+impl<Role, S: Unpin> Read for TlsStream<Role, S> {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        mut buf: ReadBufCursor<'_>,
+    ) -> Poll<Result<(), io::Error>> {
+        let mut read_buf = ReceiveBuffer::new_uninit(unsafe {
+            // We will not uninitialize anything outside this buffer.
+            buf.as_mut()
+        });
+        let status = match self
+            .conn
+            .inner
+            .as_pin_mut()
+            .async_poll_read(&mut read_buf, cx)
+        {
+            Ok(Some(status)) => status,
+            Ok(None) => return Poll::Pending,
+            Err(e) => return Poll::Ready(Err(io::Error::other(e))),
+        };
+        match status {
+            IoStatus::Ok(bytes) => {
+                debug_assert_eq!(bytes, read_buf.written());
+                unsafe {
+                    // Safety: BoringSSL has successfully written and initialized `bytes` in the buffer.
+                    buf.advance(bytes);
+                }
+                Poll::Ready(Ok(()))
+            }
+            IoStatus::EndOfStream => Poll::Ready(Ok(())),
+            _ => Poll::Ready(Err(io::Error::other("Unexpected I/O status"))),
+        }
+    }
+}
+
+impl<Role, S: Unpin> Write for TlsStream<Role, S> {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, io::Error>> {
+        let status = match self.conn.inner.as_pin_mut().async_poll_write(buf, cx) {
+            Ok(Some(status)) => status,
+            Ok(None) => return Poll::Pending,
+            Err(e) => return Poll::Ready(Err(io::Error::other(e))),
+        };
+        match status {
+            IoStatus::Ok(bytes) => Poll::Ready(Ok(bytes)),
+            IoStatus::EndOfStream => Poll::Ready(Ok(0)),
+            _ => Poll::Ready(Err(io::Error::other("Unexpected I/O status"))),
+        }
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        let status = match self.conn.inner.as_pin_mut().async_poll_flush(cx) {
+            Ok(Some(status)) => status,
+            Ok(None) => return Poll::Pending,
+            Err(e) => return Poll::Ready(Err(io::Error::other(e))),
+        };
+        match status {
+            IoStatus::Ok(_) => Poll::Ready(Ok(())),
+            IoStatus::EndOfStream => Poll::Ready(Ok(())),
+            _ => Poll::Ready(Err(io::Error::other("Unexpected I/O status"))),
+        }
+    }
+
+    fn poll_shutdown(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), io::Error>> {
+        match self.conn.inner.as_pin_mut().async_poll_shutdown(cx) {
+            Ok(Some(ShutdownStatus::CloseNotifyReceived)) => Poll::Ready(Ok(())),
+            Ok(Some(ShutdownStatus::RemainingApplicationData)) => {
+                Poll::Ready(Err(io::Error::other(
+                    "caller needs to drain application data before polling on shutdown again",
+                )))
+            }
+            Ok(Some(ShutdownStatus::EndOfStream)) => Poll::Ready(Err(io::Error::new(
+                io::ErrorKind::UnexpectedEof,
+                "unexpected eof while waiting for peek close_notify",
+            ))),
+            Ok(Some(ShutdownStatus::CloseNotifyPosted)) => {
+                unreachable!()
+            }
+            Ok(None) => Poll::Pending,
+            Err(e) => Poll::Ready(Err(io::Error::other(e))),
+        }
+    }
+}
diff --git a/rust/bssl-tls-tokio/src/lib.rs b/rust/bssl-tls-tokio/src/lib.rs
index f8065a7..0041951 100644
--- a/rust/bssl-tls-tokio/src/lib.rs
+++ b/rust/bssl-tls-tokio/src/lib.rs
@@ -119,6 +119,7 @@
         DerefMut, //
     },
     pin::Pin,
+    str::FromStr,
     task::{
         Context,
         Poll,
@@ -164,7 +165,7 @@
 mod tests;
 
 /// Translates a `std::io::Error` into an `AbstractSocketResult`.
-fn translate_stdio_err(err: io::Error) -> AbstractSocketResult {
+pub(crate) fn translate_stdio_err(err: io::Error) -> AbstractSocketResult {
     match err.kind() {
         io::ErrorKind::WouldBlock => AbstractSocketResult::Retry,
         io::ErrorKind::ConnectionReset
@@ -407,7 +408,7 @@
 
 /// A wrapper around [`TlsConnection`] that implements Tokio's async I/O traits.
 pub struct TokioTlsConnection<Role> {
-    inner: TlsConnection<Role, TlsMode>,
+    pub(crate) inner: TlsConnection<Role, TlsMode>,
 }
 
 impl<Role> TokioTlsConnection<Role> {
@@ -449,7 +450,7 @@
         let status = match self.inner.as_pin_mut().async_poll_read(&mut recv_buf, cx) {
             Ok(Some(status)) => status,
             Ok(None) => return Poll::Pending,
-            Err(e) => return Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, e))),
+            Err(e) => return Poll::Ready(Err(io::Error::other(e))),
         };
         match status {
             IoStatus::Ok(bytes) => {
@@ -461,10 +462,7 @@
                 Poll::Ready(Ok(()))
             }
             IoStatus::EndOfStream => Poll::Ready(Ok(())),
-            _ => Poll::Ready(Err(io::Error::new(
-                io::ErrorKind::Other,
-                "Unexpected I/O status",
-            ))),
+            _ => Poll::Ready(Err(io::Error::other("Unexpected I/O status"))),
         }
     }
 }
@@ -478,15 +476,12 @@
         let status = match self.inner.as_pin_mut().async_poll_write(buf, cx) {
             Ok(Some(status)) => status,
             Ok(None) => return Poll::Pending,
-            Err(e) => return Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, e))),
+            Err(e) => return Poll::Ready(Err(io::Error::other(e))),
         };
         match status {
             IoStatus::Ok(bytes) => Poll::Ready(Ok(bytes)),
             IoStatus::EndOfStream => Poll::Ready(Ok(0)),
-            _ => Poll::Ready(Err(io::Error::new(
-                io::ErrorKind::Other,
-                "Unexpected I/O status",
-            ))),
+            _ => Poll::Ready(Err(io::Error::other("Unexpected I/O status"))),
         }
     }
 
@@ -494,25 +489,23 @@
         let status = match self.inner.as_pin_mut().async_poll_flush(cx) {
             Ok(Some(status)) => status,
             Ok(None) => return Poll::Pending,
-            Err(e) => return Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, e))),
+            Err(e) => return Poll::Ready(Err(io::Error::other(e))),
         };
         match status {
             IoStatus::Ok(_) => Poll::Ready(Ok(())),
             IoStatus::EndOfStream => Poll::Ready(Ok(())),
-            _ => Poll::Ready(Err(io::Error::new(
-                io::ErrorKind::Other,
-                "Unexpected I/O status",
-            ))),
+            _ => Poll::Ready(Err(io::Error::other("Unexpected I/O status"))),
         }
     }
 
     fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
         match self.inner.as_pin_mut().async_poll_shutdown(cx) {
             Ok(Some(ShutdownStatus::CloseNotifyReceived)) => Poll::Ready(Ok(())),
-            Ok(Some(ShutdownStatus::RemainingApplicationData)) => Poll::Ready(Err(io::Error::new(
-                io::ErrorKind::Other,
-                "caller needs to drain application data before polling on shutdown again",
-            ))),
+            Ok(Some(ShutdownStatus::RemainingApplicationData)) => {
+                Poll::Ready(Err(io::Error::other(
+                    "caller needs to drain application data before polling on shutdown again",
+                )))
+            }
             Ok(Some(ShutdownStatus::EndOfStream)) => Poll::Ready(Err(io::Error::new(
                 io::ErrorKind::UnexpectedEof,
                 "unexpected eof while waiting for peek close_notify",
@@ -521,7 +514,7 @@
                 unreachable!()
             }
             Ok(None) => Poll::Pending,
-            Err(e) => Poll::Ready(Err(io::Error::new(io::ErrorKind::Other, e))),
+            Err(e) => Poll::Ready(Err(io::Error::other(e))),
         }
     }
 }
@@ -543,16 +536,15 @@
         S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
     {
         let mut conn = self.ctx.new_client_connection().build();
-        conn.in_handshake().unwrap().set_host(domain)?;
-
+        let mut in_handshake = conn.in_handshake().expect("we are handshaking");
+        in_handshake.set_host(domain)?;
+        if std::net::IpAddr::from_str(domain).is_err() {
+            in_handshake.set_tlsext_host_name(domain)?;
+        }
         conn.set_io(TokioIo(stream))?;
-
         conn.async_handshake().await?;
 
-        Ok(TlsStream {
-            conn: TokioTlsConnection::new(conn),
-            _marker: PhantomData,
-        })
+        Ok(TlsStream::new(TokioTlsConnection::new(conn)))
     }
 }
 
@@ -573,25 +565,27 @@
         S: AsyncRead + AsyncWrite + Send + Unpin + 'static,
     {
         let mut conn = self.ctx.new_server_connection().build();
-
         conn.set_io(TokioIo(stream))?;
-
         conn.async_handshake().await?;
 
-        Ok(TlsStream {
-            conn: TokioTlsConnection::new(conn),
-            _marker: PhantomData,
-        })
+        Ok(TlsStream::new(TokioTlsConnection::new(conn)))
     }
 }
 
 /// A TLS stream driven by Tokio I/O.
 pub struct TlsStream<Role, Stream> {
-    conn: TokioTlsConnection<Role>,
-    _marker: PhantomData<Stream>,
+    pub(crate) conn: TokioTlsConnection<Role>,
+    pub(crate) _marker: PhantomData<Stream>,
 }
 
 impl<Role, S> TlsStream<Role, S> {
+    pub(crate) fn new(conn: TokioTlsConnection<Role>) -> Self {
+        Self {
+            conn,
+            _marker: PhantomData,
+        }
+    }
+
     /// Get a reference to the underlying `TlsConnection`.
     pub fn get_ref(&self) -> &TlsConnection<Role> {
         &self.conn
diff --git a/rust/bssl-tls-tokio/tests/hyper.rs b/rust/bssl-tls-tokio/tests/hyper.rs
index 4c42c43..599d349 100644
--- a/rust/bssl-tls-tokio/tests/hyper.rs
+++ b/rust/bssl-tls-tokio/tests/hyper.rs
@@ -97,7 +97,7 @@
 }
 
 impl Service<hyper::http::Uri> for MockTcpConnector {
-    type Response = TcpStream;
+    type Response = HyperTokioIo<TcpStream>;
     type Error = std::io::Error;
     type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
 
@@ -107,7 +107,10 @@
 
     fn call(&mut self, _: hyper::http::Uri) -> Self::Future {
         let addr = self.addr;
-        Box::pin(TcpStream::connect(addr))
+        Box::pin(async move {
+            let stream = TcpStream::connect(addr).await?;
+            Ok(HyperTokioIo::new(stream))
+        })
     }
 }
 
@@ -134,7 +137,7 @@
 
         hyper::server::conn::http2::Builder::new(hyper_util::rt::TokioExecutor::new())
             .serve_connection(
-                HyperTokioIo::new(tls_stream),
+                tls_stream,
                 service_fn(|_req| async {
                     Ok::<_, hyper::Error>(hyper::Response::new(SimpleBody::new(
                         "hello from h2 server",
@@ -160,12 +163,10 @@
             .await
             .unwrap();
 
-        let (mut sender, conn) = hyper::client::conn::http2::handshake(
-            hyper_util::rt::TokioExecutor::new(),
-            HyperTokioIo::new(tls_stream),
-        )
-        .await
-        .unwrap();
+        let (mut sender, conn) =
+            hyper::client::conn::http2::handshake(hyper_util::rt::TokioExecutor::new(), tls_stream)
+                .await
+                .unwrap();
 
         // Drive the connection in the background.
         tokio::spawn(async move {
diff --git a/rust/bssl-tls/src/connection/io.rs b/rust/bssl-tls/src/connection/io.rs
index 070b55b..16fa610 100644
--- a/rust/bssl-tls/src/connection/io.rs
+++ b/rust/bssl-tls/src/connection/io.rs
@@ -26,11 +26,10 @@
 };
 
 use crate::{
-    ReceiveBuffer,
     connection::{
-        TlsConnection,
         lifecycle::ShutdownStatus,
         methods::HasTlsConnectionMethod, //
+        TlsConnection,
     },
     context::{
         HasDatagramIo,
@@ -44,6 +43,7 @@
     },
     ffi::slice_into_ffi_raw_parts,
     io::IoStatus, //
+    ReceiveBuffer,
 };
 
 impl<R, M> TlsConnection<R, M>