From 5404e143dcd19e879e0cfebb4e7328336fdaba9b Mon Sep 17 00:00:00 2001
From: "Dustin C. Hatch" <dustin@hatch.name>
Date: Sat, 4 Nov 2023 16:22:24 -0500
Subject: [PATCH] server: Cache machine IDs for 60 seconds

Since hosts have multiple keys that they will want to have signed, they
will need to make multiple requests, either sequentially or in parallel.
Since each request must be authenticated individually, this would result
in a libvirt connection and lookup for each one.  To avoid this
overhead, the server will now cache machine IDs in memory for 60
seconds.
---
 server/src/server/mod.rs | 35 +++++++++++++++++++++++++++++------
 1 file changed, 29 insertions(+), 6 deletions(-)

diff --git a/server/src/server/mod.rs b/server/src/server/mod.rs
index b5146e6..08f1a53 100644
--- a/server/src/server/mod.rs
+++ b/server/src/server/mod.rs
@@ -1,6 +1,8 @@
 mod host;
 
+use std::collections::HashMap;
 use std::sync::Arc;
+use std::time::{Duration, Instant};
 
 use axum::async_trait;
 use axum::extract::FromRequestParts;
@@ -11,7 +13,9 @@ use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
 use axum::routing::{get, post};
 use axum::{RequestPartsExt, Router, TypedHeader};
+use tokio::sync::RwLock;
 use tracing::debug;
+use uuid::Uuid;
 
 use crate::auth::{self, Claims};
 use crate::config::Configuration;
@@ -19,6 +23,7 @@ use crate::machine_id;
 
 struct Context {
     config: Arc<Configuration>,
+    cache: RwLock<HashMap<String, (Instant, Uuid)>>,
 }
 
 type State = Arc<Context>;
@@ -57,12 +62,10 @@ impl FromRequestParts<Arc<Context>> for Claims {
                 AuthError
             })?;
         let machine_id =
-            machine_id::get_machine_id(&hostname, ctx.config.clone())
-                .await
-                .ok_or_else(|| {
-                    debug!("No machine ID found for host {}", hostname);
-                    AuthError
-                })?;
+            get_machine_id(&hostname, ctx).await.ok_or_else(|| {
+                debug!("No machine ID found for host {}", hostname);
+                AuthError
+            })?;
         let claims = auth::validate_token(
             bearer.token(),
             &hostname,
@@ -81,6 +84,7 @@ impl FromRequestParts<Arc<Context>> for Claims {
 pub fn make_app(config: Configuration) -> Router {
     let ctx = Arc::new(Context {
         config: config.into(),
+        cache: RwLock::new(Default::default()),
     });
     Router::new()
         .route("/", get(|| async { "UP" }))
@@ -88,6 +92,25 @@ pub fn make_app(config: Configuration) -> Router {
         .with_state(ctx)
 }
 
+async fn get_machine_id(hostname: &str, ctx: &State) -> Option<Uuid> {
+    let cache = ctx.cache.read().await;
+    if let Some((ts, m)) = cache.get(hostname) {
+        if ts.elapsed() < Duration::from_secs(60) {
+            debug!("Found cached machine ID for {}", hostname);
+            return Some(*m);
+        } else {
+            debug!("Cached machine ID for {} has expired", hostname);
+        }
+    }
+    drop(cache);
+    let machine_id =
+        machine_id::get_machine_id(hostname, ctx.config.clone()).await?;
+    let mut cache = ctx.cache.write().await;
+    debug!("Caching machine ID for {}", hostname);
+    cache.insert(hostname.into(), (Instant::now(), machine_id));
+    Some(machine_id)
+}
+
 #[cfg(test)]
 mod test {
     use axum::body::Body;