diff --git a/src/client/access_token.rs b/src/client/access_token.rs index a1d6939..63c1f6a 100644 --- a/src/client/access_token.rs +++ b/src/client/access_token.rs @@ -2,6 +2,7 @@ use anyhow::{Result, anyhow}; use chrono::Utc; use indexmap::IndexMap; use parking_lot::RwLock; +use std::collections::HashMap; use std::sync::LazyLock; type AccessTokenEntry = (String, i64, Option); @@ -9,6 +10,12 @@ type AccessTokenEntry = (String, i64, Option); static ACCESS_TOKENS: LazyLock>> = LazyLock::new(|| RwLock::new(IndexMap::new())); +/// Tokens a provider rejected (401) despite being locally unexpired. +/// Maps client name → the exact rejected token so a concurrently-refreshed +/// different token is never distrusted by mistake. +static REJECTED_TOKENS: LazyLock>> = + LazyLock::new(|| RwLock::new(HashMap::new())); + pub fn get_access_token(client_name: &str) -> Result { ACCESS_TOKENS .read() @@ -30,7 +37,7 @@ pub fn is_valid_access_token(client_name: &str) -> bool { Some(v) => v, None => return false, }; - !token.is_empty() && Utc::now().timestamp() < *expires_at + !token.is_empty() && Utc::now().timestamp() < *expires_at && !is_rejected(client_name, token) } pub fn set_access_token( @@ -45,3 +52,112 @@ pub fn set_access_token( entry.1 = expires_at; entry.2 = account_id; } + +/// Compare-and-invalidate a provider-rejected token. +/// +/// Only if the currently-cached token EQUALS `rejected` is the cache entry +/// removed and the rejection marker recorded; a concurrently-refreshed +/// different token is left untouched and no marker is set. +/// +/// Returns true if a cache entry existed for this client at all (whether or +/// not it matched `rejected`) — i.e. the client is token-authed and a retry +/// after refresh is worthwhile. Returns false when there is no entry +/// (API-key clients). +#[allow(dead_code)] // Called by the 401-retry path once it lands. +pub fn distrust_access_token(client_name: &str, rejected: &str) -> bool { + let mut access_tokens = ACCESS_TOKENS.write(); + let (token, _, _) = match access_tokens.get(client_name) { + Some(v) => v, + None => return false, + }; + if token == rejected { + access_tokens.shift_remove(client_name); + REJECTED_TOKENS + .write() + .insert(client_name.to_string(), rejected.to_string()); + } + true +} + +pub fn is_rejected(client_name: &str, token: &str) -> bool { + REJECTED_TOKENS + .read() + .get(client_name) + .is_some_and(|rejected| rejected == token) +} + +pub fn clear_rejected(client_name: &str) { + REJECTED_TOKENS.write().remove(client_name); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn distrust_removes_matching_token_and_sets_marker() { + let client = "distrust-match-test"; + set_access_token(client, "at-1".into(), Utc::now().timestamp() + 3600, None); + + assert!(distrust_access_token(client, "at-1")); + + assert!(get_access_token(client).is_err(), "cache entry not removed"); + assert!(is_rejected(client, "at-1"), "marker not set"); + } + + #[test] + fn distrust_keeps_differing_token_and_skips_marker() { + let client = "distrust-differ-test"; + set_access_token(client, "at-new".into(), Utc::now().timestamp() + 3600, None); + + assert!(distrust_access_token(client, "at-old")); + + assert_eq!(get_access_token(client).unwrap(), "at-new"); + assert!(!is_rejected(client, "at-old"), "marker set for stale token"); + } + + #[test] + fn distrust_returns_false_without_cache_entry() { + let client = "distrust-missing-test"; + + assert!(!distrust_access_token(client, "at-1")); + assert!(!is_rejected(client, "at-1")); + } + + #[test] + fn is_valid_access_token_false_for_rejected_token() { + let client = "rejected-valid-test"; + set_access_token(client, "at-1".into(), Utc::now().timestamp() + 3600, None); + assert!(is_valid_access_token(client)); + + distrust_access_token(client, "at-1"); + // A concurrent in-flight prepare re-caches the rejected file token + // between mark and refresh; it must still be treated as invalid. + set_access_token(client, "at-1".into(), Utc::now().timestamp() + 3600, None); + + assert!(!is_valid_access_token(client)); + } + + #[test] + fn clear_rejected_clears_marker_and_clients_are_isolated() { + let client_a = "rejected-isolation-a"; + let client_b = "rejected-isolation-b"; + set_access_token(client_a, "at-1".into(), Utc::now().timestamp() + 3600, None); + distrust_access_token(client_a, "at-1"); + + assert!(is_rejected(client_a, "at-1")); + assert!( + !is_rejected(client_b, "at-1"), + "marker leaked across clients" + ); + + clear_rejected(client_b); + assert!( + is_rejected(client_a, "at-1"), + "wrong client's marker cleared" + ); + + clear_rejected(client_a); + assert!(!is_rejected(client_a, "at-1")); + } +} diff --git a/src/client/oauth.rs b/src/client/oauth.rs index 60d86ba..7102948 100644 --- a/src/client/oauth.rs +++ b/src/client/oauth.rs @@ -1,4 +1,4 @@ -use super::access_token::{is_valid_access_token, set_access_token}; +use super::access_token::{clear_rejected, is_rejected, is_valid_access_token, set_access_token}; use super::openai_compatible_oauth::OpenAICompatibleOAuthProvider; use super::{ClientConfig, ProviderModels}; use crate::config::paths; @@ -743,7 +743,9 @@ pub async fn prepare_oauth_access_token( None => return Ok(false), }; - let tokens = if Utc::now().timestamp() >= tokens.expires_at { + let tokens = if Utc::now().timestamp() >= tokens.expires_at + || is_rejected(client_name, &tokens.access_token) + { let guard = refresh_guard(client_name); let _guard = guard.lock().await; @@ -759,7 +761,9 @@ pub async fn prepare_oauth_access_token( None => return Ok(false), }; - if Utc::now().timestamp() >= tokens.expires_at { + if Utc::now().timestamp() >= tokens.expires_at + || is_rejected(client_name, &tokens.access_token) + { match provider.flow() { OAuthFlow::Pkce | OAuthFlow::DeviceCode => { refresh_oauth_token(client, provider, client_name, &tokens).await? @@ -784,6 +788,9 @@ pub async fn prepare_oauth_access_token( tokens.expires_at, tokens.account_id, ); + // Clear even when the refresh returned the same token (some IdPs reuse + // JWTs within validity); otherwise every request re-hits the token endpoint. + clear_rejected(client_name); Ok(true) } @@ -1032,6 +1039,7 @@ mod tests { use std::time::UNIX_EPOCH; use super::*; + use crate::client::access_token::{distrust_access_token, get_access_token}; use crate::client::openai_compatible::OpenAICompatibleConfig; use crate::client::{ModelData, ProviderModels}; use crate::utils::get_env_name; @@ -1689,6 +1697,82 @@ scopes: }); } + #[test] + #[serial] + fn prepare_rejected_valid_file_token_attempts_refresh_branch() { + with_temp_cache(|| { + let client_name = "prepare-rejected-branch-test"; + let expires_at = Utc::now().timestamp() + 3600; + save_oauth_tokens( + client_name, + &OAuthTokens { + access_token: "rejected-at".into(), + refresh_token: None, + expires_at, + account_id: None, + }, + ) + .unwrap(); + set_access_token(client_name, "rejected-at".into(), expires_at, None); + assert!(distrust_access_token(client_name, "rejected-at")); + + let err = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(prepare_oauth_access_token( + &ReqwestClient::new(), + &ResourceStubProvider, + client_name, + )) + .unwrap_err() + .to_string(); + + // The timestamp-valid but rejected file token must not be trusted; + // the refresh branch is taken and bails on the missing refresh token. + assert!(err.contains("No refresh token"), "unexpected error: {err}"); + }); + } + + #[test] + #[serial] + fn prepare_trusts_differing_unmarked_valid_file_token() { + with_temp_cache(|| { + let client_name = "prepare-differing-token-test"; + let expires_at = Utc::now().timestamp() + 3600; + set_access_token(client_name, "rejected-at".into(), expires_at, None); + assert!(distrust_access_token(client_name, "rejected-at")); + save_oauth_tokens( + client_name, + &OAuthTokens { + access_token: "fresh-at".into(), + refresh_token: None, + expires_at, + account_id: None, + }, + ) + .unwrap(); + + let ready = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(prepare_oauth_access_token( + &ReqwestClient::new(), + &ResourceStubProvider, + client_name, + )) + .unwrap(); + + assert!(ready); + assert_eq!(get_access_token(client_name).unwrap(), "fresh-at"); + assert!( + !is_rejected(client_name, "rejected-at"), + "marker not cleared after successful prepare" + ); + }); + } + #[test] fn parse_refresh_response_invalid_grant_redacts_and_prompts_reauth() { let response = serde_json::json!({