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(auth): unify refresh checks and error handling (#1236)

Browse source 
* fix(auth): unify refresh checks and error handling

Claude-Session: https://claude.ai/code/session_011pHFfoTygeG84mCDXzCcmw

* test(auth): add live authorization-server checks

Claude-Session: https://claude.ai/code/session_011pHFfoTygeG84mCDXzCcmw
This commit is contained in:
Dale Seo 2026-09-08 19:04:15 -04:00 • committed by GitHub
commit 04c2f3fac9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 502 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -294,6 +294,7 @@ impl CredentialRefreshGuard {
/// Implementations of this trait can provide custom storage backends
/// for OAuth2 credentials, such as file-based storage, keychain integration,
/// or database storage.
///
/// Return [`AuthError::CredentialStoreError`] for backend or locking failures
/// so they remain distinct from errors requiring reauthorization.
#[async_trait]
@ -2248,12 +2249,19 @@ impl AuthorizationManager {
.as_ref()
.ok_or_else(|| AuthError::InternalError("OAuth client not configured".to_string()))?;
let refresh_guard = self.credential_store.acquire_refresh_guard().await?;
// Held for the rest of this function so the load, the exchange, and the
// save stay inside one guarded section.
let _refresh_guard = self.credential_store.acquire_refresh_guard().await?;
let stored = self.credential_store.load().await?;
let stored_credentials = stored.ok_or(AuthError::AuthorizationRequired)?;
if refresh_guard.is_some()
&& stored_credentials.client_id != oauth_client.client_id().as_str()
{
// Refreshing with another client's stored token would put that token on a
// request authenticated as this client.
if stored_credentials.client_id != oauth_client.client_id().as_str() {
tracing::warn!(
stored_client_id = stored_credentials.client_id.as_str(),
configured_client_id = oauth_client.client_id().as_str(),
"stored credentials belong to a different client; reauthorization required"
);
return Err(AuthError::AuthorizationRequired);
}
let current_credentials = stored_credentials
@ -2271,8 +2279,8 @@ impl AuthorizationManager {
// RFC 8707: the resource indicator is required on token requests, including refreshes
.add_extra_param("resource", self.oauth_resource().await);
let mut refresh_scopes = stored_credentials.granted_scopes;
let authoritative_scopes = refresh_guard.is_some().then(|| refresh_scopes.clone());
self.add_offline_access_if_supported(&mut refresh_scopes);
let requested_scopes = refresh_scopes.clone();
for scope in refresh_scopes {
refresh_request = refresh_request.add_scope(Scope::new(scope));
}
@ -2298,10 +2306,12 @@ impl AuthorizationManager {
token_result.set_refresh_token(Some(refresh_token_value));
}
let granted_scopes: Vec<String> = match (token_result.scopes(), authoritative_scopes) {
(Some(scopes), _) => scopes.iter().map(|s| s.to_string()).collect(),
(None, Some(scopes)) => scopes,
(None, None) => self.current_scopes.read().await.clone(),
let response_scopes = token_result
.scopes()
.map(|scopes| scopes.iter().map(|s| s.to_string()).collect());
let granted_scopes = {
let current = self.current_scopes.read().await;
Self::resolve_granted_scopes(response_scopes, &requested_scopes, &current)
};
*self.current_scopes.write().await = granted_scopes.clone();
@ -8515,17 +8525,17 @@ mod tests {
}
}
fn refresh_store() -> RefreshStore {
async fn refresh_store() -> RefreshStore {
let credentials = StoredCredentials::new(
"my-client".into(),
Some(make_token_response_with_refresh("old-token", "old-refresh")),
vec!["read".into()],
Some(AuthorizationManager::now_epoch_secs()),
);
let credential_store = InMemoryCredentialStore::new();
credential_store.save(credentials).await.unwrap();
RefreshStore {
credentials: InMemoryCredentialStore {
credentials: Arc::new(tokio::sync::RwLock::new(Some(credentials))),
},
credentials: credential_store,
lock: Arc::new(Mutex::new(())),
events: Arc::new(StdMutex::new(Vec::new())),
guard_requested: Arc::new(Semaphore::new(0)),
@ -8604,7 +8614,7 @@ mod tests {
#[tokio::test]
async fn refresh_guard_spans_load_exchange_and_completed_save() {
let store = refresh_store();
let store = refresh_store().await;
let manager = refresh_manager(store.clone(), refresh_http_client(&store)).await;
manager.refresh_token().await.unwrap();
@ -8631,7 +8641,7 @@ mod tests {
#[tokio::test]
async fn concurrent_refreshes_wait_for_save_and_use_the_latest_token() {
let mut store = refresh_store();
let mut store = refresh_store().await;
let save_gate = Arc::new(Semaphore::new(0));
store.save_gate = Some(save_gate.clone());
let http_client = refresh_http_client(&store);
@ -8677,8 +8687,8 @@ mod tests {
}
#[tokio::test]
async fn guarded_refresh_rejects_credentials_for_another_client() {
let store = refresh_store();
async fn refresh_rejects_credentials_for_another_client() {
let store = refresh_store().await;
let mut credentials = store.credentials.load().await.unwrap().unwrap();
credentials.client_id = "other-client".into();
store.credentials.save(credentials).await.unwrap();
@ -8693,6 +8703,78 @@ mod tests {
assert!(store.lock.try_lock().is_ok());
}
#[tokio::test]
async fn refresh_rejects_credentials_for_another_client_without_a_guard() {
let (base_url, captured) = start_token_server().await;
let mut manager = manager_with_metadata(Some(AuthorizationMetadata {
authorization_endpoint: format!("{base_url}/authorize"),
token_endpoint: format!("{base_url}/token"),
..Default::default()
}))
.await;
manager.configure_client(test_client_config()).unwrap();
manager
.credential_store
.save(StoredCredentials::new(
"other-client".into(),
Some(make_token_response_with_refresh("old-token", "old-refresh")),
vec!["read".into()],
Some(AuthorizationManager::now_epoch_secs()),
))
.await
.unwrap();
let error = manager.refresh_token().await.unwrap_err();
assert!(
matches!(error, AuthError::AuthorizationRequired),
"a client mismatch must require reauthorization, got: {error:?}"
);
assert!(
captured.lock().unwrap().is_none(),
"a client mismatch must be caught before the refresh token leaves the process"
);
}
#[tokio::test]
async fn refresh_without_a_guard_keeps_stored_scopes_when_response_omits_them() {
// start_token_server answers without a `scope`, matching a provider that
// grants the request in full.
let (base_url, _captured) = start_token_server().await;
let mut manager = manager_with_metadata(Some(AuthorizationMetadata {
authorization_endpoint: format!("{base_url}/authorize"),
token_endpoint: format!("{base_url}/token"),
..Default::default()
}))
.await;
manager.configure_client(test_client_config()).unwrap();
manager
.credential_store
.save(StoredCredentials::new(
"my-client".into(),
Some(make_token_response_with_refresh("old-token", "old-refresh")),
vec!["read".into()],
Some(AuthorizationManager::now_epoch_secs()),
))
.await
.unwrap();
*manager.current_scopes.write().await = vec!["stale".into()];
manager.refresh_token().await.unwrap();
let saved = manager.credential_store.load().await.unwrap().unwrap();
assert_eq!(
saved.granted_scopes,
["read"],
"the stored grant outranks the per-process scope cache"
);
assert_eq!(
manager.get_current_scopes().await,
["read"],
"the refreshed grant must replace the stale scope cache"
);
}
#[rstest]
#[case("guard", 0)]
#[case("load", 0)]
@ -8702,7 +8784,7 @@ mod tests {
#[case] phase: &'static str,
#[case] provider_requests: usize,
) {
let mut store = refresh_store();
let mut store = refresh_store().await;
store.fail_at = Some(phase);
let http_client = refresh_http_client(&store);
let manager = refresh_manager(store.clone(), http_client.clone()).await;

118
crates/rmcp/src/transport/common/auth/streamable_http_client.rs
View file

@ -19,7 +19,10 @@ where
/// 401 propagates as [`StreamableHttpError::AuthRequired`] carrying the
/// `WWW-Authenticate` challenge for the caller to authorize with;
/// - a token the server rejects (e.g. revoked) → one silent refresh, one
/// retry, then the challenge propagates.
/// retry, then the challenge propagates;
/// - a refresh that fails for any other reason (credential store, network,
/// provider) → that error propagates so the caller can retry instead of
/// being sent through a new authorization.
async fn call_reacting_to_challenges<T, F, Fut>(
&self,
auth_token: Option<String>,
@ -54,11 +57,13 @@ where
match refreshed {
Ok(fresh_token) if fresh_token != sent_token => call(Some(fresh_token)).await,
Ok(_) => Err(StreamableHttpError::AuthRequired(challenge)),
Err(error @ AuthError::CredentialStoreError(_)) => Err(error.into()),
Err(error) => {
debug!("token refresh after server rejection failed: {error}");
// `try_refresh_or_reauth` already reports the cases that need a
// new authorization; anything else is retryable or infrastructural.
Err(AuthError::AuthorizationRequired) => {
debug!("token refresh after server rejection requires authorization");
Err(StreamableHttpError::AuthRequired(challenge))
}
Err(error) => Err(error.into()),
}
}
result => result,
@ -212,11 +217,16 @@ where
#[cfg(all(test, feature = "transport-streamable-http-client-reqwest"))]
mod tests {
use std::sync::Arc;
use oauth2::{AccessToken, RefreshToken, basic::BasicTokenType};
use super::*;
use crate::transport::{
auth::{
AuthorizationManager, AuthorizationMetadata, CredentialRefreshGuard, CredentialStore,
StoredCredentials,
InMemoryCredentialStore, OAuthHttpClient, OAuthHttpClientFuture, OAuthHttpRequest,
OAuthTokenResponse, StoredCredentials, VendorExtraTokenFields,
},
streamable_http_client::AuthRequiredError,
};
@ -269,4 +279,102 @@ mod tests {
StreamableHttpError::Auth(AuthError::CredentialStoreError(message))
if message == "guard unavailable"));
}
struct UnreachableTokenEndpoint;
impl OAuthHttpClient for UnreachableTokenEndpoint {
fn execute(&self, _: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> {
Box::pin(async { Err("token endpoint unreachable".into()) })
}
}
struct RejectingTokenEndpoint;
impl OAuthHttpClient for RejectingTokenEndpoint {
fn execute(&self, _: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> {
Box::pin(async {
Ok(oauth2::http::Response::builder()
.status(400)
.header("content-type", "application/json")
.body(br#"{"error":"invalid_grant"}"#.to_vec())
.unwrap())
})
}
}
/// A manager holding a refresh token the given token endpoint will answer for.
async fn manager_with_stored_refresh_token(
token_endpoint: Arc<dyn OAuthHttpClient>,
) -> AuthorizationManager {
let mut manager = AuthorizationManager::new_with_oauth_http_client(
"https://mcp.example.com/mcp",
token_endpoint,
)
.await
.unwrap();
manager.set_metadata(AuthorizationMetadata {
authorization_endpoint: "https://auth.example.com/authorize".into(),
token_endpoint: "https://auth.example.com/token".into(),
..Default::default()
});
manager.configure_client_id("client").unwrap();
let mut token_response = OAuthTokenResponse::new(
AccessToken::new("old-token".into()),
BasicTokenType::Bearer,
VendorExtraTokenFields::default(),
);
token_response.set_refresh_token(Some(RefreshToken::new("stored-refresh".into())));
let store = InMemoryCredentialStore::new();
store
.save(StoredCredentials::new(
"client".into(),
Some(token_response),
vec![],
None,
))
.await
.unwrap();
manager.set_credential_store(store);
manager
}
/// Drive one call whose server answer is a 401 challenge.
async fn challenge_once(manager: AuthorizationManager) -> StreamableHttpError<reqwest::Error> {
AuthClient::new(reqwest::Client::new(), manager)
.call_reacting_to_challenges(Some("old-token".into()), |_| async {
Err::<(), _>(StreamableHttpError::AuthRequired(AuthRequiredError::new(
"Bearer".into(),
)))
})
.await
.unwrap_err()
}
#[tokio::test]
async fn reactive_refresh_propagates_retryable_refresh_failure() {
let manager = manager_with_stored_refresh_token(Arc::new(UnreachableTokenEndpoint)).await;
let error = challenge_once(manager).await;
assert!(
matches!(
error,
StreamableHttpError::Auth(AuthError::TokenRefreshFailed(_))
),
"a retryable refresh failure must reach the caller instead of asking for a new authorization, got: {error:?}"
);
}
#[tokio::test]
async fn reactive_refresh_reports_a_rejected_refresh_token_as_a_challenge() {
let manager = manager_with_stored_refresh_token(Arc::new(RejectingTokenEndpoint)).await;
let error = challenge_once(manager).await;
assert!(
matches!(error, StreamableHttpError::AuthRequired(_)),
"a definitively rejected refresh token must surface the challenge, got: {error:?}"
);
}
}