modelcontextprotocol/rust-sdk
Watch
1
0
Fork
You've already forked rust-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

fix(oauth): fix oauth credential refresh (#509)

Browse source 
auth.rs was using an in-memory expires-at which is only set on initial token exchange.
Instead, this PR switches it to use the expires-at set in the credentials that are passed in.
This commit is contained in:
Gabriel Peal 2025-11-04 13:17:11 -08:00 • committed by GitHub
commit dcbeb9b9bc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 10 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

36
crates/rmcp/src/transport/auth.rs
View file

@ -1,8 +1,4 @@
use std::{ use std::{collections::HashMap, sync::Arc, time::Duration};
collections::HashMap,
sync::Arc,
time::{Duration, Instant},
};
use oauth2::{ use oauth2::{
AuthUrl, AuthorizationCode, ClientId, ClientSecret, CsrfToken, EmptyExtraTokenFields, AuthUrl, AuthorizationCode, ClientId, ClientSecret, CsrfToken, EmptyExtraTokenFields,
@ -157,7 +153,6 @@ pub struct AuthorizationManager {
oauth_client: Option<OAuthClient>, oauth_client: Option<OAuthClient>,
credentials: RwLock<Option<OAuthTokenResponse>>, credentials: RwLock<Option<OAuthTokenResponse>>,
state: RwLock<Option<AuthorizationState>>, state: RwLock<Option<AuthorizationState>>,
expires_at: RwLock<Option<Instant>>,
base_url: Url, base_url: Url,
} }
@ -229,7 +224,6 @@ impl AuthorizationManager {
oauth_client: None, oauth_client: None,
credentials: RwLock::new(None), credentials: RwLock::new(None),
state: RwLock::new(None), state: RwLock::new(None),
expires_at: RwLock::new(None),
base_url, base_url,
}; };
@ -484,12 +478,6 @@ impl AuthorizationManager {
} }
}; };
// get expires_in from token response
let expires_in = token_result.expires_in();
if let Some(expires_in) = expires_in {
let expires_at = Instant::now() + expires_in;
*self.expires_at.write().await = Some(expires_at);
}
debug!("exchange token result: {:?}", token_result); debug!("exchange token result: {:?}", token_result);
// store credentials // store credentials
*self.credentials.write().await = Some(token_result.clone()); *self.credentials.write().await = Some(token_result.clone());
@ -503,13 +491,15 @@ impl AuthorizationManager {
if let Some(creds) = credentials.as_ref() { if let Some(creds) = credentials.as_ref() {
// check if the token is expire // check if the token is expire
if let Some(expires_at) = *self.expires_at.read().await { let expires_in = creds.expires_in().unwrap_or(Duration::from_secs(0));
if expires_at < Instant::now() { if expires_in <= Duration::from_secs(0) {
// token expired, try to refresh , release the lock tracing::info!("Access token expired, refreshing.");
drop(credentials); // token expired, try to refresh , release the lock
let new_creds = self.refresh_token().await?; drop(credentials);
return Ok(new_creds.access_token().secret().to_string());
} let new_creds = self.refresh_token().await?;
tracing::info!("Refreshed access token.");
return Ok(new_creds.access_token().secret().to_string());
} }
Ok(creds.access_token().secret().to_string()) Ok(creds.access_token().secret().to_string())
@ -548,12 +538,6 @@ impl AuthorizationManager {
// store new credentials // store new credentials
*self.credentials.write().await = Some(token_result.clone()); *self.credentials.write().await = Some(token_result.clone());
// get expires_in from token response
let expires_in = token_result.expires_in();
if let Some(expires_in) = expires_in {
let expires_at = Instant::now() + expires_in;
*self.expires_at.write().await = Some(expires_at);
}
Ok(token_result) Ok(token_result)
} }