From 1d0ee252d0b443f052932a355185a7bc24a5310c Mon Sep 17 00:00:00 2001 From: Quentin Gliech Date: Tue, 17 Feb 2026 17:10:09 +0100 Subject: [PATCH] feat: auto-refresh expired upstream tokens during token exchange When a stored upstream OAuth token is expired and a refresh token is available, automatically refresh it before returning to the client. Co-Authored-By: Claude Opus 4.6 --- crates/handlers/src/lib.rs | 1 + crates/handlers/src/oauth2/token.rs | 88 +++++++++++++++++++++- crates/handlers/src/upstream_oauth2/mod.rs | 4 +- 3 files changed, 87 insertions(+), 6 deletions(-) diff --git a/crates/handlers/src/lib.rs b/crates/handlers/src/lib.rs index 20d17d982..f5b04acf6 100644 --- a/crates/handlers/src/lib.rs +++ b/crates/handlers/src/lib.rs @@ -216,6 +216,7 @@ where SiteConfig: FromRef, Templates: FromRef, Arc: FromRef, + MetadataCache: FromRef, BoxClock: FromRequestParts, BoxRng: FromRequestParts, Policy: FromRequestParts, diff --git a/crates/handlers/src/oauth2/token.rs b/crates/handlers/src/oauth2/token.rs index 74913d714..56cd5358e 100644 --- a/crates/handlers/src/oauth2/token.rs +++ b/crates/handlers/src/oauth2/token.rs @@ -174,7 +174,6 @@ pub(crate) enum RouteError { NoUpstreamToken, #[error("failed to refresh upstream token")] - #[expect(dead_code, reason = "constructed once auto-refresh is implemented")] UpstreamTokenRefreshFailed(#[source] anyhow::Error), } @@ -338,6 +337,7 @@ pub(crate) async fn post( State(homeserver): State>, State(site_config): State, State(encrypter): State, + State(metadata_cache): State, State(templates): State, policy: Policy, user_agent: Option>, @@ -447,11 +447,15 @@ pub(crate) async fn post( } AccessTokenRequest::TokenExchange(grant) => { token_exchange_grant( + &mut rng, &clock, &activity_tracker, &grant, &client, &encrypter, + &http_client, + &key_store, + &metadata_cache, repo, policy, user_agent, @@ -1112,11 +1116,15 @@ const ACCESS_TOKEN_TYPE_URN: &str = "urn:ietf:params:oauth:token-type:access_tok skip_all, )] async fn token_exchange_grant( + rng: &mut BoxRng, clock: &impl Clock, activity_tracker: &BoundActivityTracker, grant: &TokenExchangeGrant, client: &Client, encrypter: &Encrypter, + http_client: &reqwest::Client, + key_store: &Keystore, + metadata_cache: &crate::upstream_oauth2::cache::MetadataCache, mut repo: BoxRepository, mut policy: Policy, user_agent: Option, @@ -1237,20 +1245,92 @@ async fn token_exchange_grant( .ok_or(RouteError::NoUpstreamLink)?; // 7. Find the stored token for this link - let link_token = repo + let mut link_token = repo .upstream_oauth_link_token() .find_by_link(&link) .await? .ok_or(RouteError::NoUpstreamToken)?; - // 8. Decrypt the upstream access token + // 8. If the token is expired and we have a refresh token, auto-refresh + if link_token.is_expired(clock.now()) && link_token.has_refresh_token() { + let encrypted_refresh_token = link_token + .encrypted_refresh_token + .as_deref() + .expect("refresh token presence already checked"); + let refresh_token_bytes = encrypter + .decrypt_string(encrypted_refresh_token) + .map_err(|e| RouteError::Internal(Box::new(e)))?; + let refresh_token = String::from_utf8(refresh_token_bytes) + .map_err(|e| RouteError::Internal(Box::new(e)))?; + + let mut lazy_metadata = crate::upstream_oauth2::cache::LazyProviderInfos::new( + metadata_cache, + &provider, + http_client, + ); + let token_endpoint = lazy_metadata + .token_endpoint() + .await + .map_err(|e| RouteError::UpstreamTokenRefreshFailed(e.into()))? + .clone(); + + let client_credentials = crate::upstream_oauth2::client_credentials_for_provider( + &provider, + &token_endpoint, + key_store, + encrypter, + ) + .map_err(|e| RouteError::UpstreamTokenRefreshFailed(e.into()))?; + + let (token_response, _id_token) = + mas_oidc_client::requests::refresh_token::refresh_access_token( + http_client, + client_credentials, + &token_endpoint, + refresh_token, + None, + None, + None, + clock.now(), + rng, + ) + .await + .map_err(|e| RouteError::UpstreamTokenRefreshFailed(e.into()))?; + + // Encrypt and store the new tokens + let new_encrypted_access_token = encrypter + .encrypt_to_string(token_response.access_token.as_bytes()) + .map_err(|e| RouteError::Internal(Box::new(e)))?; + + let new_encrypted_refresh_token = token_response + .refresh_token + .as_deref() + .map(|rt| encrypter.encrypt_to_string(rt.as_bytes())) + .transpose() + .map_err(|e| RouteError::Internal(Box::new(e)))?; + + let new_expires_at = token_response.expires_in.map(|d| clock.now() + d); + + link_token = repo + .upstream_oauth_link_token() + .update_tokens( + clock, + link_token, + new_encrypted_access_token, + new_encrypted_refresh_token, + new_expires_at, + ) + .await?; + } + + // 9. Decrypt the upstream access token let decrypted = encrypter .decrypt_string(&link_token.encrypted_access_token) .map_err(|e| RouteError::Internal(Box::new(e)))?; let upstream_access_token = String::from_utf8(decrypted).map_err(|e| RouteError::Internal(Box::new(e)))?; - // 9. Build the response + // 10. Build the response let mut response = AccessTokenResponse::new(upstream_access_token); response.issued_token_type = Some(ACCESS_TOKEN_TYPE_URN.to_owned()); diff --git a/crates/handlers/src/upstream_oauth2/mod.rs b/crates/handlers/src/upstream_oauth2/mod.rs index dab366094..cab448ceb 100644 --- a/crates/handlers/src/upstream_oauth2/mod.rs +++ b/crates/handlers/src/upstream_oauth2/mod.rs @@ -26,7 +26,7 @@ mod template; use self::cookie::UpstreamSessions as UpstreamSessionsCookie; #[derive(Debug, Error)] -enum ProviderCredentialsError { +pub(crate) enum ProviderCredentialsError { #[error("Provider doesn't have a client secret")] MissingClientSecret, @@ -62,7 +62,7 @@ pub struct SignInWithApple { pub key_id: String, } -fn client_credentials_for_provider( +pub(crate) fn client_credentials_for_provider( provider: &UpstreamOAuthProvider, token_endpoint: &Url, keystore: &Keystore,