From bec47f8c7de4ef25d9dbfcc8a4cc98338ff24f7e Mon Sep 17 00:00:00 2001 From: Sharang Parnerkar <30073382+mighty840@users.noreply.github.com> Date: Wed, 17 Jun 2026 21:38:06 +0200 Subject: [PATCH] feat(dashboard): proactively refresh expired Keycloak tokens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dashboard stored a refresh_token in the session at login (auth.rs) but never used it. Once the access_token's 5-minute lifespan ran out, every subsequent agent call failed with 401 ExpiredSignature. The UI showed "unable to load X" until the user logged out and back in. Fix: before attaching the bearer, decode the JWT's `exp` claim and proactively refresh via the stored refresh_token if the token is expired or within REFRESH_SKEW_SECS (30s) of expiry. Updates the session with the new access_token (and rotated refresh_token if KC sends one). Refresh failures fall through with the stale token so the agent's 401 surfaces to the UI rather than failing the request at the dashboard layer. Why "proactive" instead of "retry on 401" - Saves a wasted round-trip on every agent call once the token has aged past 5 min. - Doesn't require cloning RequestBuilder bodies for retry. - Same end state — fresh token reaches the agent. Test plan - cargo test -p compliance-dashboard --features server --no-default-features infrastructure::agent_client::tests — 5 pass: * expired JWT → refresh * near-expiry within skew window → refresh * fresh JWT → no refresh * malformed/empty JWT → refresh (defensive) * JWT without exp claim → refresh (defensive) - Manual after deploy: dashboard works past the 5-min token lifespan without manual re-login. Note - The refresh code addresses the ExpiredSignature failure mode. The separate "JWT is missing tenant_id claim" 401 is a Keycloak realm config issue (the user logging in lacks the M7.1 attributes that the protocol mappers consume) and is fixed by realm/attribute config, not by this PR. Co-Authored-By: Claude Opus 4.7 --- .../src/infrastructure/agent_client.rs | 165 +++++++++++++++++- 1 file changed, 157 insertions(+), 8 deletions(-) diff --git a/compliance-dashboard/src/infrastructure/agent_client.rs b/compliance-dashboard/src/infrastructure/agent_client.rs index f2354b0..679fb98 100644 --- a/compliance-dashboard/src/infrastructure/agent_client.rs +++ b/compliance-dashboard/src/infrastructure/agent_client.rs @@ -9,7 +9,16 @@ //! When Keycloak is not configured (dev convenience), the helper //! returns an unauthenticated builder — matching the agent's //! pass-through behavior in the same state. +//! +//! **Token refresh**: KC access tokens are short-lived (5 min default +//! in the certifai realm). Before attaching, we decode the JWT's `exp` +//! claim and proactively refresh via the stored refresh_token if the +//! access token is expired or about to expire. The session is updated +//! with the new pair. If refresh fails, we send the (stale) token +//! anyway — the agent's 401 will surface to the UI, which can prompt +//! re-login. +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine}; use dioxus::prelude::ServerFnError; use dioxus_fullstack::FullstackContext; use reqwest::Method; @@ -18,6 +27,11 @@ use super::auth::LOGGED_IN_USER_SESS_KEY; use super::server_state::ServerState; use super::user_state::UserStateInner; +/// Seconds before the JWT's `exp` time at which we consider it stale +/// enough to refresh. Covers clock skew + the round-trip to the agent +/// so the token doesn't expire mid-flight. +const REFRESH_SKEW_SECS: i64 = 30; + /// Build a `RequestBuilder` for `` with the /// session's access token attached. `path` should include a leading /// `/`, e.g. `"/api/v1/repositories"`. @@ -38,10 +52,9 @@ pub async fn agent_get(path: &str) -> Result req.bearer_auth(u.access_token), - None => req, - }) + let Some(mut user) = user else { + return Ok(req); + }; + + if token_needs_refresh(&user.access_token) { + tracing::debug!("Access token expired or near-expiring; refreshing"); + match refresh_tokens(state, &user.refresh_token).await { + Ok((new_access, new_refresh)) => { + user.access_token = new_access; + if let Some(rt) = new_refresh { + user.refresh_token = rt; + } + if let Err(e) = session.insert(LOGGED_IN_USER_SESS_KEY, &user).await { + tracing::warn!("Failed to persist refreshed tokens: {e}"); + } + } + Err(e) => { + tracing::warn!("Token refresh failed: {e}; sending current token anyway"); + // Fall through — the agent will 401 and the UI will + // prompt re-login. Better than failing the request at + // the dashboard layer with no helpful UX cue. + } + } + } + + Ok(req.bearer_auth(user.access_token)) +} + +/// Decode the JWT's payload (no signature verification — the agent +/// does that) and check the `exp` claim. Treats malformed tokens as +/// expired so the refresh path runs. +fn token_needs_refresh(jwt: &str) -> bool { + let Some(payload_b64) = jwt.split('.').nth(1) else { + return true; + }; + let Ok(bytes) = URL_SAFE_NO_PAD.decode(payload_b64) else { + return true; + }; + #[derive(serde::Deserialize)] + struct ExpClaim { + exp: i64, + } + let Ok(claims) = serde_json::from_slice::(&bytes) else { + return true; + }; + let now = chrono::Utc::now().timestamp(); + claims.exp - REFRESH_SKEW_SECS <= now +} + +/// Exchange a refresh_token for a new access_token. Returns the new +/// access_token and (optionally) the new refresh_token KC issued. +/// KC may rotate refresh_tokens on use; we honor whatever it sends. +async fn refresh_tokens( + state: &ServerState, + refresh_token: &str, +) -> Result<(String, Option), String> { + let kc = state + .keycloak + .ok_or_else(|| "Keycloak not configured".to_string())?; + if refresh_token.is_empty() { + return Err("no refresh_token in session".to_string()); + } + + #[derive(serde::Deserialize)] + struct TokenResp { + access_token: String, + refresh_token: Option, + } + + let resp = reqwest::Client::new() + .post(kc.token_endpoint()) + .form(&[ + ("grant_type", "refresh_token"), + ("client_id", kc.client_id.as_str()), + ("refresh_token", refresh_token), + ]) + .send() + .await + .map_err(|e| format!("refresh request failed: {e}"))?; + + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(format!("refresh rejected ({status}): {body}")); + } + + let r: TokenResp = resp + .json() + .await + .map_err(|e| format!("refresh response parse failed: {e}"))?; + Ok((r.access_token, r.refresh_token)) +} + +#[cfg(test)] +mod tests { + use super::*; + use base64::Engine; + + /// Build a JWT-shaped string (header.payload.sig) with the given + /// payload. Signature is bogus — we never verify it locally. + fn make_jwt(payload: &serde_json::Value) -> String { + let payload_b64 = URL_SAFE_NO_PAD.encode(serde_json::to_vec(payload).unwrap()); + format!("hdr.{payload_b64}.sig") + } + + #[test] + fn token_needs_refresh_true_when_expired() { + let exp = chrono::Utc::now().timestamp() - 60; + let jwt = make_jwt(&serde_json::json!({ "exp": exp })); + assert!(token_needs_refresh(&jwt)); + } + + #[test] + fn token_needs_refresh_true_within_skew_window() { + // 10 seconds left; less than the 30s skew → must refresh. + let exp = chrono::Utc::now().timestamp() + 10; + let jwt = make_jwt(&serde_json::json!({ "exp": exp })); + assert!(token_needs_refresh(&jwt)); + } + + #[test] + fn token_needs_refresh_false_with_plenty_of_life() { + let exp = chrono::Utc::now().timestamp() + 600; + let jwt = make_jwt(&serde_json::json!({ "exp": exp })); + assert!(!token_needs_refresh(&jwt)); + } + + #[test] + fn token_needs_refresh_true_on_malformed_jwt() { + assert!(token_needs_refresh("")); + assert!(token_needs_refresh("not.a.jwt")); + assert!(token_needs_refresh("only-one-segment")); + assert!(token_needs_refresh("hdr.not-base64!.sig")); + } + + #[test] + fn token_needs_refresh_true_when_exp_missing() { + let jwt = make_jwt(&serde_json::json!({ "sub": "abc" })); + assert!(token_needs_refresh(&jwt)); + } }