use std::{ io::{self, IoSlice}, ops::{Deref, DerefMut}, pin::Pin, task::{Context, Poll}, }; use actix_codec::{AsyncRead, AsyncWrite, ReadBuf}; 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; use super::MAX_CONN_COUNTER; /// Wrapper type for `tokio_native_tls::TlsStream` in order to impl `ActixStream` trait. pub struct TlsStream(tokio_native_tls::TlsStream); impl From> for TlsStream { fn from(stream: tokio_native_tls::TlsStream) -> Self { Self(stream) } } impl Deref for TlsStream { type Target = tokio_native_tls::TlsStream; fn deref(&self) -> &Self::Target { &self.0 } } impl DerefMut for TlsStream { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.0 } } impl AsyncRead for TlsStream { fn poll_read( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { Pin::new(&mut **self.get_mut()).poll_read(cx, buf) } } impl AsyncWrite for TlsStream { fn poll_write( self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { Pin::new(&mut **self.get_mut()).poll_write(cx, buf) } fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut **self.get_mut()).poll_flush(cx) } fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut **self.get_mut()).poll_shutdown(cx) } fn poll_write_vectored( self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>], ) -> Poll> { Pin::new(&mut **self.get_mut()).poll_write_vectored(cx, bufs) } fn is_write_vectored(&self) -> bool { (&**self).is_write_vectored() } } impl ActixStream for TlsStream { fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll> { T::poll_read_ready((&**self).get_ref().get_ref().get_ref(), cx) } fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll> { T::poll_write_ready((&**self).get_ref().get_ref().get_ref(), cx) } } /// Accept TLS connections via `native-tls` package. /// /// `native-tls` feature enables this `Acceptor` type. pub struct Acceptor { acceptor: TlsAcceptor, } impl Acceptor { /// Create `native-tls` based `Acceptor` service factory. #[inline] pub fn new(acceptor: TlsAcceptor) -> Self { Acceptor { acceptor } } } impl Clone for Acceptor { #[inline] fn clone(&self) -> Self { Self { acceptor: self.acceptor.clone(), } } } impl ServiceFactory for Acceptor { type Response = TlsStream; type Error = Error; type Config = (); type Service = NativeTlsAcceptorService; type InitError = (); type Future = LocalBoxFuture<'static, Result>; fn new_service(&self, _: ()) -> Self::Future { let res = MAX_CONN_COUNTER.with(|conns| { Ok(NativeTlsAcceptorService { acceptor: self.acceptor.clone(), conns: conns.clone(), }) }); Box::pin(async { res }) } } pub struct NativeTlsAcceptorService { acceptor: TlsAcceptor, conns: Counter, } impl Clone for NativeTlsAcceptorService { fn clone(&self) -> Self { Self { acceptor: self.acceptor.clone(), conns: self.conns.clone(), } } } impl Service for NativeTlsAcceptorService { type Response = TlsStream; type Error = Error; type Future = LocalBoxFuture<'static, Result, Error>>; fn poll_ready(&self, cx: &mut Context<'_>) -> Poll> { if self.conns.available(cx) { Poll::Ready(Ok(())) } else { Poll::Pending } } fn call(&self, io: T) -> Self::Future { let guard = self.conns.get(); let this = self.clone(); Box::pin(async move { let io = this.acceptor.accept(io).await; drop(guard); io.map(Into::into) }) } }