diff --git a/crates/basilica-sdk/src/auth/simple_manager.rs b/crates/basilica-sdk/src/auth/simple_manager.rs index 8cfc54587..48d01519b 100644 --- a/crates/basilica-sdk/src/auth/simple_manager.rs +++ b/crates/basilica-sdk/src/auth/simple_manager.rs @@ -7,7 +7,6 @@ use super::refresh::refresh_access_token; use super::token_store::TokenStore; use super::types::{get_sdk_data_dir, AuthError, AuthMethod, AuthResult, TokenSet}; use std::sync::Arc; -use std::time::Duration; use tokio::sync::Mutex; use tracing::{debug, info}; @@ -19,9 +18,6 @@ pub struct TokenManager { } impl TokenManager { - /// Pre-emptive refresh threshold (60 minutes before expiry) - const REFRESH_THRESHOLD: Duration = Duration::from_secs(3600); - /// Create a new token manager with direct tokens pub fn new_direct(access_token: String, refresh_token: String) -> Self { let tokens = TokenSet::new(access_token, refresh_token); @@ -108,12 +104,10 @@ impl TokenManager { } /// Check if token should be refreshed + /// + /// Defers to `TokenSet`, so this manager and the CLI renew on the same + /// schedule instead of each carrying its own threshold. fn should_refresh(&self, token_set: &TokenSet) -> bool { - if token_set.is_expired() { - return true; - } - - // Pre-emptive refresh if expiring within threshold - token_set.expires_within(Self::REFRESH_THRESHOLD) + token_set.is_expired() || token_set.needs_refresh() } } diff --git a/crates/basilica-sdk/src/auth/types.rs b/crates/basilica-sdk/src/auth/types.rs index 4ee05e1d3..3e19a1db8 100644 --- a/crates/basilica-sdk/src/auth/types.rs +++ b/crates/basilica-sdk/src/auth/types.rs @@ -7,7 +7,13 @@ use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine}; use etcetera::{choose_base_strategy, BaseStrategy}; use serde::{Deserialize, Serialize}; use std::path::PathBuf; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +/// Ceiling on how early a token is renewed, however long its lifetime. +const MAX_REFRESH_LEAD: Duration = Duration::from_secs(60 * 60); + +/// Lead used when the token carries no `iat`, so its lifetime is unknown. +const FALLBACK_REFRESH_LEAD: Duration = Duration::from_secs(5 * 60); /// Result type for authentication operations pub type AuthResult = Result; @@ -51,9 +57,8 @@ impl TokenSet { } } - /// Extract expiration from JWT token - /// Returns the exp claim from the JWT if it can be decoded - fn decode_jwt_exp(token: &str) -> Option { + /// Read a numeric claim from the access token's payload. + fn decode_jwt_claim(token: &str, claim: &str) -> Option { // JWT has three parts: header.payload.signature let parts: Vec<&str> = token.split('.').collect(); if parts.len() != 3 { @@ -66,15 +71,41 @@ impl TokenSet { // Decode base64url without padding (JWT uses base64url encoding) let decoded = URL_SAFE_NO_PAD.decode(payload).ok()?; - // Parse JSON and extract exp claim + // Parse JSON and extract the requested claim let json: serde_json::Value = serde_json::from_slice(&decoded).ok()?; - json.get("exp")?.as_u64() + json.get(claim)?.as_u64() } /// Get the expiration time by decoding JWT fn get_expiration(&self) -> Option { // Always decode from JWT token - Self::decode_jwt_exp(&self.access_token) + Self::decode_jwt_claim(&self.access_token, "exp") + } + + /// Total lifetime the issuer gave this token, from its own claims. + fn lifetime(&self) -> Option { + let exp = Self::decode_jwt_claim(&self.access_token, "exp")?; + let iat = Self::decode_jwt_claim(&self.access_token, "iat")?; + exp.checked_sub(iat).map(Duration::from_secs) + } + + /// How long before expiry a refresh should happen. + /// + /// Half the token's own lifetime, capped at an hour. A fixed lead is wrong + /// for short-lived tokens: with a one-hour token, a fixed one-hour lead is + /// satisfied the instant the token is minted, so every single call would + /// refresh. Scaling with the lifetime keeps roughly half the token usable + /// before the first renewal, whatever the issuer configures. + /// + /// A 24-hour token still gets the one-hour lead it had before this was + /// derived, so long-lived tokens behave exactly as they always did. + fn refresh_lead(&self) -> Duration { + match self.lifetime() { + Some(lifetime) => MAX_REFRESH_LEAD.min(lifetime / 2), + // No `iat` to measure against; fall back to the same small buffer + // the token stores use rather than assuming a long lifetime. + None => FALLBACK_REFRESH_LEAD, + } } /// Check if the access token is expired @@ -91,13 +122,15 @@ impl TokenSet { } } - /// Check if the token needs refresh (expires within 60 minutes) + /// Check whether the token is close enough to expiry to renew. + /// + /// The window scales with the token's own lifetime; see [`Self::refresh_lead`]. pub fn needs_refresh(&self) -> bool { - self.expires_within(std::time::Duration::from_secs(60 * 60)) + self.expires_within(self.refresh_lead()) } /// Check if the token expires within the specified duration - pub fn expires_within(&self, duration: std::time::Duration) -> bool { + pub fn expires_within(&self, duration: Duration) -> bool { match self.get_expiration() { Some(expires_at) => { let now = SystemTime::now() @@ -112,7 +145,7 @@ impl TokenSet { } /// Get time until token expiration - pub fn time_until_expiry(&self) -> Option { + pub fn time_until_expiry(&self) -> Option { match self.get_expiration() { Some(expires_at) => { let now = SystemTime::now() @@ -120,9 +153,9 @@ impl TokenSet { .unwrap() .as_secs(); if expires_at > now { - Some(std::time::Duration::from_secs(expires_at - now)) + Some(Duration::from_secs(expires_at - now)) } else { - Some(std::time::Duration::from_secs(0)) // Already expired + Some(Duration::from_secs(0)) // Already expired } } None => None, // No expiration time @@ -219,3 +252,72 @@ pub fn get_sdk_data_dir() -> AuthResult { // Use the same path as the CLI for consistency Ok(strategy.data_dir().join("basilica")) } + +#[cfg(test)] +mod tests { + use super::*; + + /// Build an access token whose payload carries the given claims. Only the + /// payload is read, so the header and signature can be anything. + fn token_with(iat: u64, exp: u64) -> TokenSet { + let payload = URL_SAFE_NO_PAD.encode(format!(r#"{{"iat":{iat},"exp":{exp}}}"#)); + TokenSet::new(format!("header.{payload}.signature"), "refresh".to_string()) + } + + fn now() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() + } + + #[test] + fn lead_is_half_the_lifetime_for_a_short_token() { + // One-hour token: renew in the final half hour, not immediately. + let t = token_with(now(), now() + 3600); + assert_eq!(t.refresh_lead(), Duration::from_secs(1800)); + } + + #[test] + fn a_freshly_minted_short_token_does_not_need_refresh() { + // The bug this replaces: a fixed one-hour lead made every call to a + // one-hour token refresh, because it was always "expiring within" it. + let t = token_with(now(), now() + 3600); + assert!(!t.needs_refresh()); + } + + #[test] + fn a_short_token_past_halfway_needs_refresh() { + let t = token_with(now() - 1900, now() + 1700); + assert!(t.needs_refresh()); + } + + #[test] + fn long_lived_tokens_keep_the_one_hour_lead() { + // A 24-hour token behaves exactly as it did before the lead was derived. + let t = token_with(now(), now() + 86_400); + assert_eq!(t.refresh_lead(), Duration::from_secs(3600)); + assert!(!t.needs_refresh()); + } + + #[test] + fn a_long_token_inside_the_last_hour_needs_refresh() { + let t = token_with(now() - 84_000, now() + 2_400); + assert!(t.needs_refresh()); + } + + #[test] + fn without_iat_the_lead_falls_back_to_the_small_buffer() { + let payload = URL_SAFE_NO_PAD.encode(format!(r#"{{"exp":{}}}"#, now() + 3600)); + let t = TokenSet::new(format!("header.{payload}.signature"), "refresh".to_string()); + assert_eq!(t.refresh_lead(), FALLBACK_REFRESH_LEAD); + assert!(!t.needs_refresh()); + } + + #[test] + fn an_expired_token_needs_refresh() { + let t = token_with(now() - 7200, now() - 3600); + assert!(t.is_expired()); + assert!(t.needs_refresh()); + } +}