From 7652e4c4732f624e3866f8a1ccf891f9629e226a Mon Sep 17 00:00:00 2001
From: ad hoc <postma.marin@protonmail.com>
Date: Fri, 15 Sep 2023 11:44:25 +0200
Subject: [PATCH] add AcceptorBuilder::with_acceptor method

---
 src/acceptor.rs         | 40 +++++++++++++++++++++++-----------------
 src/acceptor/builder.rs | 10 +++++++++-
 2 files changed, 32 insertions(+), 18 deletions(-)

diff --git a/src/acceptor.rs b/src/acceptor.rs
index 4cf816d..037b64e 100644
--- a/src/acceptor.rs
+++ b/src/acceptor.rs
@@ -17,9 +17,9 @@ pub use builder::AcceptorBuilder;
 use builder::WantsTlsConfig;
 
 /// A TLS acceptor that can be used with hyper servers.
-pub struct TlsAcceptor {
+pub struct TlsAcceptor<A = AddrIncoming> {
     config: Arc<ServerConfig>,
-    incoming: AddrIncoming,
+    acceptor: A,
 }
 
 /// An Acceptor for the `https` scheme.
@@ -31,12 +31,19 @@ impl TlsAcceptor {
 
     /// Creates a new `TlsAcceptor` from a `ServerConfig` and an `AddrIncoming`.
     pub fn new(config: Arc<ServerConfig>, incoming: AddrIncoming) -> Self {
-        Self { config, incoming }
+        Self {
+            config,
+            acceptor: incoming,
+        }
     }
 }
 
-impl Accept for TlsAcceptor {
-    type Conn = TlsStream;
+impl<A> Accept for TlsAcceptor<A>
+where
+    A: Accept<Error = io::Error> + Unpin,
+    A::Conn: AsyncRead + AsyncWrite + Unpin,
+{
+    type Conn = TlsStream<A::Conn>;
     type Error = io::Error;
 
     fn poll_accept(
@@ -44,7 +51,7 @@ impl Accept for TlsAcceptor {
         cx: &mut Context<'_>,
     ) -> Poll<Option<Result<Self::Conn, Self::Error>>> {
         let pin = self.get_mut();
-        Poll::Ready(match ready!(Pin::new(&mut pin.incoming).poll_accept(cx)) {
+        Poll::Ready(match ready!(Pin::new(&mut pin.acceptor).poll_accept(cx)) {
             Some(Ok(sock)) => Some(Ok(TlsStream::new(sock, pin.config.clone()))),
             Some(Err(e)) => Some(Err(e)),
             None => None,
@@ -66,22 +73,21 @@ where
 // 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 by handshaking with tokio_rustls::Accept first
-pub struct TlsStream {
-    state: State,
+pub struct TlsStream<C = AddrStream> {
+    state: State<C>,
 }
 
-impl TlsStream {
-    fn new(stream: AddrStream, config: Arc<ServerConfig>) -> Self {
+impl<C: AsyncRead + AsyncWrite + Unpin> TlsStream<C> {
+    fn new(stream: C, config: Arc<ServerConfig>) -> Self {
         let accept = tokio_rustls::TlsAcceptor::from(config).accept(stream);
         Self {
             state: State::Handshaking(accept),
         }
     }
-
     /// Returns a reference to the underlying IO stream.
     ///
     /// This should always return `Some`, except if an error has already been yielded.
-    pub fn io(&self) -> Option<&AddrStream> {
+    pub fn io(&self) -> Option<&C> {
         match &self.state {
             State::Handshaking(accept) => accept.get_ref(),
             State::Streaming(stream) => Some(stream.get_ref().0),
@@ -99,7 +105,7 @@ impl TlsStream {
     }
 }
 
-impl AsyncRead for TlsStream {
+impl<C: AsyncRead + AsyncWrite + Unpin> AsyncRead for TlsStream<C> {
     fn poll_read(
         self: Pin<&mut Self>,
         cx: &mut Context,
@@ -122,7 +128,7 @@ impl AsyncRead for TlsStream {
     }
 }
 
-impl AsyncWrite for TlsStream {
+impl<C: AsyncRead + AsyncWrite + Unpin> AsyncWrite for TlsStream<C> {
     fn poll_write(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
@@ -159,7 +165,7 @@ impl AsyncWrite for TlsStream {
     }
 }
 
-enum State {
-    Handshaking(tokio_rustls::Accept<AddrStream>),
-    Streaming(tokio_rustls::server::TlsStream<AddrStream>),
+enum State<C> {
+    Handshaking(tokio_rustls::Accept<C>),
+    Streaming(tokio_rustls::server::TlsStream<C>),
 }
diff --git a/src/acceptor/builder.rs b/src/acceptor/builder.rs
index 70c0ca7..9b95a82 100644
--- a/src/acceptor/builder.rs
+++ b/src/acceptor/builder.rs
@@ -90,9 +90,17 @@ impl AcceptorBuilder<WantsIncoming> {
     /// Passes a [`AddrIncoming`] to configure the TLS connection and
     /// creates the [`TlsAcceptor`]
     pub fn with_incoming(self, incoming: impl Into<AddrIncoming>) -> TlsAcceptor {
+        self.with_acceptor(incoming.into())
+    }
+
+    /// Passes an acceptor implementing [`Accept`] to configure the TLS connection and
+    /// creates the [`TlsAcceptor`]
+    ///
+    /// [`Accept`]: hyper::server::accept::Accept
+    pub fn with_acceptor<A>(self, acceptor: A) -> TlsAcceptor<A> {
         TlsAcceptor {
             config: Arc::new(self.0 .0),
-            incoming: incoming.into(),
+            acceptor,
         }
     }
 }