From 3d5f43b2f573fc016bf955455a1738f9511ec132 Mon Sep 17 00:00:00 2001
From: Sadovoi Grisha <gsad1030@gmail.com>
Date: Fri, 3 Jul 2026 11:22:37 +0300
Subject: [PATCH] Fix JWK refresh cache check

When a caller requests a missing desired key, the cache guard used Option::map and short-circuited every Some(kid) request before checking whether the key existed. This prevented key rotation from loading new JWKS entries until restart.

Use and_then so only an actual cached entry skips the fetch, and add a regression test that refreshes from a JWKS endpoint after the requested kid changes.

Fixes #78

Signed-off-by: Sadovoi Grisha <gsad1030@gmail.com>
---
 lore-server/src/auth/jwk.rs | 76 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 75 insertions(+), 1 deletion(-)

diff --git a/lore-server/src/auth/jwk.rs b/lore-server/src/auth/jwk.rs
index 57bacabf..f83cc511 100644
--- a/lore-server/src/auth/jwk.rs
+++ b/lore-server/src/auth/jwk.rs
@@ -86,7 +86,7 @@ impl JwkServiceImpl {
         let mut cache = self.cached_set.write().await;
 
         // Check to see if the desired key was fetched while we waited for the lock.
-        if desired.map(|d| cache.get(d)).is_some() {
+        if desired.and_then(|d| cache.get(d)).is_some() {
             return Ok(());
         }
 
@@ -189,3 +189,77 @@ impl InstrumentProvider for JwkServiceImpl {
         "urc.auth.jwk_service"
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::net::SocketAddr;
+    use std::sync::Arc;
+    use std::sync::atomic::AtomicUsize;
+    use std::sync::atomic::Ordering;
+
+    use axum::Json;
+    use axum::Router;
+    use axum::extract::State;
+    use axum::routing::get;
+    use serde_json::json;
+    use tokio::net::TcpListener;
+
+    use super::*;
+
+    async fn jwks_handler(State(requests): State<Arc<AtomicUsize>>) -> Json<serde_json::Value> {
+        let kid = if requests.fetch_add(1, Ordering::SeqCst) == 0 {
+            "old-kid"
+        } else {
+            "new-kid"
+        };
+
+        Json(json!({
+            "keys": [{
+                "kty": "oct",
+                "use": "sig",
+                "kid": kid,
+                "alg": "HS256",
+                "k": "c2VjcmV0"
+            }]
+        }))
+    }
+
+    async fn spawn_jwks_server(requests: Arc<AtomicUsize>) -> SocketAddr {
+        let app = Router::new()
+            .route("/jwks", get(jwks_handler))
+            .with_state(requests);
+        let listener = TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("bind test jwks server");
+        let address = listener.local_addr().expect("get test jwks server address");
+
+        tokio::spawn(async move {
+            axum::serve(listener, app)
+                .await
+                .expect("serve test jwks server");
+        });
+
+        address
+    }
+
+    #[tokio::test]
+    async fn fetch_new_keys_refreshes_when_desired_key_is_missing() {
+        let requests = Arc::new(AtomicUsize::new(0));
+        let address = spawn_jwks_server(requests.clone()).await;
+        let service = JwkServiceImpl::new(JWKServiceSettings {
+            endpoint: format!("http://{address}/jwks"),
+        });
+
+        service
+            .fetch_new_keys(None)
+            .await
+            .expect("initial key fetch should succeed");
+
+        service
+            .get_key("new-kid")
+            .await
+            .expect("missing desired key should trigger a refresh");
+
+        assert_eq!(requests.load(Ordering::SeqCst), 2);
+    }
+}