Skip to content

Commit dde1d0a

Browse files
fix(auth): address initial scope selection review
1 parent b9b9fea commit dde1d0a

1 file changed

Lines changed: 147 additions & 52 deletions

File tree

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

Lines changed: 147 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -813,10 +813,10 @@ impl OAuthClientConfig {
813813
pub struct AuthorizationRequest {
814814
/// Redirect URI for the authorization-code flow.
815815
pub redirect_uri: String,
816-
/// Scopes to request. When empty, the SDK selects scopes from the
817-
/// server's `WWW-Authenticate` challenge, Protected Resource Metadata,
818-
/// or authorization server metadata.
816+
/// Scopes to request. When empty, the SDK uses the configured scope-selection policy.
819817
pub scopes: Vec<String>,
818+
/// Use initial-login scope selection when scopes are empty. Defaults to false.
819+
pub use_initial_scopes: bool,
820820
/// Human-readable client name, used for Dynamic Client Registration.
821821
pub client_name: Option<String>,
822822
/// Pre-registered client ID obtained from the authorization server out of
@@ -845,6 +845,7 @@ impl AuthorizationRequest {
845845
Self {
846846
redirect_uri: redirect_uri.into(),
847847
scopes: Vec::new(),
848+
use_initial_scopes: false,
848849
client_name: None,
849850
client_id: None,
850851
client_secret: None,
@@ -854,8 +855,14 @@ impl AuthorizationRequest {
854855
}
855856
}
856857

858+
/// Opt into [`AuthorizationManager::select_initial_scopes`]. Nonempty explicit scopes take precedence.
859+
pub fn with_initial_scope_selection(mut self) -> Self {
860+
self.use_initial_scopes = true;
861+
self
862+
}
863+
857864
/// Set the scopes to request. When not set, the SDK auto-selects scopes
858-
/// using its normal scope-selection policy.
865+
/// using its configured scope-selection policy.
859866
pub fn with_scopes<I, S>(mut self, scopes: I) -> Self
860867
where
861868
I: IntoIterator<Item = S>,
@@ -1125,6 +1132,14 @@ impl Default for ScopeUpgradeConfig {
11251132
}
11261133
}
11271134

1135+
#[derive(Default)]
1136+
struct InitialScopeState {
1137+
/// Initial resource challenge only; None means the resource has not been observed yet.
1138+
challenge_scopes: Option<Vec<String>>,
1139+
/// Protected resource scopes from the current discovery, including an empty result.
1140+
resource_scopes: Vec<String>,
1141+
}
1142+
11281143
/// oauth2 auth manager
11291144
pub struct AuthorizationManager {
11301145
http_client: Arc<dyn OAuthHttpClient>,
@@ -1140,10 +1155,7 @@ pub struct AuthorizationManager {
11401155
scope_upgrade_config: ScopeUpgradeConfig,
11411156
/// scopes from the initial 401 WWW-Authenticate header, used by select_scopes()
11421157
www_auth_scopes: RwLock<Vec<String>>,
1143-
/// Initial resource challenge only; None means the resource has not been observed yet.
1144-
initial_challenge_scopes: RwLock<Option<Vec<String>>>,
1145-
/// Protected resource scopes from the current discovery, including an empty result.
1146-
initial_resource_scopes: RwLock<Vec<String>>,
1158+
initial_scopes: RwLock<InitialScopeState>,
11471159
/// scopes_supported from protected resource metadata (RFC 9728)
11481160
resource_scopes: RwLock<Vec<String>>,
11491161
/// resource indicator from protected resource metadata, used for RFC 8707 `resource`
@@ -1390,8 +1402,7 @@ impl AuthorizationManager {
13901402
scope_upgrade_attempts: RwLock::new(0),
13911403
scope_upgrade_config: ScopeUpgradeConfig::default(),
13921404
www_auth_scopes: RwLock::new(Vec::new()),
1393-
initial_challenge_scopes: RwLock::new(None),
1394-
initial_resource_scopes: RwLock::new(Vec::new()),
1405+
initial_scopes: RwLock::new(InitialScopeState::default()),
13951406
resource_scopes: RwLock::new(Vec::new()),
13961407
discovered_resource: RwLock::new(None),
13971408
application_type: Some(DEFAULT_APPLICATION_TYPE.to_string()),
@@ -1547,8 +1558,7 @@ impl AuthorizationManager {
15471558
/// # }
15481559
/// ```
15491560
pub async fn resolve_metadata(&self) -> Result<AuthorizationMetadataResolution, AuthError> {
1550-
*self.initial_challenge_scopes.write().await = None;
1551-
self.initial_resource_scopes.write().await.clear();
1561+
self.reset_initial_scopes().await;
15521562
self.resolve_metadata_inner().await
15531563
}
15541564

@@ -1591,8 +1601,7 @@ impl AuthorizationManager {
15911601
};
15921602
let params = WWWAuthenticateParams::parse(www_authenticate, &self.base_url);
15931603

1594-
*self.initial_challenge_scopes.write().await = None;
1595-
self.initial_resource_scopes.write().await.clear();
1604+
self.reset_initial_scopes().await;
15961605
self.record_initial_challenge_scope(Some(&params)).await;
15971606
self.record_challenge_scope(&params).await;
15981607

@@ -1610,10 +1619,14 @@ impl AuthorizationManager {
16101619
self.resolve_metadata_inner().await
16111620
}
16121621

1622+
async fn reset_initial_scopes(&self) {
1623+
*self.initial_scopes.write().await = InitialScopeState::default();
1624+
}
1625+
16131626
async fn record_initial_challenge_scope(&self, params: Option<&WWWAuthenticateParams>) {
1614-
let mut initial_scopes = self.initial_challenge_scopes.write().await;
1615-
if initial_scopes.is_none() {
1616-
*initial_scopes = Some(
1627+
let mut initial_scopes = self.initial_scopes.write().await;
1628+
if initial_scopes.challenge_scopes.is_none() {
1629+
initial_scopes.challenge_scopes = Some(
16171630
params
16181631
.and_then(|params| params.scope.as_deref())
16191632
.map(|scope| scope.split_whitespace().map(str::to_string).collect())
@@ -1978,13 +1991,15 @@ impl AuthorizationManager {
19781991
/// Select scopes for initial authorization after resolving and setting metadata.
19791992
///
19801993
/// Prefer the initial resource challenge, then protected resource metadata. If
1981-
/// neither supplies scopes, retain the authorization-server metadata and caller
1982-
/// defaults as compatibility fallbacks. Append `offline_access` when advertised
1983-
/// by the authorization server, as with [`Self::select_scopes`].
1994+
/// neither supplies scopes, use the caller-provided `default_scopes`. Pass an
1995+
/// empty slice to omit scopes when discovery supplies none. Authorization-server
1996+
/// scopes are not used as a fallback. When the selected scopes are nonempty,
1997+
/// append `offline_access` if advertised by the authorization server, as with
1998+
/// [`Self::select_scopes`].
19841999
///
19852000
/// Unlike [`Self::select_scopes`], this does not accumulate previously requested
1986-
/// scopes or challenges from metadata endpoints. Callers with explicit scopes
1987-
/// should continue to pass those scopes directly to authorization instead.
2001+
/// scopes or challenges from metadata endpoints. Callers that need to override
2002+
/// discovered scopes should pass those scopes directly to authorization instead.
19882003
///
19892004
/// # Examples
19902005
///
@@ -1999,17 +2014,20 @@ impl AuthorizationManager {
19992014
/// # }
20002015
/// ```
20012016
pub async fn select_initial_scopes(&self, default_scopes: &[&str]) -> Vec<String> {
2002-
let mut scopes = self
2003-
.initial_challenge_scopes
2004-
.read()
2005-
.await
2006-
.clone()
2007-
.unwrap_or_default();
2008-
if scopes.is_empty() {
2009-
scopes = self.initial_resource_scopes.read().await.clone();
2010-
}
2017+
let mut scopes = {
2018+
let initial_scopes = self.initial_scopes.read().await;
2019+
initial_scopes
2020+
.challenge_scopes
2021+
.as_ref()
2022+
.filter(|scopes| !scopes.is_empty())
2023+
.unwrap_or(&initial_scopes.resource_scopes)
2024+
.clone()
2025+
};
20112026
if scopes.is_empty() {
2012-
scopes = self.select_fallback_scopes(default_scopes);
2027+
scopes = default_scopes
2028+
.iter()
2029+
.map(|scope| scope.to_string())
2030+
.collect();
20132031
}
20142032
let mut scopes = Self::dedup_scopes(scopes);
20152033
self.add_offline_access_if_supported(&mut scopes);
@@ -2688,7 +2706,7 @@ impl AuthorizationManager {
26882706
.await
26892707
.replace(resource_metadata.resource.clone().unwrap_or_default());
26902708

2691-
*self.initial_resource_scopes.write().await = resource_metadata
2709+
self.initial_scopes.write().await.resource_scopes = resource_metadata
26922710
.scopes_supported
26932711
.clone()
26942712
.unwrap_or_default();
@@ -3661,7 +3679,7 @@ impl AuthorizationSession {
36613679
///
36623680
/// The manager must already have discovered authorization server metadata.
36633681
/// If `request.scopes` is empty, scopes are selected using the SDK's
3664-
/// normal scope-selection policy.
3682+
/// configured scope-selection policy.
36653683
///
36663684
/// On failure, the manager is returned alongside the error so callers can
36673685
/// retry without losing the original configuration and stores.
@@ -3681,7 +3699,11 @@ impl AuthorizationSession {
36813699
}
36823700

36833701
if request.scopes.is_empty() {
3684-
request.scopes = auth_manager.select_scopes(None, &[]);
3702+
request.scopes = if request.use_initial_scopes {
3703+
auth_manager.select_initial_scopes(&[]).await
3704+
} else {
3705+
auth_manager.select_scopes(None, &[])
3706+
};
36853707
} else {
36863708
auth_manager.add_offline_access_if_supported(&mut request.scopes);
36873709
}
@@ -4015,7 +4037,7 @@ impl OAuthState {
40154037
/// advertises a `registration_endpoint`
40164038
///
40174039
/// If `request.scopes` is empty, scopes are selected using the SDK's
4018-
/// normal scope-selection policy.
4040+
/// configured scope-selection policy.
40194041
///
40204042
/// On failure, the state returns to `Unauthorized` so callers can retry
40214043
/// without losing the original configuration and stores.
@@ -4262,9 +4284,9 @@ mod tests {
42624284
AuthError, AuthorizationCallback, AuthorizationManager, AuthorizationMetadata,
42634285
AuthorizationMetadataSource, AuthorizationRequest, AuthorizationSession,
42644286
CredentialRefreshGuard, CredentialStore, InMemoryCredentialStore, InMemoryStateStore,
4265-
OAuthClientConfig, OAuthHttpClient, OAuthHttpClientError, OAuthHttpClientFuture,
4266-
OAuthHttpRedirectPolicy, OAuthHttpRequest, ScopeUpgradeConfig, StateStore,
4267-
StoredAuthorizationState, is_https_url,
4287+
InitialScopeState, OAuthClientConfig, OAuthHttpClient, OAuthHttpClientError,
4288+
OAuthHttpClientFuture, OAuthHttpRedirectPolicy, OAuthHttpRequest, ScopeUpgradeConfig,
4289+
StateStore, StoredAuthorizationState, is_https_url,
42684290
};
42694291
use crate::transport::auth::VendorExtraTokenFields;
42704292

@@ -5964,6 +5986,55 @@ mod tests {
59645986
assert_eq!(query.get("scope").unwrap(), "read write offline_access");
59655987
}
59665988

5989+
#[rstest]
5990+
#[case::legacy(false, Some("read"), Some(&["read", "write"][..]), &[], Some("read write offline_access"))]
5991+
#[case::initial(true, Some("read"), Some(&["read", "write"][..]), &[], Some("read offline_access"))]
5992+
#[case::resource(true, None, Some(&["read", "write"][..]), &[], Some("read write offline_access"))]
5993+
#[case::missing(true, None, None, &[], None)]
5994+
#[case::explicit(true, Some("read"), Some(&["read", "write"][..]), &["custom"], Some("custom offline_access"))]
5995+
#[tokio::test]
5996+
async fn oauth_state_uses_requested_scope_selection(
5997+
#[case] opt_in: bool,
5998+
#[case] challenge_scope: Option<&str>,
5999+
#[case] resource_scopes: Option<&[&str]>,
6000+
#[case] explicit_scopes: &[&str],
6001+
#[case] expected_scope: Option<&str>,
6002+
) {
6003+
let mut challenge = r#"Bearer resource_metadata="/metadata""#.to_string();
6004+
if let Some(scope) = challenge_scope {
6005+
challenge.push_str(&format!(r#", scope="{scope}""#));
6006+
}
6007+
let mut resource_metadata = serde_json::json!({
6008+
"resource": "https://mcp.example.com/mcp",
6009+
"authorization_servers": ["https://auth.example.com"]
6010+
});
6011+
if let Some(scopes) = resource_scopes {
6012+
resource_metadata["scopes_supported"] = serde_json::json!(scopes);
6013+
}
6014+
let client = RecordingOAuthHttpClient::with_responses(vec![
6015+
challenge_response(&challenge),
6016+
http_response(200, resource_metadata),
6017+
preregistered_as_metadata_response(),
6018+
]);
6019+
let mut state = super::OAuthState::new_with_oauth_http_client(
6020+
"https://mcp.example.com/mcp",
6021+
Arc::new(client),
6022+
)
6023+
.await
6024+
.unwrap();
6025+
6026+
let mut request = AuthorizationRequest::new("http://localhost:8080/callback")
6027+
.with_preregistered_client("preregistered-client")
6028+
.with_scopes(explicit_scopes.iter().copied());
6029+
if opt_in {
6030+
request = request.with_initial_scope_selection();
6031+
}
6032+
state.start_authorization(request).await.unwrap();
6033+
6034+
let query = auth_url_query(&state.get_authorization_url().await.unwrap());
6035+
assert_eq!(query.get("scope").map(String::as_str), expected_scope);
6036+
}
6037+
59676038
#[tokio::test]
59686039
async fn preregistered_client_uses_explicit_scopes_and_adds_offline_access() {
59696040
let client = RecordingOAuthHttpClient::with_responses(preregistered_discovery_responses());
@@ -7968,14 +8039,9 @@ mod tests {
79688039
assert_eq!(scopes, ["read"]);
79698040
}
79708041

7971-
#[rstest]
7972-
#[case::initial_scope(Some("read"), vec!["read"])]
7973-
#[case::no_initial_scope(None, vec!["read", "write"])]
7974-
#[tokio::test]
7975-
async fn select_initial_scopes_ignores_later_metadata_challenges(
7976-
#[case] initial_scope: Option<&str>,
7977-
#[case] expected: Vec<&str>,
7978-
) {
8042+
async fn manager_with_later_metadata_challenge(
8043+
initial_scope: Option<&str>,
8044+
) -> AuthorizationManager {
79798045
let initial_response = match initial_scope {
79808046
Some(scope) => challenge_response(&format!(r#"Bearer scope="{scope}""#)),
79818047
None => empty_response(404),
@@ -7995,7 +8061,31 @@ mod tests {
79958061
let resolution = manager.resolve_metadata().await.unwrap();
79968062
manager.set_metadata(resolution.metadata);
79978063

8064+
manager
8065+
}
8066+
8067+
#[rstest]
8068+
#[case::initial_scope(Some("read"), vec!["read"])]
8069+
#[case::no_initial_scope(None, vec!["read", "write"])]
8070+
#[tokio::test]
8071+
async fn select_initial_scopes_ignores_later_metadata_challenges(
8072+
#[case] initial_scope: Option<&str>,
8073+
#[case] expected: Vec<&str>,
8074+
) {
8075+
let manager = manager_with_later_metadata_challenge(initial_scope).await;
8076+
79988077
assert_eq!(manager.select_initial_scopes(&[]).await, expected);
8078+
}
8079+
8080+
#[rstest]
8081+
#[case::initial_scope(Some("read"))]
8082+
#[case::no_initial_scope(None)]
8083+
#[tokio::test]
8084+
async fn select_scopes_preserves_later_metadata_challenges(
8085+
#[case] initial_scope: Option<&str>,
8086+
) {
8087+
let manager = manager_with_later_metadata_challenge(initial_scope).await;
8088+
79998089
assert!(
80008090
manager
80018091
.select_scopes(None, &[])
@@ -8028,10 +8118,13 @@ mod tests {
80288118
#[rstest]
80298119
#[case::challenge(&["read", "read"], &["write"], &["admin"], &["default"], &["read"])]
80308120
#[case::resource(&[], &["read", "read"], &["admin"], &["default"], &["read"])]
8031-
#[case::authorization_server(&[], &[], &["admin", "admin"], &["default"], &["admin"])]
8121+
#[case::defaults_ignore_authorization_server(&[], &[], &["admin", "admin"], &["default"], &["default"])]
8122+
#[case::authorization_server_ignored(&[], &[], &["admin"], &[], &[])]
80328123
#[case::defaults(&[], &[], &[], &["default", "default"], &["default"])]
80338124
#[case::empty(&[], &[], &[], &[], &[])]
80348125
#[case::offline_access(&["read"], &["write"], &["offline_access"], &[], &["read", "offline_access"])]
8126+
#[case::offline_access_without_selected_scopes(&[], &[], &["admin", "offline_access"], &[], &[])]
8127+
#[case::defaults_with_offline_access(&[], &[], &["admin", "offline_access"], &["default", "default"], &["default", "offline_access"])]
80358128
#[tokio::test]
80368129
async fn select_initial_scopes_uses_fallbacks_and_preserves_refresh_support(
80378130
#[case] challenge: &[&str],
@@ -8046,17 +8139,19 @@ mod tests {
80468139
..Default::default()
80478140
}))
80488141
.await;
8049-
*manager.initial_challenge_scopes.write().await = Some(strings(challenge));
8050-
*manager.initial_resource_scopes.write().await = strings(resource);
8142+
*manager.initial_scopes.write().await = InitialScopeState {
8143+
challenge_scopes: Some(strings(challenge)),
8144+
resource_scopes: strings(resource),
8145+
};
80518146
*manager.current_scopes.write().await = vec!["previous".to_string()];
80528147

80538148
assert_eq!(manager.select_initial_scopes(defaults).await, expected);
80548149
}
80558150

80568151
#[rstest]
80578152
#[case::challenge(Some("read"), Some(&["read", "write"][..]), &["read"], &["read", "write"])]
8058-
#[case::missing_resource_scopes(None, None, &["read", "write"], &["admin"])]
8059-
#[case::empty_resource_scopes(None, Some(&[][..]), &["read", "write"], &["admin"])]
8153+
#[case::missing_resource_scopes(None, None, &["read", "write"], &[])]
8154+
#[case::empty_resource_scopes(None, Some(&[][..]), &["read", "write"], &[])]
80608155
#[tokio::test]
80618156
async fn select_initial_scopes_resets_on_repeated_discovery(
80628157
#[case] initial_scope: Option<&str>,

0 commit comments

Comments
 (0)