Better handle get_cert for https request (faster path, hostname verification)

1 files changed, 36 insertions, 6 deletions

diff --git a/src/cert_store.rs b/src/cert_store.rs
index a2f67ec..f1b7d2b 100644
--- a/src/cert_store.rs
+++ b/src/cert_store.rs
@@ -19,20 +19,21 @@ use crate::proxy_config::ProxyConfig;
 pub struct CertStore {
 	consul: Consul,
 	certs: RwLock<HashMap<String, Arc<Cert>>>,
+	rx_proxy_config: watch::Receiver<Arc<ProxyConfig>>,
 }
 
 impl CertStore {
-	pub fn new(consul: Consul) -> Arc<Self> {
+	pub fn new(consul: Consul, rx_proxy_config: watch::Receiver<Arc<ProxyConfig>>) -> Arc<Self> {
 		Arc::new(Self {
 			consul,
 			certs: RwLock::new(HashMap::new()),
+			rx_proxy_config,
 		})
 	}
 
-	pub async fn watch_proxy_config(
-		self: Arc<Self>,
-		mut rx_proxy_config: watch::Receiver<Arc<ProxyConfig>>,
-	) {
+	pub async fn watch_proxy_config(self: Arc<Self>) {
+		let mut rx_proxy_config = self.rx_proxy_config.clone();
+
 		while rx_proxy_config.changed().await.is_ok() {
 			let mut domains: HashSet<String> = HashSet::new();
 
@@ -50,6 +51,35 @@ impl CertStore {
 		}
 	}
 
+	pub fn get_cert_for_https(self: &Arc<Self>, domain: &str) -> Result<Arc<Cert>> {
+		// Check if domain is authorized
+		if !self
+			.rx_proxy_config
+			.borrow()
+			.entries
+			.iter()
+			.any(|ent| ent.host == domain)
+		{
+			bail!("Domain {} should not have a TLS certificate.", domain);
+		}
+
+		// Check in local memory if it exists
+		let certs = self.certs.read().unwrap();
+		if let Some(cert) = certs.get(domain) {
+			if !cert.is_old() {
+				return Ok(cert.clone());
+			}
+		}
+
+		// Not found in local memory
+		tokio::spawn(self.clone().get_cert_task(domain.to_string()));
+		bail!("Certificate not found (will try to get it in background)");
+	}
+
+	pub async fn get_cert_task(self: Arc<Self>, domain: String) -> Result<Arc<Cert>> {
+		self.get_cert(domain.as_str()).await
+	}
+
 	pub async fn get_cert(self: &Arc<Self>, domain: &str) -> Result<Arc<Cert>> {
 		// First, try locally.
 		{
@@ -196,7 +226,7 @@ pub struct StoreResolver(pub Arc<CertStore>);
 impl rustls::server::ResolvesServerCert for StoreResolver {
 	fn resolve(&self, client_hello: rustls::server::ClientHello<'_>) -> Option<Arc<CertifiedKey>> {
 		let domain = client_hello.server_name()?;
-		let cert = futures::executor::block_on(self.0.get_cert(domain)).ok()?;
+		let cert = self.0.get_cert_for_https(domain).ok()?;
 		Some(cert.certkey.clone())
 	}
 }