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
4 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:?}"
);
}
}

233
crates/rmcp/tests/test_live_oauth_refresh.rs Normal file
View file

@ -0,0 +1,233 @@
//! Live authorization-server checks for the refresh path.
//!
//! These are `#[ignore]`d, so they never run in CI. Provision the authorization
//! server with `./scripts/keycloak-oauth-fixture.sh`, then run them with
//! `cargo test -p rmcp --all-features --test test_live_oauth_refresh -- --ignored`.
//! `KC_BASE` overrides the server location for both the script and these tests.
#![cfg(feature = "auth")]
use std::sync::Arc;
use rmcp::transport::auth::{
AuthError, AuthorizationManager, AuthorizationMetadata, CredentialRefreshGuard,
CredentialStore, InMemoryCredentialStore, OAuthClientConfig, OAuthTokenResponse,
StoredCredentials,
};
use tokio::sync::Mutex;
const REALM: &str = "rmcp";
const CLIENT_ID: &str = "rmcp-client";
const CLIENT_SECRET: &str = "rmcp-secret";
fn kc_base() -> String {
std::env::var("KC_BASE").unwrap_or_else(|_| "http://localhost:8081".to_string())
}
fn token_endpoint() -> String {
format!("{}/realms/{REALM}/protocol/openid-connect/token", kc_base())
}
/// Ask Keycloak for a genuine token pair through the direct access grant, so the
/// stored credentials hold a refresh token the server will actually honor.
async fn issue_real_credentials() -> OAuthTokenResponse {
let form = format!(
"client_id={CLIENT_ID}&client_secret={CLIENT_SECRET}\
&username=alice&password=alice-pw&grant_type=password&scope=openid+profile"
);
let body = reqwest::Client::new()
.post(token_endpoint())
.header("content-type", "application/x-www-form-urlencoded")
.body(form)
.send()
.await
.expect("keycloak unreachable")
.text()
.await
.unwrap();
serde_json::from_str(&body).unwrap_or_else(|e| panic!("unexpected token response {body}: {e}"))
}
fn metadata() -> AuthorizationMetadata {
// `AuthorizationMetadata` is `#[non_exhaustive]`, so fill it field by field.
let mut metadata = AuthorizationMetadata::default();
metadata.authorization_endpoint =
format!("{}/realms/{REALM}/protocol/openid-connect/auth", kc_base());
metadata.token_endpoint = token_endpoint();
metadata
}
fn client_config() -> OAuthClientConfig {
let mut config = OAuthClientConfig::new(CLIENT_ID, "http://localhost/callback");
config.client_secret = Some(CLIENT_SECRET.to_string());
config
}
async fn manager_with_store<S: CredentialStore + 'static>(store: S) -> AuthorizationManager {
let mut manager = AuthorizationManager::new(kc_base()).await.unwrap();
manager.set_metadata(metadata());
manager.configure_client(client_config()).unwrap();
manager.set_credential_store(store);
manager
}
fn stored(client_id: &str, token: OAuthTokenResponse) -> StoredCredentials {
StoredCredentials::new(
client_id.to_string(),
Some(token),
vec!["openid".into(), "profile".into()],
None,
)
}
/// A store that serializes refreshes the way a shared on-disk store would.
#[derive(Clone, Default)]
struct GuardedStore {
inner: InMemoryCredentialStore,
lock: Arc<Mutex<()>>,
}
#[async_trait::async_trait]
impl CredentialStore for GuardedStore {
async fn load(&self) -> Result<Option<StoredCredentials>, AuthError> {
self.inner.load().await
}
async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError> {
self.inner.save(credentials).await
}
async fn clear(&self) -> Result<(), AuthError> {
self.inner.clear().await
}
async fn acquire_refresh_guard(&self) -> Result<Option<CredentialRefreshGuard>, AuthError> {
Ok(Some(CredentialRefreshGuard::new(
self.lock.clone().lock_owned().await,
)))
}
}
#[tokio::test]
#[ignore = "requires a live Keycloak"]
async fn live_refresh_rotates_the_stored_token() {
use oauth2::TokenResponse;
let issued = issue_real_credentials().await;
let original_refresh = issued.refresh_token().unwrap().secret().clone();
let store = InMemoryCredentialStore::new();
store.save(stored(CLIENT_ID, issued)).await.unwrap();
let manager = manager_with_store(store.clone()).await;
let refreshed = manager.refresh_token().await.expect("live refresh failed");
let saved = store.load().await.unwrap().unwrap();
let saved_token = saved.token_response.unwrap();
assert_ne!(
saved_token.refresh_token().unwrap().secret(),
&original_refresh,
"keycloak rotates refresh tokens, so the store must hold the new one"
);
assert_eq!(
saved_token.access_token().secret(),
refreshed.access_token().secret(),
"the saved credentials must match what the caller received"
);
assert!(
saved.granted_scopes.contains(&"openid".to_string()),
"granted scopes should come from the provider response, got: {:?}",
saved.granted_scopes
);
}
#[tokio::test]
#[ignore = "requires a live Keycloak"]
async fn live_refresh_rejects_credentials_for_another_client() {
use oauth2::TokenResponse;
let issued = issue_real_credentials().await;
let untouched_refresh = issued.refresh_token().unwrap().secret().clone();
let store = InMemoryCredentialStore::new();
store
.save(stored("some-other-client", issued))
.await
.unwrap();
let manager = manager_with_store(store.clone()).await;
let error = manager.refresh_token().await.unwrap_err();
assert!(
matches!(error, AuthError::AuthorizationRequired),
"a client mismatch must require reauthorization, got: {error:?}"
);
let saved = store.load().await.unwrap().unwrap();
assert_eq!(
saved
.token_response
.unwrap()
.refresh_token()
.unwrap()
.secret(),
&untouched_refresh,
"a rejected refresh must leave the stored token untouched"
);
}
#[tokio::test]
#[ignore = "requires a live Keycloak"]
async fn live_concurrent_refreshes_survive_refresh_token_rotation() {
use oauth2::TokenResponse;
let issued = issue_real_credentials().await;
let store = GuardedStore::default();
store.save(stored(CLIENT_ID, issued)).await.unwrap();
let first = manager_with_store(store.clone()).await;
let second = manager_with_store(store.clone()).await;
// Without the guard the second caller would reuse the refresh token the first
// one already consumed, and Keycloak would answer `invalid_grant`.
let (a, b) = tokio::join!(
tokio::spawn(async move { first.refresh_token().await }),
tokio::spawn(async move { second.refresh_token().await })
);
let a = a.unwrap().expect("first concurrent refresh failed");
let b = b.unwrap().expect("second concurrent refresh failed");
assert_ne!(
a.access_token().secret(),
b.access_token().secret(),
"each caller performs its own exchange, so the tokens must differ"
);
}
/// Deterministic proof that this realm really does invalidate a rotated refresh
/// token, which is what makes the coordination above load-bearing.
#[tokio::test]
#[ignore = "requires a live Keycloak with revokeRefreshToken enabled"]
async fn live_reusing_a_rotated_refresh_token_is_rejected() {
let issued = issue_real_credentials().await;
let first_store = InMemoryCredentialStore::new();
first_store
.save(stored(CLIENT_ID, issued.clone()))
.await
.unwrap();
manager_with_store(first_store)
.await
.refresh_token()
.await
.expect("the first refresh should succeed");
// A second caller that never saw the rotation still holds the consumed token.
let stale_store = InMemoryCredentialStore::new();
stale_store.save(stored(CLIENT_ID, issued)).await.unwrap();
let error = manager_with_store(stale_store)
.await
.refresh_token()
.await
.unwrap_err();
assert!(
matches!(error, AuthError::TokenRefreshRejected(_)),
"reusing a rotated refresh token must be rejected, got: {error:?}"
);
}

56
scripts/keycloak-oauth-fixture.sh Executable file
View file

@ -0,0 +1,56 @@
#!/usr/bin/env bash
# =============================================================================
# keycloak-oauth-fixture.sh — Live OAuth fixture for rmcp refresh tests
#
# Starts a throwaway Keycloak and provisions the realm that
# crates/rmcp/tests/test_live_oauth_refresh.rs expects. Those tests are
# #[ignore]d, so nothing here runs in CI.
#
# Provision: ./scripts/keycloak-oauth-fixture.sh
# Run tests: cargo test -p rmcp --all-features \
# --test test_live_oauth_refresh -- --ignored
# Tear down: docker rm -f kc-rmcp-test
#
# Requires Docker. Override the port with KC_BASE (default localhost:8081);
# the tests read the same variable.
# =============================================================================
set -euo pipefail
KC=${KC_BASE:-http://localhost:8081}
PORT=${KC##*:}
CONTAINER=kc-rmcp-test
docker rm -f "$CONTAINER" >/dev/null 2>&1 || true
docker run -d --name "$CONTAINER" -p "$PORT:8080" \
-e KC_BOOTSTRAP_ADMIN_USERNAME=admin -e KC_BOOTSTRAP_ADMIN_PASSWORD=admin \
quay.io/keycloak/keycloak:26.0 start-dev >/dev/null
echo "waiting for keycloak at $KC ..."
until curl -sf -o /dev/null "$KC/realms/master/.well-known/openid-configuration"; do
sleep 3
done
token=$(curl -s -X POST "$KC/realms/master/protocol/openid-connect/token" \
-d client_id=admin-cli -d username=admin -d password=admin -d grant_type=password |
python3 -c 'import sys,json;print(json.load(sys.stdin)["access_token"])')
provision() {
curl -s -X POST "$KC/admin/realms$1" \
-H "Authorization: Bearer $token" -H "Content-Type: application/json" \
-d "$2" -o /dev/null -w " ${1:-/} -> %{http_code}\n"
}
# revokeRefreshToken makes refresh tokens single-use. The concurrency test
# depends on it: without the refresh guard the second caller replays a consumed
# token and Keycloak answers invalid_grant.
provision "" '{"realm":"rmcp","enabled":true,"revokeRefreshToken":true,"refreshTokenMaxReuse":0}'
provision "/rmcp/clients" '{"clientId":"rmcp-client","secret":"rmcp-secret","publicClient":false,
"directAccessGrantsEnabled":true,"standardFlowEnabled":true,
"redirectUris":["http://localhost/callback"]}'
# The profile fields and empty requiredActions keep the direct access grant from
# failing with "Account is not fully set up".
provision "/rmcp/users" '{"username":"alice","enabled":true,"emailVerified":true,
"email":"alice@example.com","firstName":"Alice","lastName":"Example","requiredActions":[],
"credentials":[{"type":"password","value":"alice-pw","temporary":false}]}'
echo "ready"