feat(transport): Fix TLS accept w/ peer certs (#535)

* feat(transport): Fix TLS accept w/ peer certs

* fix unused var

* fix feature flag imports

* spawn accept task

examples/src/tls_client_auth/server.rs

@@ -17,9 +17,11 @@ pub struct EchoServer;
 #[tonic::async_trait]
 impl pb::echo_server::Echo for EchoServer {
     async fn unary_echo(&self, request: Request<EchoRequest>) -> EchoResult<EchoResponse> {
-        if let Some(certs) = request.peer_certs() {
+        let certs = request
-            println!("Got {} peer certs!", certs.len());
+            .peer_certs()
-        }
+            .expect("Client did not send its certs!");
+
+        println!("Got {} peer certs!", certs.len());
 
         let message = request.into_inner().message;
         Ok(Response::new(EchoResponse { message }))

tonic/Cargo.toml

@@ -31,6 +31,7 @@ transport = [
     "tokio",
     "tower",
     "tracing-futures",
+    "tokio/macros"
 ]
 tls = ["transport", "tokio-rustls"]
 tls-roots = ["tls", "rustls-native-certs"]

tonic/src/transport/server/conn.rs

@@ -1,11 +1,9 @@
-#[cfg(feature = "tls")]
-use super::TlsStream;
 use crate::transport::Certificate;
 use hyper::server::conn::AddrStream;
 use std::net::SocketAddr;
 use tokio::net::TcpStream;
 #[cfg(feature = "tls")]
-use tokio_rustls::rustls::Session;
+use tokio_rustls::{rustls::Session, server::TlsStream};
 
 /// Trait that connected IO resources implement.
 ///
@@ -39,24 +37,20 @@ impl Connected for TcpStream {
 #[cfg(feature = "tls")]
 impl<T: Connected> Connected for TlsStream<T> {
     fn remote_addr(&self) -> Option<SocketAddr> {
-        if let Some((inner, _)) = self.get_ref() {
+        let (inner, _) = self.get_ref();
-            inner.remote_addr()
+
-        } else {
+        inner.remote_addr()
-            None
-        }
     }
 
     fn peer_certs(&self) -> Option<Vec<Certificate>> {
-        if let Some((_, session)) = self.get_ref() {
+        let (_, session) = self.get_ref();
-            if let Some(certs) = session.get_peer_certificates() {
+
-                let certs = certs
+        if let Some(certs) = session.get_peer_certificates() {
-                    .into_iter()
+            let certs = certs
-                    .map(|c| Certificate::from_pem(c.0))
+                .into_iter()
-                    .collect();
+                .map(|c| Certificate::from_pem(c.0))
-                Some(certs)
+                .collect();
-            } else {
+            Some(certs)
-                None
-            }
         } else {
             None
         }

tonic/src/transport/server/incoming.rs

@@ -14,7 +14,27 @@ use std::{
 };
 use tokio::io::{AsyncRead, AsyncWrite};
 
-#[cfg_attr(not(feature = "tls"), allow(unused_variables))]
+#[cfg(not(feature = "tls"))]
+pub(crate) fn tcp_incoming<IO, IE>(
+    incoming: impl Stream<Item = Result<IO, IE>>,
+    _server: Server,
+) -> impl Stream<Item = Result<ServerIo, crate::Error>>
+where
+    IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
+    IE: Into<crate::Error>,
+{
+    async_stream::try_stream! {
+        futures_util::pin_mut!(incoming);
+
+
+        while let Some(stream) = incoming.try_next().await? {
+
+            yield ServerIo::new(stream);
+        }
+    }
+}
+
+#[cfg(feature = "tls")]
 pub(crate) fn tcp_incoming<IO, IE>(
     incoming: impl Stream<Item = Result<IO, IE>>,
     server: Server,
@@ -26,21 +46,89 @@ where
     async_stream::try_stream! {
         futures_util::pin_mut!(incoming);
 
-        while let Some(stream) = incoming.try_next().await? {
+        #[cfg(feature = "tls")]
-            #[cfg(feature = "tls")]
+        let mut tasks = futures_util::stream::futures_unordered::FuturesUnordered::new();
-            {
+
-                if let Some(tls) = &server.tls {
+        loop {
-                    let io = tls.accept(stream);
+            match select(&mut incoming, &mut tasks).await {
-                    yield ServerIo::new(io);
+                SelectOutput::Incoming(stream) => {
-                    continue;
+                    if let Some(tls) = &server.tls {
+                        let tls = tls.clone();
+
+                        let accept = tokio::spawn(async move {
+                            let io = tls.accept(stream).await?;
+                            Ok(ServerIo::new(io))
+                        });
+
+                        tasks.push(accept);
+                    } else {
+                        yield ServerIo::new(stream);
+                    }
+                }
+
+                SelectOutput::Io(io) => {
+                    yield io;
+                }
+
+                SelectOutput::Err(e) => {
+                    tracing::error!(message = "Accept loop error.", error = %e);
+                }
+
+                SelectOutput::Done => {
+                    break;
                 }
             }
-
-            yield ServerIo::new(stream);
         }
     }
 }
 
+#[cfg(feature = "tls")]
+async fn select<IO, IE>(
+    incoming: &mut (impl Stream<Item = Result<IO, IE>> + Unpin),
+    tasks: &mut futures_util::stream::futures_unordered::FuturesUnordered<
+        tokio::task::JoinHandle<Result<ServerIo, crate::Error>>,
+    >,
+) -> SelectOutput<IO>
+where
+    IE: Into<crate::Error>,
+{
+    use futures_util::StreamExt;
+
+    if tasks.is_empty() {
+        return match incoming.try_next().await {
+            Ok(Some(stream)) => SelectOutput::Incoming(stream),
+            Ok(None) => SelectOutput::Done,
+            Err(e) => SelectOutput::Err(e.into()),
+        };
+    }
+
+    tokio::select! {
+        stream = incoming.try_next() => {
+            match stream {
+                Ok(Some(stream)) => SelectOutput::Incoming(stream),
+                Ok(None) => SelectOutput::Done,
+                Err(e) => SelectOutput::Err(e.into()),
+            }
+        }
+
+        accept = tasks.next() => {
+            match accept.expect("FuturesUnordered stream should never end") {
+                Ok(Ok(io)) => SelectOutput::Io(io),
+                Ok(Err(e)) => SelectOutput::Err(e),
+                Err(e) => SelectOutput::Err(e.into()),
+            }
+        }
+    }
+}
+
+#[cfg(feature = "tls")]
+enum SelectOutput<A> {
+    Incoming(A),
+    Io(ServerIo),
+    Err(crate::Error),
+    Done,
+}
+
 pub(crate) struct TcpIncoming {
     inner: AddrIncoming,
 }
@@ -65,106 +153,3 @@ impl Stream for TcpIncoming {
         Pin::new(&mut self.inner).poll_accept(cx)
     }
 }
-
-// tokio_rustls::server::TlsStream doesn't expose constructor methods,
-// so we have to TlsAcceptor::accept and handshake to have access to it
-// TlsStream implements AsyncRead/AsyncWrite handshaking tokio_rustls::Accept first
-#[cfg(feature = "tls")]
-pub(crate) struct TlsStream<IO> {
-    state: State<IO>,
-}
-
-#[cfg(feature = "tls")]
-enum State<IO> {
-    Handshaking(tokio_rustls::Accept<IO>),
-    Streaming(tokio_rustls::server::TlsStream<IO>),
-}
-
-#[cfg(feature = "tls")]
-impl<IO> TlsStream<IO> {
-    pub(crate) fn new(accept: tokio_rustls::Accept<IO>) -> Self {
-        TlsStream {
-            state: State::Handshaking(accept),
-        }
-    }
-
-    pub(crate) fn get_ref(&self) -> Option<(&IO, &tokio_rustls::rustls::ServerSession)> {
-        if let State::Streaming(tls) = &self.state {
-            Some(tls.get_ref())
-        } else {
-            None
-        }
-    }
-}
-
-#[cfg(feature = "tls")]
-impl<IO> AsyncRead for TlsStream<IO>
-where
-    IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
-{
-    fn poll_read(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-        buf: &mut tokio::io::ReadBuf<'_>,
-    ) -> Poll<std::io::Result<()>> {
-        use std::future::Future;
-
-        let pin = self.get_mut();
-        match pin.state {
-            State::Handshaking(ref mut accept) => {
-                match futures_core::ready!(Pin::new(accept).poll(cx)) {
-                    Ok(mut stream) => {
-                        let result = Pin::new(&mut stream).poll_read(cx, buf);
-                        pin.state = State::Streaming(stream);
-                        result
-                    }
-                    Err(err) => Poll::Ready(Err(err)),
-                }
-            }
-            State::Streaming(ref mut stream) => Pin::new(stream).poll_read(cx, buf),
-        }
-    }
-}
-
-#[cfg(feature = "tls")]
-impl<IO> AsyncWrite for TlsStream<IO>
-where
-    IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
-{
-    fn poll_write(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-        buf: &[u8],
-    ) -> Poll<std::io::Result<usize>> {
-        use std::future::Future;
-
-        let pin = self.get_mut();
-        match pin.state {
-            State::Handshaking(ref mut accept) => {
-                match futures_core::ready!(Pin::new(accept).poll(cx)) {
-                    Ok(mut stream) => {
-                        let result = Pin::new(&mut stream).poll_write(cx, buf);
-                        pin.state = State::Streaming(stream);
-                        result
-                    }
-                    Err(err) => Poll::Ready(Err(err)),
-                }
-            }
-            State::Streaming(ref mut stream) => Pin::new(stream).poll_write(cx, buf),
-        }
-    }
-
-    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        match self.state {
-            State::Handshaking(_) => Poll::Ready(Ok(())),
-            State::Streaming(ref mut stream) => Pin::new(stream).poll_flush(cx),
-        }
-    }
-
-    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        match self.state {
-            State::Handshaking(_) => Poll::Ready(Ok(())),
-            State::Streaming(ref mut stream) => Pin::new(stream).poll_shutdown(cx),
-        }
-    }
-}

tonic/src/transport/server/mod.rs

@@ -16,7 +16,7 @@ use super::service::TlsAcceptor;
 use incoming::TcpIncoming;
 
 #[cfg(feature = "tls")]
-pub(crate) use incoming::TlsStream;
+pub(crate) use tokio_rustls::server::TlsStream;
 
 #[cfg(feature = "tls")]
 use crate::transport::Error;

tonic/src/transport/service/tls.rs

@@ -162,14 +162,12 @@ impl TlsAcceptor {
         })
     }
 
-    pub(crate) fn accept<IO>(&self, io: IO) -> TlsStream<IO>
+    pub(crate) async fn accept<IO>(&self, io: IO) -> Result<TlsStream<IO>, crate::Error>
     where
         IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
     {
         let acceptor = RustlsAcceptor::from(self.inner.clone());
-        let accept = acceptor.accept(io);
+        acceptor.accept(io).await.map_err(Into::into)
-
-        TlsStream::new(accept)
     }
 }