Merge pull request #14 from joshtriplett/custom-tls-acceptor

Add support for custom TLS acceptors

Author: jbr

src/custom_tls_acceptor.rs

@@ -0,0 +1,31 @@
+use async_rustls::server::TlsStream;
+use async_std::net::TcpStream;
+
+/// The CustomTlsAcceptor trait provides a custom implementation of accepting
+/// TLS connections from a [`TcpStream`]. tide-rustls will call the
+/// [`CustomTlsAcceptor::accept`] function for each new [`TcpStream`] it
+/// accepts, to obtain a [`TlsStream`]).
+///
+/// Implementing this trait gives you control over the TLS negotiation process,
+/// and allows you to process some TLS connections internally without passing
+/// them through to tide, such as for multiplexing or custom ALPN negotiation.
+#[tide::utils::async_trait]
+pub trait CustomTlsAcceptor: Send + Sync {
+    /// Accept a [`TlsStream`] from a [`TcpStream`].
+    ///
+    /// If TLS negotiation succeeds, but does not result in a stream that tide
+    /// should process HTTP connections from, return `Ok(None)`.
+    async fn accept(&self, stream: TcpStream) -> std::io::Result<Option<TlsStream<TcpStream>>>;
+}
+
+/// Crate-private adapter to make `async_rustls::TlsAcceptor` implement
+/// `CustomTlsAcceptor`, without creating a conflict between the two `accept`
+/// methods.
+pub(crate) struct StandardTlsAcceptor(pub(crate) async_rustls::TlsAcceptor);
+
+#[tide::utils::async_trait]
+impl CustomTlsAcceptor for StandardTlsAcceptor {
+    async fn accept(&self, stream: TcpStream) -> std::io::Result<Option<TlsStream<TcpStream>>> {
+        self.0.accept(stream).await.map(Some)
+    }
+}

src/lib.rs

@@ -28,6 +28,7 @@
     unused_qualifications
 )]
 
+mod custom_tls_acceptor;
 mod tcp_connection;
 mod tls_listener;
 mod tls_listener_builder;
@@ -38,6 +39,7 @@ pub(crate) use tcp_connection::TcpConnection;
 pub(crate) use tls_listener_config::TlsListenerConfig;
 pub(crate) use tls_stream_wrapper::TlsStreamWrapper;
 
+pub use custom_tls_acceptor::CustomTlsAcceptor;
 pub use tls_listener::TlsListener;
 pub use tls_listener_builder::TlsListenerBuilder;
 

src/tls_listener.rs

@@ -1,4 +1,7 @@
-use crate::{TcpConnection, TlsListenerBuilder, TlsListenerConfig, TlsStreamWrapper};
+use crate::custom_tls_acceptor::StandardTlsAcceptor;
+use crate::{
+    CustomTlsAcceptor, TcpConnection, TlsListenerBuilder, TlsListenerConfig, TlsStreamWrapper,
+};
 
 use tide::listener::ListenInfo;
 use tide::listener::{Listener, ToListener};
@@ -79,12 +82,14 @@ impl<State> TlsListener<State> {
                     .set_single_cert(certs, keys.remove(0))
                     .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?;
 
-                TlsListenerConfig::Acceptor(TlsAcceptor::from(Arc::new(config)))
+                TlsListenerConfig::Acceptor(Arc::new(StandardTlsAcceptor(TlsAcceptor::from(
+                    Arc::new(config),
+                ))))
             }
 
-            TlsListenerConfig::ServerConfig(config) => {
-                TlsListenerConfig::Acceptor(TlsAcceptor::from(Arc::new(config)))
-            }
+            TlsListenerConfig::ServerConfig(config) => TlsListenerConfig::Acceptor(Arc::new(
+                StandardTlsAcceptor(TlsAcceptor::from(Arc::new(config))),
+            )),
 
             other @ TlsListenerConfig::Acceptor(_) => other,
 
@@ -99,7 +104,7 @@ impl<State> TlsListener<State> {
         Ok(())
     }
 
-    fn acceptor(&self) -> Option<&TlsAcceptor> {
+    fn acceptor(&self) -> Option<&Arc<dyn CustomTlsAcceptor>> {
         match self.config {
             TlsListenerConfig::Acceptor(ref a) => Some(a),
             _ => None,
@@ -125,14 +130,16 @@ impl<State> TlsListener<State> {
 fn handle_tls<State: Clone + Send + Sync + 'static>(
     app: Server<State>,
     stream: TcpStream,
-    acceptor: TlsAcceptor,
+    acceptor: Arc<dyn CustomTlsAcceptor>,
 ) {
     task::spawn(async move {
         let local_addr = stream.local_addr().ok();
         let peer_addr = stream.peer_addr().ok();
 
         match acceptor.accept(stream).await {
-            Ok(tls_stream) => {
+            Ok(None) => {}
+
+            Ok(Some(tls_stream)) => {
                 let stream = TlsStreamWrapper::new(tls_stream);
                 let fut = async_h1::accept(stream, |mut req| async {
                     if req.url_mut().set_scheme("https").is_err() {

src/tls_listener_builder.rs

@@ -3,11 +3,12 @@ use async_std::net::TcpListener;
 
 use rustls::ServerConfig;
 
-use super::{TcpConnection, TlsListener, TlsListenerConfig};
+use super::{CustomTlsAcceptor, TcpConnection, TlsListener, TlsListenerConfig};
 
 use std::marker::PhantomData;
 use std::net::{SocketAddr, ToSocketAddrs};
 use std::path::{Path, PathBuf};
+use std::sync::Arc;
 
 /// # A builder for TlsListeners
 ///
@@ -38,6 +39,7 @@ pub struct TlsListenerBuilder<State> {
     key: Option<PathBuf>,
     cert: Option<PathBuf>,
     config: Option<ServerConfig>,
+    tls_acceptor: Option<Arc<dyn CustomTlsAcceptor>>,
     tcp: Option<TcpListener>,
     addrs: Option<Vec<SocketAddr>>,
     _state: PhantomData<State>,
@@ -49,6 +51,7 @@ impl<State> Default for TlsListenerBuilder<State> {
             key: None,
             cert: None,
             config: None,
+            tls_acceptor: None,
             tcp: None,
             addrs: None,
             _state: PhantomData,
@@ -69,6 +72,14 @@ impl<State> std::fmt::Debug for TlsListenerBuilder<State> {
                     "None"
                 },
             )
+            .field(
+                "tls_acceptor",
+                &if self.tls_acceptor.is_some() {
+                    "Some(_)"
+                } else {
+                    "None"
+                },
+            )
             .field("tcp", &self.tcp)
             .field("addrs", &self.addrs)
             .finish()
@@ -108,6 +119,17 @@ impl<State> TlsListenerBuilder<State> {
         self
     }
 
+    /// Provides a custom acceptor for TLS connections.  This is mutually
+    /// exclusive with any of [`TlsListenerBuilder::key`],
+    /// [`TlsListenerBuilder::cert`], and [`TlsListenerBuilder::config`], but
+    /// gives total control over accepting TLS connections, including
+    /// multiplexing other streams or ALPN negotiations on the same TLS
+    /// connection that tide should ignore.
+    pub fn tls_acceptor(mut self, acceptor: Arc<dyn CustomTlsAcceptor>) -> Self {
+        self.tls_acceptor = Some(acceptor);
+        self
+    }
+
     /// Provides a bound tcp listener (either async-std or std) to
     /// build this tls listener on. This is mutually exclusive with
     /// [`TlsListenerBuilder::addrs`], but one of them is mandatory.
@@ -134,26 +156,29 @@ impl<State> TlsListenerBuilder<State> {
     /// * either of these is provided, but not both
     ///   * [`TlsListenerBuilder::tcp`]
     ///   * [`TlsListenerBuilder::addrs`]
-    /// * either of these is provided, but not both
+    /// * exactly one of these is provided
     ///   * both [`TlsListenerBuilder::cert`] AND [`TlsListenerBuilder::key`]
     ///   * [`TlsListenerBuilder::config`]
+    ///   * [`TlsListenerBuilder::tls_acceptor`]
     pub fn finish(self) -> io::Result<TlsListener<State>> {
         let Self {
             key,
             cert,
             config,
+            tls_acceptor,
             tcp,
             addrs,
             ..
         } = self;
 
-        let config = match (key, cert, config) {
-            (Some(key), Some(cert), None) => TlsListenerConfig::Paths { key, cert },
-            (None, None, Some(config)) => TlsListenerConfig::ServerConfig(config),
+        let config = match (key, cert, config, tls_acceptor) {
+            (Some(key), Some(cert), None, None) => TlsListenerConfig::Paths { key, cert },
+            (None, None, Some(config), None) => TlsListenerConfig::ServerConfig(config),
+            (None, None, None, Some(tls_acceptor)) => TlsListenerConfig::Acceptor(tls_acceptor),
             _ => {
                 return Err(io::Error::new(
                     io::ErrorKind::InvalidInput,
-                    "either cert + key are required or a ServerConfig",
+                    "need exactly one of cert + key, ServerConfig, or TLS acceptor",
                 ))
             }
         };

src/tls_listener_config.rs

@@ -1,9 +1,11 @@
 use std::fmt::{self, Debug, Formatter};
 
-use async_rustls::TlsAcceptor;
 use rustls::ServerConfig;
 
+use super::CustomTlsAcceptor;
+
 use std::path::PathBuf;
+use std::sync::Arc;
 
 impl Default for TlsListenerConfig {
     fn default() -> Self {
@@ -12,7 +14,7 @@ impl Default for TlsListenerConfig {
 }
 pub(crate) enum TlsListenerConfig {
     Unconfigured,
-    Acceptor(TlsAcceptor),
+    Acceptor(Arc<dyn CustomTlsAcceptor>),
     ServerConfig(ServerConfig),
     Paths { cert: PathBuf, key: PathBuf },
 }