Skip to content

Commit ae2f9c9

Browse files
authored
[rmcp] Notify credential stores when a refresh token is rejected (#1285)
* fix(auth): notify credential stores when refresh tokens are rejected * fix(auth): preserve refresh rejection when credential callback fails
1 parent bdd93dc commit ae2f9c9

1 file changed

Lines changed: 143 additions & 6 deletions

File tree

‎crates/rmcp/src/transport/auth.rs‎

Lines changed: 143 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,20 @@ pub trait CredentialStore: Send + Sync {
305305

306306
async fn clear(&self) -> Result<(), AuthError>;
307307

308+
/// Optionally handle the provider's definitive rejection of a refresh token.
309+
///
310+
/// Called only for a definitive `invalid_grant` response, with the credentials used
311+
/// for that exchange, before its refresh guard (if any) is released. Implementations
312+
/// must not reacquire that guard or change credentials that have since been replaced.
313+
/// The default leaves credentials unchanged. Callback errors are logged, and the
314+
/// manager still returns [`AuthError::TokenRefreshRejected`] so callers can reauthorize.
315+
async fn on_refresh_token_rejected(
316+
&self,
317+
_credentials: &StoredCredentials,
318+
) -> Result<(), AuthError> {
319+
Ok(())
320+
}
321+
308322
/// Optionally coordinate refreshes that share these credentials.
309323
///
310324
/// The manager acquires this guard before loading credentials and retains it
@@ -2279,6 +2293,7 @@ impl AuthorizationManager {
22792293
}
22802294
let current_credentials = stored_credentials
22812295
.token_response
2296+
.as_ref()
22822297
.ok_or(AuthError::AuthorizationRequired)?;
22832298

22842299
let refresh_token = current_credentials
@@ -2291,26 +2306,39 @@ impl AuthorizationManager {
22912306
.exchange_refresh_token(&refresh_token_value)
22922307
// RFC 8707: the resource indicator is required on token requests, including refreshes
22932308
.add_extra_param("resource", self.oauth_resource().await);
2294-
let mut refresh_scopes = stored_credentials.granted_scopes;
2309+
let mut refresh_scopes = stored_credentials.granted_scopes.clone();
22952310
self.add_offline_access_if_supported(&mut refresh_scopes);
22962311
let requested_scopes = refresh_scopes.clone();
22972312
for scope in refresh_scopes {
22982313
refresh_request = refresh_request.add_scope(Scope::new(scope));
22992314
}
2300-
let mut token_result = refresh_request
2315+
let mut token_result = match refresh_request
23012316
.request_async(&OAuth2HttpClient {
23022317
client: self.http_client.as_ref(),
23032318
redirect_policy: self.refresh_redirect_policy,
23042319
})
23052320
.await
2306-
.map_err(|error| match &error {
2321+
{
2322+
Ok(token_result) => token_result,
2323+
Err(error) => match &error {
23072324
RequestTokenError::ServerResponse(response)
23082325
if response.error() == &BasicErrorResponseType::InvalidGrant =>
23092326
{
2310-
AuthError::TokenRefreshRejected(error.to_string())
2327+
if let Err(store_error) = self
2328+
.credential_store
2329+
.on_refresh_token_rejected(&stored_credentials)
2330+
.await
2331+
{
2332+
tracing::warn!(
2333+
error = %store_error,
2334+
"Failed to handle rejected refresh token"
2335+
);
2336+
}
2337+
return Err(AuthError::TokenRefreshRejected(error.to_string()));
23112338
}
2312-
_ => AuthError::TokenRefreshFailed(error.to_string()),
2313-
})?;
2339+
_ => return Err(AuthError::TokenRefreshFailed(error.to_string())),
2340+
},
2341+
};
23142342

23152343
// RFC 6749 section 6: issuing a new refresh token on refresh is optional.
23162344
// When the response omits one, keep the existing refresh token rather than
@@ -8679,13 +8707,19 @@ mod tests {
86798707
#[tokio::test]
86808708
async fn invalid_grant_refresh_requires_reauthorization() {
86818709
let manager = manager_with_refresh_error("invalid_grant").await;
8710+
let before = manager.credential_store.load().await.unwrap();
86828711

86838712
let err = manager.try_refresh_or_reauth().await.unwrap_err();
86848713

86858714
assert!(
86868715
matches!(err, AuthError::AuthorizationRequired),
86878716
"expected AuthorizationRequired when the refresh token is rejected, got: {err:?}"
86888717
);
8718+
assert_eq!(
8719+
serde_json::to_value(manager.credential_store.load().await.unwrap()).unwrap(),
8720+
serde_json::to_value(before).unwrap(),
8721+
"the default rejection callback must leave credentials unchanged"
8722+
);
86898723
}
86908724

86918725
#[tokio::test]
@@ -9196,6 +9230,7 @@ mod tests {
91969230
credentials: InMemoryCredentialStore,
91979231
lock: Arc<Mutex<()>>,
91989232
events: Arc<StdMutex<Vec<&'static str>>>,
9233+
rejected_credentials: Arc<StdMutex<Option<StoredCredentials>>>,
91999234
guard_requested: Arc<Semaphore>,
92009235
save_started: Arc<Semaphore>,
92019236
save_gate: Option<Arc<Semaphore>>,
@@ -9241,6 +9276,19 @@ mod tests {
92419276
self.credentials.clear().await
92429277
}
92439278

9279+
async fn on_refresh_token_rejected(
9280+
&self,
9281+
credentials: &StoredCredentials,
9282+
) -> Result<(), AuthError> {
9283+
assert!(self.lock.try_lock().is_err());
9284+
self.events.lock().unwrap().push("rejected");
9285+
*self.rejected_credentials.lock().unwrap() = Some(credentials.clone());
9286+
if self.fail_at == Some("rejected") {
9287+
return Err(AuthError::CredentialStoreError("rejection failed".into()));
9288+
}
9289+
Ok(())
9290+
}
9291+
92449292
async fn acquire_refresh_guard(&self) -> Result<Option<CredentialRefreshGuard>, AuthError> {
92459293
self.events.lock().unwrap().push("acquire");
92469294
self.guard_requested.add_permits(1);
@@ -9269,6 +9317,7 @@ mod tests {
92699317
credentials: credential_store,
92709318
lock: Arc::new(Mutex::new(())),
92719319
events: Arc::new(StdMutex::new(Vec::new())),
9320+
rejected_credentials: Arc::new(StdMutex::new(None)),
92729321
guard_requested: Arc::new(Semaphore::new(0)),
92739322
save_started: Arc::new(Semaphore::new(0)),
92749323
save_gate: None,
@@ -9311,6 +9360,16 @@ mod tests {
93119360
})
93129361
}
93139362

9363+
fn refresh_error_http_client(store: &RefreshStore, error: &str) -> Arc<RefreshHttpClient> {
9364+
Arc::new(RefreshHttpClient {
9365+
recording: RecordingOAuthHttpClient::with_responses(vec![http_response(
9366+
400,
9367+
serde_json::json!({"error": error}),
9368+
)]),
9369+
events: store.events.clone(),
9370+
})
9371+
}
9372+
93149373
async fn refresh_manager(
93159374
store: RefreshStore,
93169375
http_client: Arc<RefreshHttpClient>,
@@ -9370,6 +9429,84 @@ mod tests {
93709429
assert!(store.lock.try_lock().is_ok());
93719430
}
93729431

9432+
#[rstest]
9433+
#[case(None)]
9434+
#[case(Some("rejected"))]
9435+
#[tokio::test]
9436+
async fn rejected_refresh_notifies_store_with_attempted_credentials_under_guard(
9437+
#[case] fail_at: Option<&'static str>,
9438+
) {
9439+
let mut store = refresh_store().await;
9440+
store.fail_at = fail_at;
9441+
let attempted = store.credentials.load().await.unwrap();
9442+
let http_client = refresh_error_http_client(&store, "invalid_grant");
9443+
let manager = refresh_manager(store.clone(), http_client.clone()).await;
9444+
9445+
let error = manager.refresh_token().await.unwrap_err();
9446+
9447+
assert!(matches!(error, AuthError::TokenRefreshRejected(_)));
9448+
assert_eq!(
9449+
serde_json::to_value(&*store.rejected_credentials.lock().unwrap()).unwrap(),
9450+
serde_json::to_value(attempted).unwrap()
9451+
);
9452+
assert_eq!(
9453+
*store.events.lock().unwrap(),
9454+
[
9455+
"acquire", "acquired", "load", "provider", "rejected", "release"
9456+
]
9457+
);
9458+
assert_eq!(http_client.recording.requests().len(), 1);
9459+
assert!(store.lock.try_lock().is_ok());
9460+
}
9461+
9462+
#[tokio::test]
9463+
async fn get_access_token_requires_reauth_when_rejection_callback_fails() {
9464+
let mut store = refresh_store().await;
9465+
store.fail_at = Some("rejected");
9466+
let mut credentials = store.credentials.load().await.unwrap().unwrap();
9467+
credentials
9468+
.token_response
9469+
.as_mut()
9470+
.unwrap()
9471+
.set_expires_in(Some(&std::time::Duration::ZERO));
9472+
store.credentials.save(credentials).await.unwrap();
9473+
let http_client = refresh_error_http_client(&store, "invalid_grant");
9474+
let manager = refresh_manager(store.clone(), http_client.clone()).await;
9475+
9476+
assert!(matches!(
9477+
manager.get_access_token().await,
9478+
Err(AuthError::AuthorizationRequired)
9479+
));
9480+
9481+
assert_eq!(
9482+
*store.events.lock().unwrap(),
9483+
[
9484+
"load", "acquire", "acquired", "load", "provider", "rejected", "release"
9485+
]
9486+
);
9487+
assert_eq!(http_client.recording.requests().len(), 1);
9488+
assert!(store.lock.try_lock().is_ok());
9489+
}
9490+
9491+
#[tokio::test]
9492+
async fn transient_refresh_failure_does_not_notify_store_of_rejection() {
9493+
let store = refresh_store().await;
9494+
let http_client = refresh_error_http_client(&store, "temporarily_unavailable");
9495+
let manager = refresh_manager(store.clone(), http_client).await;
9496+
9497+
assert!(matches!(
9498+
manager.refresh_token().await,
9499+
Err(AuthError::TokenRefreshFailed(_))
9500+
));
9501+
9502+
assert!(store.rejected_credentials.lock().unwrap().is_none());
9503+
assert_eq!(
9504+
*store.events.lock().unwrap(),
9505+
["acquire", "acquired", "load", "provider", "release"]
9506+
);
9507+
assert!(store.lock.try_lock().is_ok());
9508+
}
9509+
93739510
#[tokio::test]
93749511
async fn concurrent_refreshes_wait_for_save_and_use_the_latest_token() {
93759512
let mut store = refresh_store().await;

0 commit comments

Comments
 (0)