Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
118 changes: 100 additions & 18 deletions crates/rmcp/src/transport/auth.rs
Original file line number Diff line number Diff line change
Expand Up @@ -278,6 +278,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]
Expand Down Expand Up @@ -2236,12 +2237,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
Expand All @@ -2259,8 +2267,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));
}
Expand All @@ -2286,10 +2294,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();
Expand Down Expand Up @@ -8451,17 +8461,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)),
Expand Down Expand Up @@ -8540,7 +8550,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();
Expand All @@ -8567,7 +8577,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);
Expand Down Expand Up @@ -8613,8 +8623,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();
Expand All @@ -8629,6 +8639,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)]
Expand All @@ -8638,7 +8720,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;
Expand Down
118 changes: 113 additions & 5 deletions crates/rmcp/src/transport/common/auth/streamable_http_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
};
Expand Down Expand Up @@ -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:?}"
);
}
}
Loading