From a2245c1751043146c4dc862db5cb5d5973239998 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Thu, 27 Sep 2018 21:06:27 -0700
Subject: [PATCH] add rust-tls support

---
 .travis.yml       |   6 +--
 src/ssl/mod.rs    |  12 +++---
 src/ssl/rustls.rs | 103 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 112 insertions(+), 9 deletions(-)
 create mode 100644 src/ssl/rustls.rs

diff --git a/.travis.yml b/.travis.yml
index 48a1414ac..767e9fde6 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -32,12 +32,12 @@ script:
   - |
     if [[ "$TRAVIS_RUST_VERSION" != "nightly" ]]; then
     cargo clean
-    cargo test --features="ssl,tls" -- --nocapture
+    cargo test --features="ssl,tls,rust-tls" -- --nocapture
     fi
   - |
     if [[ "$TRAVIS_RUST_VERSION" == "nightly" ]]; then
     RUSTFLAGS="--cfg procmacro2_semver_exempt" cargo install -f cargo-tarpaulin
-    cargo tarpaulin --features="ssl,tls" --out Xml
+    cargo tarpaulin --features="ssl,tls,rust-tls" --out Xml
     bash <(curl -s https://codecov.io/bash)
     echo "Uploaded code coverage"
     fi
@@ -46,7 +46,7 @@ script:
 after_success:
   - |
     if [[ "$TRAVIS_OS_NAME" == "linux" && "$TRAVIS_PULL_REQUEST" = "false" && "$TRAVIS_BRANCH" == "master" && "$TRAVIS_RUST_VERSION" == "beta" ]]; then
-      cargo doc --features "ssl,tls" --no-deps &&
+      cargo doc --features "ssl,tls,rust-tls" --no-deps &&
       echo "<meta http-equiv=refresh content=0;url=os_balloon/index.html>" > target/doc/index.html &&
       git clone https://github.com/davisp/ghp-import.git &&
       ./ghp-import/ghp_import.py -n -p -f -m "Documentation upload" -r https://"$GH_TOKEN"@github.com/"$TRAVIS_REPO_SLUG.git" target/doc &&
diff --git a/src/ssl/mod.rs b/src/ssl/mod.rs
index 81a364df5..906ee615c 100644
--- a/src/ssl/mod.rs
+++ b/src/ssl/mod.rs
@@ -13,7 +13,10 @@ mod nativetls;
 #[cfg(feature = "tls")]
 pub use self::nativetls::{NativeTlsAcceptor, TlsStream};
 
-pub(crate) const MAX_CONN: AtomicUsize = AtomicUsize::new(256);
+#[cfg(feature = "rust-tls")]
+mod rustls;
+#[cfg(feature = "rust-tls")]
+pub use self::rustls::RustlsAcceptor;
 
 /// Sets the maximum per-worker concurrent ssl connection establish process.
 ///
@@ -25,11 +28,8 @@ pub fn max_concurrent_ssl_connect(num: usize) {
     MAX_CONN.store(num, Ordering::Relaxed);
 }
 
+pub(crate) const MAX_CONN: AtomicUsize = AtomicUsize::new(256);
+
 thread_local! {
     static MAX_CONN_COUNTER: Counter = Counter::new(MAX_CONN.load(Ordering::Relaxed));
 }
-
-// #[cfg(feature = "rust-tls")]
-// mod rustls;
-// #[cfg(feature = "rust-tls")]
-// pub use self::rustls::RustlsAcceptor;
diff --git a/src/ssl/rustls.rs b/src/ssl/rustls.rs
new file mode 100644
index 000000000..ef4c966be
--- /dev/null
+++ b/src/ssl/rustls.rs
@@ -0,0 +1,103 @@
+use std::io;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+use futures::{future::ok, future::FutureResult, Async, Future, Poll};
+use rustls::{ServerConfig, ServerSession};
+use tokio_io::{AsyncRead, AsyncWrite};
+use tokio_rustls::{AcceptAsync, ServerConfigExt, TlsStream};
+
+use super::MAX_CONN_COUNTER;
+use counter::{Counter, CounterGuard};
+use service::{NewService, Service};
+
+/// Support `SSL` connections via rustls package
+///
+/// `rust-tls` feature enables `RustlsAcceptor` type
+pub struct RustlsAcceptor<T> {
+    config: Arc<ServerConfig>,
+    io: PhantomData<T>,
+}
+
+impl<T: AsyncRead + AsyncWrite> RustlsAcceptor<T> {
+    /// Create `RustlsAcceptor` new service
+    pub fn new(config: ServerConfig) -> Self {
+        RustlsAcceptor {
+            config: Arc::new(config),
+            io: PhantomData,
+        }
+    }
+}
+
+impl<T> Clone for RustlsAcceptor<T> {
+    fn clone(&self) -> Self {
+        Self {
+            config: self.config.clone(),
+            io: PhantomData,
+        }
+    }
+}
+
+impl<T: AsyncRead + AsyncWrite> NewService for RustlsAcceptor<T> {
+    type Request = T;
+    type Response = TlsStream<T, ServerSession>;
+    type Error = io::Error;
+    type Service = RustlsAcceptorService<T>;
+    type InitError = ();
+    type Future = FutureResult<Self::Service, Self::InitError>;
+
+    fn new_service(&self) -> Self::Future {
+        MAX_CONN_COUNTER.with(|conns| {
+            ok(RustlsAcceptorService {
+                config: self.config.clone(),
+                conns: conns.clone(),
+                io: PhantomData,
+            })
+        })
+    }
+}
+
+pub struct RustlsAcceptorService<T> {
+    config: Arc<ServerConfig>,
+    io: PhantomData<T>,
+    conns: Counter,
+}
+
+impl<T: AsyncRead + AsyncWrite> Service for RustlsAcceptorService<T> {
+    type Request = T;
+    type Response = TlsStream<T, ServerSession>;
+    type Error = io::Error;
+    type Future = RustlsAcceptorServiceFut<T>;
+
+    fn poll_ready(&mut self) -> Poll<(), Self::Error> {
+        if self.conns.available() {
+            Ok(Async::Ready(()))
+        } else {
+            Ok(Async::NotReady)
+        }
+    }
+
+    fn call(&mut self, req: Self::Request) -> Self::Future {
+        RustlsAcceptorServiceFut {
+            _guard: self.conns.get(),
+            fut: ServerConfigExt::accept_async(&self.config, req),
+        }
+    }
+}
+
+pub struct RustlsAcceptorServiceFut<T>
+where
+    T: AsyncRead + AsyncWrite,
+{
+    fut: AcceptAsync<T>,
+    _guard: CounterGuard,
+}
+
+impl<T: AsyncRead + AsyncWrite> Future for RustlsAcceptorServiceFut<T> {
+    type Item = TlsStream<T, ServerSession>;
+    type Error = io::Error;
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        self.fut.poll()
+    }
+}