add actix stream trait

actix · robjtede · Feb 20, 2021 · Feb 8, 2021 · Feb 8, 2021 · Feb 20, 2021
commit 350fbd053aab943c036ebc4eb6efd6de0703017e
diff --git a/actix-rt/src/lib.rs b/actix-rt/src/lib.rs
@@ -77,6 +77,44 @@ pub mod net {
 
     #[cfg(unix)]
     pub use tokio::net::{UnixDatagram, UnixListener, UnixStream};
+
+    use std::task::{Context, Poll};
+
+    use tokio::io::{AsyncRead, AsyncWrite};
+
+    /// Trait for generic over tokio stream types and various wrapper types around them.
+    pub trait ActixStream: AsyncRead + AsyncWrite + Unpin + 'static {
+        /// poll stream and check read readiness of Self.
+        ///
+        /// See [tokio::net::TcpStream::poll_read_ready] for detail
+        fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<std::io::Result<()>>;
+
+        /// poll stream and check write readiness of Self.
+        ///
+        /// See [tokio::net::TcpStream::poll_write_ready] for detail
+        fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<std::io::Result<()>>;
+    }
+
+    impl ActixStream for TcpStream {
+        fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+            TcpStream::poll_read_ready(self, cx)
+        }
+
+        fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+            TcpStream::poll_write_ready(self, cx)
+        }
+    }
+
+    #[cfg(unix)]
+    impl ActixStream for UnixStream {
+        fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+            UnixStream::poll_read_ready(self, cx)
+        }
+
+        fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+            UnixStream::poll_write_ready(self, cx)
+        }
+    }
 }
 
 pub mod time {

diff --git a/actix-tls/examples/basic.rs b/actix-tls/examples/basic.rs
@@ -29,9 +29,10 @@ use std::{
     },
 };
 
+use actix_rt::net::TcpStream;
 use actix_server::Server;
 use actix_service::pipeline_factory;
-use actix_tls::accept::rustls::Acceptor as RustlsAcceptor;
+use actix_tls::accept::rustls::{Acceptor as RustlsAcceptor, TlsStream};
 use futures_util::future::ok;
 use log::info;
 use rustls::{
@@ -74,9 +75,9 @@ async fn main() -> io::Result<()> {
             // Set up TLS service factory
             pipeline_factory(tls_acceptor.clone())
                 .map_err(|err| println!("Rustls error: {:?}", err))
-                .and_then(move |stream| {
+                .and_then(move |stream: TlsStream<TcpStream>| {
                     let num = count.fetch_add(1, Ordering::Relaxed);
-                    info!("[{}] Got TLS connection: {:?}", num, stream);
+                    info!("[{}] Got TLS connection: {:?}", num, &*stream);
                     ok(())
                 })
         })?

diff --git a/actix-tls/src/accept/nativetls.rs b/actix-tls/src/accept/nativetls.rs
@@ -1,14 +1,92 @@
+use std::io;
+use std::ops::{Deref, DerefMut};
 use std::task::{Context, Poll};
 
-use actix_codec::{AsyncRead, AsyncWrite};
+use actix_rt::net::ActixStream;
 use actix_service::{Service, ServiceFactory};
 use actix_utils::counter::Counter;
 use futures_core::future::LocalBoxFuture;
 
 pub use tokio_native_tls::native_tls::Error;
-pub use tokio_native_tls::{TlsAcceptor, TlsStream};
+pub use tokio_native_tls::TlsAcceptor;
 
 use super::MAX_CONN_COUNTER;
+use actix_codec::{AsyncRead, AsyncWrite, ReadBuf};
+use std::io::IoSlice;
+use std::pin::Pin;
+
+/// wrapper type for `tokio_native_tls::TlsStream` in order to impl `ActixStream` trait.
+pub struct TlsStream<T>(tokio_native_tls::TlsStream<T>);
+
+impl<T> From<tokio_native_tls::TlsStream<T>> for TlsStream<T> {
+    fn from(stream: tokio_native_tls::TlsStream<T>) -> Self {
+        Self(stream)
+    }
+}
+
+impl<T: ActixStream> Deref for TlsStream<T> {
+    type Target = tokio_native_tls::TlsStream<T>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<T: ActixStream> DerefMut for TlsStream<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl<T: ActixStream> AsyncRead for TlsStream<T> {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        Pin::new(&mut **self.get_mut()).poll_read(cx, buf)
+    }
+}
+
+impl<T: ActixStream> AsyncWrite for TlsStream<T> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        Pin::new(&mut **self.get_mut()).poll_write(cx, buf)
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut **self.get_mut()).poll_flush(cx)
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut **self.get_mut()).poll_shutdown(cx)
+    }
+
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[IoSlice<'_>],
+    ) -> Poll<io::Result<usize>> {
+        Pin::new(&mut **self.get_mut()).poll_write_vectored(cx, bufs)
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        (&**self).is_write_vectored()
+    }
+}
+
+impl<T: ActixStream> ActixStream for TlsStream<T> {
+    fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        T::poll_read_ready((&**self).get_ref().get_ref().get_ref(), cx)
+    }
+
+    fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        T::poll_write_ready((&**self).get_ref().get_ref().get_ref(), cx)
+    }
+}
 
 /// Accept TLS connections via `native-tls` package.
 ///
@@ -34,10 +112,7 @@ impl Clone for Acceptor {
     }
 }
 
-impl<T> ServiceFactory<T> for Acceptor
-where
-    T: AsyncRead + AsyncWrite + Unpin + 'static,
-{
+impl<T: ActixStream> ServiceFactory<T> for Acceptor {
     type Response = TlsStream<T>;
     type Error = Error;
     type Config = ();
@@ -71,10 +146,7 @@ impl Clone for NativeTlsAcceptorService {
     }
 }
 
-impl<T> Service<T> for NativeTlsAcceptorService
-where
-    T: AsyncRead + AsyncWrite + Unpin + 'static,
-{
+impl<T: ActixStream> Service<T> for NativeTlsAcceptorService {
     type Response = TlsStream<T>;
     type Error = Error;
     type Future = LocalBoxFuture<'static, Result<TlsStream<T>, Error>>;
@@ -93,7 +165,7 @@ where
         Box::pin(async move {
             let io = this.acceptor.accept(io).await;
             drop(guard);
-            io
+            io.map(Into::into)
         })
     }
 }
diff --git a/actix-tls/src/accept/openssl.rs b/actix-tls/src/accept/openssl.rs
@@ -1,21 +1,96 @@
 use std::{
     future::Future,
+    io::{self, IoSlice},
+    ops::{Deref, DerefMut},
     pin::Pin,
     task::{Context, Poll},
 };
 
-use actix_codec::{AsyncRead, AsyncWrite};
+use actix_codec::{AsyncRead, AsyncWrite, ReadBuf};
+use actix_rt::net::ActixStream;
 use actix_service::{Service, ServiceFactory};
 use actix_utils::counter::{Counter, CounterGuard};
 use futures_core::{future::LocalBoxFuture, ready};
 
 pub use openssl::ssl::{
     AlpnError, Error as SslError, HandshakeError, Ssl, SslAcceptor, SslAcceptorBuilder,
 };
-pub use tokio_openssl::SslStream;
 
 use super::MAX_CONN_COUNTER;
 
+/// wrapper type for `tokio_openssl::SslStream` in order to impl `ActixStream` trait.
+pub struct SslStream<T>(tokio_openssl::SslStream<T>);
+
+impl<T> From<tokio_openssl::SslStream<T>> for SslStream<T> {
+    fn from(stream: tokio_openssl::SslStream<T>) -> Self {
+        Self(stream)
+    }
+}
+
+impl<T> Deref for SslStream<T> {
+    type Target = tokio_openssl::SslStream<T>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<T> DerefMut for SslStream<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl<T: ActixStream> AsyncRead for SslStream<T> {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        Pin::new(&mut **self.get_mut()).poll_read(cx, buf)
+    }
+}
+
+impl<T: ActixStream> AsyncWrite for SslStream<T> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<io::Result<usize>> {
+        Pin::new(&mut **self.get_mut()).poll_write(cx, buf)
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut **self.get_mut()).poll_flush(cx)
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut **self.get_mut()).poll_shutdown(cx)
+    }
+
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[IoSlice<'_>],
+    ) -> Poll<io::Result<usize>> {
+        Pin::new(&mut **self.get_mut()).poll_write_vectored(cx, bufs)
+    }
+
+    fn is_write_vectored(&self) -> bool {
+        (&**self).is_write_vectored()
+    }
+}
+
+impl<T: ActixStream> ActixStream for SslStream<T> {
+    fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        T::poll_read_ready((&**self).get_ref(), cx)
+    }
+
+    fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        T::poll_write_ready((&**self).get_ref(), cx)
+    }
+}
+
 /// Accept TLS connections via `openssl` package.
 ///
 /// `openssl` feature enables this `Acceptor` type.
@@ -40,10 +115,7 @@ impl Clone for Acceptor {
     }
 }
 
-impl<T> ServiceFactory<T> for Acceptor
-where
-    T: AsyncRead + AsyncWrite + Unpin + 'static,
-{
+impl<T: ActixStream> ServiceFactory<T> for Acceptor {
     type Response = SslStream<T>;
     type Error = SslError;
     type Config = ();
@@ -67,10 +139,7 @@ pub struct AcceptorService {
     conns: Counter,
 }
 
-impl<T> Service<T> for AcceptorService
-where
-    T: AsyncRead + AsyncWrite + Unpin + 'static,
-{
+impl<T: ActixStream> Service<T> for AcceptorService {
     type Response = SslStream<T>;
     type Error = SslError;
     type Future = AcceptorServiceResponse<T>;
@@ -88,24 +157,25 @@ where
         let ssl = Ssl::new(ssl_ctx).expect("Provided SSL acceptor was invalid.");
         AcceptorServiceResponse {
             _guard: self.conns.get(),
-            stream: Some(SslStream::new(ssl, io).unwrap()),
+            stream: Some(tokio_openssl::SslStream::new(ssl, io).unwrap()),
         }
     }
 }
 
-pub struct AcceptorServiceResponse<T>
-where
-    T: AsyncRead + AsyncWrite,
-{
-    stream: Option<SslStream<T>>,
+pub struct AcceptorServiceResponse<T: ActixStream> {
+    stream: Option<tokio_openssl::SslStream<T>>,
     _guard: CounterGuard,
 }
 
-impl<T: AsyncRead + AsyncWrite + Unpin> Future for AcceptorServiceResponse<T> {
+impl<T: ActixStream> Future for AcceptorServiceResponse<T> {
     type Output = Result<SslStream<T>, SslError>;
 
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         ready!(Pin::new(self.stream.as_mut().unwrap()).poll_accept(cx))?;
-        Poll::Ready(Ok(self.stream.take().expect("SSL connect has resolved.")))
+        Poll::Ready(Ok(self
+            .stream
+            .take()
+            .expect("SSL connect has resolved.")
+            .into()))
     }
 }