From 3bdb417bfb87d7ef3381be2d56346a7995c54dde Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Wed, 8 Dec 2021 17:50:40 +0100
Subject: [PATCH] Exit more agressively on certain errors

---
 src/cert_store.rs    |  4 +++-
 src/http.rs          |  5 +----
 src/main.rs          | 29 ++++++++++++++++++++++-------
 src/proxy_config.rs  |  2 +-
 src/reverse_proxy.rs |  6 ++++--
 5 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/src/cert_store.rs b/src/cert_store.rs
index 6cc3ea9..e2ad62b 100644
--- a/src/cert_store.rs
+++ b/src/cert_store.rs
@@ -39,7 +39,7 @@ impl CertStore {
 		})
 	}
 
-	pub async fn watch_proxy_config(self: Arc<Self>) {
+	pub async fn watch_proxy_config(self: Arc<Self>) -> Result<()> {
 		let mut rx_proxy_config = self.rx_proxy_config.clone();
 
 		while rx_proxy_config.changed().await.is_ok() {
@@ -59,6 +59,8 @@ impl CertStore {
 				}
 			}
 		}
+
+		bail!("rx_proxy_config closed");
 	}
 
 	pub fn get_cert_for_https(self: &Arc<Self>, domain: &str) -> Result<Arc<Cert>> {
diff --git a/src/http.rs b/src/http.rs
index 2b26e6d..05d7440 100644
--- a/src/http.rs
+++ b/src/http.rs
@@ -12,10 +12,7 @@ use crate::consul::Consul;
 
 const CHALLENGE_PREFIX: &str = "/.well-known/acme-challenge/";
 
-pub async fn serve_http(
-	bind_addr: SocketAddr,
-	consul: Consul,
-) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+pub async fn serve_http(bind_addr: SocketAddr, consul: Consul) -> Result<()> {
 	let consul = Arc::new(consul);
 	// For every connection, we must make a `Service` to handle all
 	// incoming HTTP requests on said connection.
diff --git a/src/main.rs b/src/main.rs
index d495fb2..987c3ec 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,7 @@
 #[macro_use]
 extern crate anyhow;
 
+use futures::TryFutureExt;
 use std::net::SocketAddr;
 use structopt::StructOpt;
 
@@ -65,6 +66,12 @@ async fn main() {
 	}
 	pretty_env_logger::init();
 
+	// Abort on panic (same behavior as in Go)
+	std::panic::set_hook(Box::new(|panic_info| {
+		error!("{}", panic_info.to_string());
+		std::process::abort();
+	}));
+
 	let opt = Opt::from_args();
 
 	info!("Starting Tricot");
@@ -77,14 +84,17 @@ async fn main() {
 		rx_proxy_config.clone(),
 		opt.letsencrypt_email.clone(),
 	);
-	tokio::spawn(cert_store.clone().watch_proxy_config());
+	tokio::spawn(cert_store.clone().watch_proxy_config().map_err(exit_on_err));
 
-	tokio::spawn(http::serve_http(opt.http_bind_addr, consul.clone()));
-	tokio::spawn(https::serve_https(
-		opt.https_bind_addr,
-		cert_store.clone(),
-		rx_proxy_config.clone(),
-	));
+	tokio::spawn(http::serve_http(opt.http_bind_addr, consul.clone()).map_err(exit_on_err));
+	tokio::spawn(
+		https::serve_https(
+			opt.https_bind_addr,
+			cert_store.clone(),
+			rx_proxy_config.clone(),
+		)
+		.map_err(exit_on_err),
+	);
 
 	while rx_proxy_config.changed().await.is_ok() {
 		info!("Proxy config:");
@@ -93,3 +103,8 @@ async fn main() {
 		}
 	}
 }
+
+fn exit_on_err(e: anyhow::Error) -> () {
+	error!("{}", e);
+	std::process::exit(1);
+}
diff --git a/src/proxy_config.rs b/src/proxy_config.rs
index 009ca07..399b52a 100644
--- a/src/proxy_config.rs
+++ b/src/proxy_config.rs
@@ -102,7 +102,7 @@ fn parse_tricot_tag(
 		Some(i) => {
 			let (host, pp) = splits[1].split_at(i);
 			(host, Some(pp.to_string()))
-		},
+		}
 		None => (splits[1], None),
 	};
 
diff --git a/src/reverse_proxy.rs b/src/reverse_proxy.rs
index 1768724..401f4b1 100644
--- a/src/reverse_proxy.rs
+++ b/src/reverse_proxy.rs
@@ -72,7 +72,6 @@ fn create_proxied_request<B>(
 
 	*builder.headers_mut().unwrap() = remove_hop_headers(request.headers());
 
-
 	// If request does not have host header, add it from original URI authority
 	let host_header_name = "host";
 	if let Some(authority) = request.uri().authority() {
@@ -100,7 +99,10 @@ fn create_proxied_request<B>(
 		}
 	}
 
-	builder.headers_mut().unwrap().insert(HeaderName::from_bytes(b"x-forwarded-proto")?, "https".try_into()?);
+	builder.headers_mut().unwrap().insert(
+		HeaderName::from_bytes(b"x-forwarded-proto")?,
+		"https".try_into()?,
+	);
 
 	if let Some(conn) = request.headers().get("connection") {
 		if conn.to_str()?.to_lowercase() == "upgrade" {