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)); + } }