From 48a421a3b6c07705b67276853643ed5601a30154 Mon Sep 17 00:00:00 2001
From: Quentin Dufour <quentin@deuxfleurs.fr>
Date: Tue, 8 Aug 2023 18:04:21 +0200
Subject: [PATCH] feat: warmup memory store when starting

---
 src/cert_store.rs | 66 ++++++++++++++++++++++++++++++++++++++++-------
 src/main.rs       | 12 +++++++++
 2 files changed, 69 insertions(+), 9 deletions(-)

diff --git a/src/cert_store.rs b/src/cert_store.rs
index abbc83d..2d1b772 100644
--- a/src/cert_store.rs
+++ b/src/cert_store.rs
@@ -145,6 +145,47 @@ impl CertStore {
 		self.gen_self_signed_certificate(domain)
 	}
 
+	pub async fn warmup_memory_store(self: &Arc<Self>) -> Result<()> {
+		let consul_certs = self
+			.consul
+			.kv_get_prefix("certs/", None)
+			.await?
+			.into_inner();
+
+		trace!(
+			"Fetched {} certificate entries from Consul",
+			consul_certs.len()
+		);
+		let mut loaded_certs: usize = 0;
+		for (domain, cert) in consul_certs {
+			let certser: CertSer = match serde_json::from_slice(&cert) {
+				Ok(cs) => cs,
+				Err(e) => {
+					warn!("Could not deserialize CertSer for {domain}: {e}");
+					continue;
+				}
+			};
+
+			let cert = match Cert::new(certser) {
+				Ok(c) => c,
+				Err(e) => {
+					warn!("Could not create Cert from CertSer for domain {domain}: {e}");
+					continue;
+				}
+			};
+
+			self.certs
+				.write()
+				.unwrap()
+				.insert(domain.to_string(), Arc::new(cert));
+
+			debug!("({domain}) Certificate loaded from Consul to the Memory Store");
+			loaded_certs += 1;
+		}
+		info!("Memory store warmed up with {loaded_certs} certificates");
+		Ok(())
+	}
+
 	pub async fn check_cert(self: &Arc<Self>, domain: &str) -> Result<()> {
 		// First, try locally.
 		{
@@ -162,16 +203,23 @@ impl CertStore {
 			.kv_get_json::<CertSer>(&format!("certs/{}", domain))
 			.await?
 		{
-			if let Ok(cert) = Cert::new(consul_cert) {
-				let cert = Arc::new(cert);
-				self.certs
-					.write()
-					.unwrap()
-					.insert(domain.to_string(), cert.clone());
-				if !cert.is_old() {
-					return Ok(());
+			match Cert::new(consul_cert) {
+				Ok(cert) => {
+					let cert = Arc::new(cert);
+					self.certs
+						.write()
+						.unwrap()
+						.insert(domain.to_string(), cert.clone());
+					debug!("({domain}) Certificate loaded from Consul to the Memory Store");
+
+					if !cert.is_old() {
+						return Ok(());
+					}
 				}
-			}
+				Err(e) => {
+					warn!("Could not create Cert from CertSer for domain {domain}: {e}");
+				}
+			};
 		}
 
 		// Third, ask from Let's Encrypt
diff --git a/src/main.rs b/src/main.rs
index c148a91..43f3447 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -101,6 +101,12 @@ struct Opt {
 		default_value = "text/html,text/plain,text/css,text/javascript,text/xml,application/javascript,application/json,application/xml,image/svg+xml,font/ttf"
 	)]
 	pub compress_mime_types: String,
+
+	#[structopt(
+		long = "warmup-cert-memory-store",
+		env = "TRICOT_WARMUP_CERT_MEMORY_STORE"
+	)]
+	pub warmup_cert_memory_store: bool,
 }
 
 #[tokio::main(flavor = "multi_thread", worker_threads = 10)]
@@ -154,6 +160,12 @@ async fn main() {
 		opt.letsencrypt_email.clone(),
 		exit_on_err.clone(),
 	);
+	if opt.warmup_cert_memory_store {
+		match cert_store.warmup_memory_store().await {
+            Err(e) => error!("An error occured while warming up the certificate memory store with Consul data, continue without caching: {e}"),
+            _ => (),
+        };
+	}
 
 	let metrics_task = tokio::spawn(
 		metrics_server