1 files changed, 39 insertions, 13 deletions
diff --git a/src/https.rs b/src/https.rs
index 34e3f85..6b1f5e7 100644
--- a/src/https.rs
+++ b/src/https.rs
@@ -7,6 +7,7 @@ use log::*;
 
 use accept_encoding_fork::Encoding;
 use async_compression::tokio::bufread::*;
+use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use futures::TryStreamExt;
 use http::header::{HeaderName, HeaderValue};
@@ -15,6 +16,7 @@ use hyper::server::conn::Http;
 use hyper::service::service_fn;
 use hyper::{header, Body, Request, Response, StatusCode};
 use tokio::net::TcpListener;
+use tokio::select;
 use tokio::sync::watch;
 use tokio_rustls::TlsAcceptor;
 use tokio_util::io::{ReaderStream, StreamReader};
@@ -33,6 +35,7 @@ pub async fn serve_https(
 	config: HttpsConfig,
 	cert_store: Arc<CertStore>,
 	rx_proxy_config: watch::Receiver<Arc<ProxyConfig>>,
+	mut must_exit: watch::Receiver<bool>,
 ) -> Result<()> {
 	let config = Arc::new(config);
 
@@ -47,28 +50,43 @@ pub async fn serve_https(
 	info!("Starting to serve on https://{}.", config.bind_addr);
 
 	let tcp = TcpListener::bind(config.bind_addr).await?;
-	loop {
-		let (socket, remote_addr) = tcp.accept().await?;
+	let mut connections = FuturesUnordered::new();
+
+	while !*must_exit.borrow() {
+		let (socket, remote_addr) = select! {
+			a = tcp.accept() => a?,
+			_ = connections.next() => continue,
+			_ = must_exit.changed() => continue,
+		};
 
 		let rx_proxy_config = rx_proxy_config.clone();
 		let tls_acceptor = tls_acceptor.clone();
 		let config = config.clone();
 
-		tokio::spawn(async move {
+		let mut must_exit_2 = must_exit.clone();
+		let conn = tokio::spawn(async move {
 			match tls_acceptor.accept(socket).await {
 				Ok(stream) => {
 					debug!("TLS handshake was successfull");
-					let http_result = Http::new()
-						.serve_connection(
-							stream,
-							service_fn(move |req: Request<Body>| {
-								let https_config = config.clone();
-								let proxy_config: Arc<ProxyConfig> =
-									rx_proxy_config.borrow().clone();
-								handle_outer(remote_addr, req, https_config, proxy_config)
-							}),
+					let http_conn = Http::new().serve_connection(
+						stream,
+						service_fn(move |req: Request<Body>| {
+							let https_config = config.clone();
+							let proxy_config: Arc<ProxyConfig> = rx_proxy_config.borrow().clone();
+							handle_outer(remote_addr, req, https_config, proxy_config)
+						}),
+					);
+					tokio::pin!(http_conn);
+					let http_result = loop {
+						select! (
+							r = &mut http_conn => break r,
+							_ = must_exit_2.changed() => {
+								if *must_exit_2.borrow() {
+									http_conn.as_mut().graceful_shutdown();
+								}
+							}
 						)
-						.await;
+					};
 					if let Err(http_err) = http_result {
 						warn!("HTTP error: {}", http_err);
 					}
@@ -76,7 +94,15 @@ pub async fn serve_https(
 				Err(e) => warn!("Error in TLS connection: {}", e),
 			}
 		});
+		connections.push(conn);
 	}
+
+	info!("HTTPS server shutting down, draining remaining connections...");
+	while !connections.is_empty() {
+		let _ = connections.next().await;
+	}
+
+	Ok(())
 }
 
 async fn handle_outer(