Skip to content

Commit dcbeb9b

Browse files
authored
fix(oauth): fix oauth credential refresh (#509)
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.
1 parent d3ddc09 commit dcbeb9b

File tree

1 file changed

+10
-26
lines changed

1 file changed

+10
-26
lines changed

crates/rmcp/src/transport/auth.rs

Lines changed: 10 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,4 @@
1-
use std::{
2-
collections::HashMap,
3-
sync::Arc,
4-
time::{Duration, Instant},
5-
};
1+
use std::{collections::HashMap, sync::Arc, time::Duration};
62

73
use oauth2::{
84
AuthUrl, AuthorizationCode, ClientId, ClientSecret, CsrfToken, EmptyExtraTokenFields,
@@ -157,7 +153,6 @@ pub struct AuthorizationManager {
157153
oauth_client: Option<OAuthClient>,
158154
credentials: RwLock<Option<OAuthTokenResponse>>,
159155
state: RwLock<Option<AuthorizationState>>,
160-
expires_at: RwLock<Option<Instant>>,
161156
base_url: Url,
162157
}
163158

@@ -229,7 +224,6 @@ impl AuthorizationManager {
229224
oauth_client: None,
230225
credentials: RwLock::new(None),
231226
state: RwLock::new(None),
232-
expires_at: RwLock::new(None),
233227
base_url,
234228
};
235229

@@ -484,12 +478,6 @@ impl AuthorizationManager {
484478
}
485479
};
486480

487-
// get expires_in from token response
488-
let expires_in = token_result.expires_in();
489-
if let Some(expires_in) = expires_in {
490-
let expires_at = Instant::now() + expires_in;
491-
*self.expires_at.write().await = Some(expires_at);
492-
}
493481
debug!("exchange token result: {:?}", token_result);
494482
// store credentials
495483
*self.credentials.write().await = Some(token_result.clone());
@@ -503,13 +491,15 @@ impl AuthorizationManager {
503491

504492
if let Some(creds) = credentials.as_ref() {
505493
// check if the token is expire
506-
if let Some(expires_at) = *self.expires_at.read().await {
507-
if expires_at < Instant::now() {
508-
// token expired, try to refresh , release the lock
509-
drop(credentials);
510-
let new_creds = self.refresh_token().await?;
511-
return Ok(new_creds.access_token().secret().to_string());
512-
}
494+
let expires_in = creds.expires_in().unwrap_or(Duration::from_secs(0));
495+
if expires_in <= Duration::from_secs(0) {
496+
tracing::info!("Access token expired, refreshing.");
497+
// token expired, try to refresh , release the lock
498+
drop(credentials);
499+
500+
let new_creds = self.refresh_token().await?;
501+
tracing::info!("Refreshed access token.");
502+
return Ok(new_creds.access_token().secret().to_string());
513503
}
514504

515505
Ok(creds.access_token().secret().to_string())
@@ -548,12 +538,6 @@ impl AuthorizationManager {
548538
// store new credentials
549539
*self.credentials.write().await = Some(token_result.clone());
550540

551-
// get expires_in from token response
552-
let expires_in = token_result.expires_in();
553-
if let Some(expires_in) = expires_in {
554-
let expires_at = Instant::now() + expires_in;
555-
*self.expires_at.write().await = Some(expires_at);
556-
}
557541
Ok(token_result)
558542
}
559543

0 commit comments

Comments
 (0)