refactor: make oauth discovery reactive (#1052)
This commit is contained in:
parent
92c0793247
commit
c284607794
6 changed files with 581 additions and 241 deletions
|
|
@ -195,25 +195,73 @@ const REDIRECT_URI: &str = "http://localhost:3000/callback";
|
|||
const SCOPE_STEP_UP_INITIAL_SCOPES: &[&str] = &["mcp:basic"];
|
||||
const SCOPE_STEP_UP_ESCALATED_SCOPES: &[&str] = &["mcp:basic", "mcp:write"];
|
||||
|
||||
/// Perform the headless OAuth authorization-code flow.
|
||||
/// Attempt the real connection unauthenticated and return the server's
|
||||
/// `WWW-Authenticate` challenge from the 401 — the reactive discovery
|
||||
/// trigger.
|
||||
///
|
||||
/// 1. Discover metadata, register (or use CIMD), get auth URL
|
||||
/// 2. Fetch the auth URL with redirect:manual → extract code from Location header
|
||||
/// 3. Exchange code for token
|
||||
/// 4. Return an `AuthClient` wrapping `reqwest::Client`
|
||||
/// `None` (server accepted the unauthenticated connection, which is then
|
||||
/// closed cleanly) is a legitimate outcome, not an error: the scope-step-up
|
||||
/// and scope-retry-limit mocks allow unauthenticated `initialize` and only
|
||||
/// enforce authorization on tool calls.
|
||||
async fn initialize_challenge(
|
||||
server_url: &str,
|
||||
lifecycle: ClientLifecycleMode,
|
||||
) -> anyhow::Result<Option<String>> {
|
||||
let transport = StreamableHttpClientTransport::from_uri(server_url);
|
||||
match BasicClientHandler
|
||||
.serve_with_lifecycle(transport, lifecycle)
|
||||
.await
|
||||
{
|
||||
Ok(client) => {
|
||||
client.cancel().await.ok();
|
||||
Ok(None)
|
||||
}
|
||||
Err(error) => match error.auth_challenge() {
|
||||
Some(challenge) => Ok(Some(challenge.to_string())),
|
||||
None => Err(error.into()),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn with_optional_challenge(
|
||||
request: AuthorizationRequest,
|
||||
challenge: Option<String>,
|
||||
) -> AuthorizationRequest {
|
||||
match challenge {
|
||||
Some(challenge) => request.with_challenge(challenge),
|
||||
None => request,
|
||||
}
|
||||
}
|
||||
|
||||
/// Perform the headless OAuth authorization-code flow, reactively:
|
||||
///
|
||||
/// 1. Attempt the real connection; take the 401's WWW-Authenticate challenge
|
||||
/// 2. Discover from the challenge, register (or use CIMD), get auth URL
|
||||
/// 3. Fetch the auth URL with redirect:manual → extract code from Location header
|
||||
/// 4. Exchange code for token
|
||||
/// 5. Return an `AuthClient` wrapping `reqwest::Client`
|
||||
async fn perform_oauth_flow(
|
||||
server_url: &str,
|
||||
_ctx: &ConformanceContext,
|
||||
) -> anyhow::Result<AuthClient<reqwest::Client>> {
|
||||
// Always the discover lifecycle here (not `conformance_lifecycle()`):
|
||||
// this flow serves `run_auth_client`, whose 2026-07-28 auth mocks require
|
||||
// the per-request MCP-Protocol-Version negotiation.
|
||||
let challenge = initialize_challenge(
|
||||
server_url,
|
||||
ClientLifecycleMode::Discover {
|
||||
preferred_versions: preferred_protocol_versions(),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
let mut oauth = OAuthState::new(server_url, None).await?;
|
||||
|
||||
// Discover + register + get auth URL
|
||||
let request = AuthorizationRequest::new(REDIRECT_URI)
|
||||
.with_client_name("conformance-client")
|
||||
.with_client_metadata_url(CIMD_CLIENT_METADATA_URL);
|
||||
oauth
|
||||
.start_authorization(
|
||||
AuthorizationRequest::new(REDIRECT_URI)
|
||||
.with_client_name("conformance-client")
|
||||
.with_client_metadata_url(CIMD_CLIENT_METADATA_URL),
|
||||
)
|
||||
.start_authorization(with_optional_challenge(request, challenge))
|
||||
.await?;
|
||||
|
||||
let auth_url = oauth.get_authorization_url().await?;
|
||||
|
|
@ -255,14 +303,14 @@ async fn perform_oauth_flow_preregistered(
|
|||
client_id: &str,
|
||||
client_secret: &str,
|
||||
) -> anyhow::Result<AuthClient<reqwest::Client>> {
|
||||
let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?;
|
||||
let mut oauth = OAuthState::new(server_url, None).await?;
|
||||
|
||||
let request = AuthorizationRequest::new(REDIRECT_URI)
|
||||
.with_preregistered_client(client_id)
|
||||
.with_client_secret(client_secret);
|
||||
oauth
|
||||
.start_authorization(
|
||||
AuthorizationRequest::new(REDIRECT_URI)
|
||||
.with_preregistered_client(client_id)
|
||||
.with_client_secret(client_secret),
|
||||
)
|
||||
.start_authorization(with_optional_challenge(request, challenge))
|
||||
.await?;
|
||||
|
||||
let auth_url = oauth.get_authorization_url().await?;
|
||||
|
|
@ -325,14 +373,14 @@ async fn run_auth_scope_step_up_client(
|
|||
server_url: &str,
|
||||
_ctx: &ConformanceContext,
|
||||
) -> anyhow::Result<()> {
|
||||
let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?;
|
||||
let mut oauth = OAuthState::new(server_url, None).await?;
|
||||
let request = AuthorizationRequest::new(REDIRECT_URI)
|
||||
.with_scopes(SCOPE_STEP_UP_INITIAL_SCOPES.iter().copied())
|
||||
.with_client_name("conformance-client")
|
||||
.with_client_metadata_url(CIMD_CLIENT_METADATA_URL);
|
||||
oauth
|
||||
.start_authorization(
|
||||
AuthorizationRequest::new(REDIRECT_URI)
|
||||
.with_scopes(SCOPE_STEP_UP_INITIAL_SCOPES.iter().copied())
|
||||
.with_client_name("conformance-client")
|
||||
.with_client_metadata_url(CIMD_CLIENT_METADATA_URL),
|
||||
)
|
||||
.start_authorization(with_optional_challenge(request, challenge))
|
||||
.await?;
|
||||
|
||||
let auth_url = oauth.get_authorization_url().await?;
|
||||
|
|
@ -427,15 +475,15 @@ async fn run_auth_scope_retry_limit_client(
|
|||
) -> anyhow::Result<()> {
|
||||
let max_retries = 3u32;
|
||||
let mut attempt = 0u32;
|
||||
let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?;
|
||||
|
||||
loop {
|
||||
let mut oauth = OAuthState::new(server_url, None).await?;
|
||||
let request = AuthorizationRequest::new(REDIRECT_URI)
|
||||
.with_client_name("conformance-client")
|
||||
.with_client_metadata_url(CIMD_CLIENT_METADATA_URL);
|
||||
oauth
|
||||
.start_authorization(
|
||||
AuthorizationRequest::new(REDIRECT_URI)
|
||||
.with_client_name("conformance-client")
|
||||
.with_client_metadata_url(CIMD_CLIENT_METADATA_URL),
|
||||
)
|
||||
.start_authorization(with_optional_challenge(request, challenge.clone()))
|
||||
.await?;
|
||||
let auth_url = oauth.get_authorization_url().await?;
|
||||
let callback = headless_authorize(&auth_url).await?;
|
||||
|
|
@ -509,7 +557,10 @@ async fn migration_token(
|
|||
return Ok(manager.get_access_token().await?);
|
||||
}
|
||||
|
||||
let resolution = manager.resolve_metadata().await?;
|
||||
let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?;
|
||||
let resolution = manager
|
||||
.resolve_metadata_from_challenge(challenge.as_deref())
|
||||
.await?;
|
||||
manager.set_metadata(resolution.metadata);
|
||||
manager
|
||||
.register_client("conformance-client", REDIRECT_URI, &[])
|
||||
|
|
@ -609,7 +660,10 @@ async fn run_client_credentials_basic(
|
|||
.unwrap_or("conformance-test-secret");
|
||||
|
||||
let mut manager = AuthorizationManager::new(server_url).await?;
|
||||
let resolution = manager.resolve_metadata().await?;
|
||||
let challenge = initialize_challenge(server_url, conformance_lifecycle()).await?;
|
||||
let resolution = manager
|
||||
.resolve_metadata_from_challenge(challenge.as_deref())
|
||||
.await?;
|
||||
let token_endpoint = resolution.metadata.token_endpoint.clone();
|
||||
manager.set_metadata(resolution.metadata);
|
||||
|
||||
|
|
|
|||
|
|
@ -86,6 +86,31 @@ impl ClientInitializeError {
|
|||
context: context.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// The `WWW-Authenticate` challenge from the 401/403 the transport hit
|
||||
/// during initialization, if that is why initialization failed.
|
||||
///
|
||||
/// This is the trigger of the reactive OAuth flow: feed the challenge to
|
||||
/// `AuthorizationRequest::with_challenge` to authorize, then reconnect.
|
||||
#[cfg(feature = "transport-streamable-http-client")]
|
||||
pub fn auth_challenge(&self) -> Option<&str> {
|
||||
use crate::transport::streamable_http_client::{AuthRequiredError, InsufficientScopeError};
|
||||
|
||||
let Self::TransportError { error, .. } = self else {
|
||||
return None;
|
||||
};
|
||||
let mut source: Option<&(dyn std::error::Error + 'static)> = Some(error.error.as_ref());
|
||||
while let Some(current) = source {
|
||||
if let Some(auth_required) = current.downcast_ref::<AuthRequiredError>() {
|
||||
return Some(&auth_required.www_authenticate_header);
|
||||
}
|
||||
if let Some(insufficient_scope) = current.downcast_ref::<InsufficientScopeError>() {
|
||||
return Some(&insufficient_scope.www_authenticate_header);
|
||||
}
|
||||
source = current.source();
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper function to get the next message from the stream
|
||||
|
|
|
|||
|
|
@ -31,12 +31,6 @@ use crate::transport::common::http_header::HEADER_MCP_PROTOCOL_VERSION;
|
|||
const DEFAULT_HTTP_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
const MAX_OAUTH_HTTP_RESPONSE_BODY_BYTES: usize = 1024 * 1024;
|
||||
const MAX_OAUTH_DISCOVERY_REDIRECTS: usize = 10;
|
||||
const RESOURCE_METADATA_POST_PROBE_BODY: &str = concat!(
|
||||
r#"{"jsonrpc":"2.0","id":"auth-discovery","method":"initialize","params":{"#,
|
||||
r#""protocolVersion":"2024-11-05","capabilities":{},"#,
|
||||
r#""clientInfo":{"name":"rmcp-auth-discovery","version":"0.0.0"}}"#,
|
||||
r#"}"#
|
||||
);
|
||||
const CLOUD_METADATA_HOSTS: &[&str] = &[
|
||||
"metadata",
|
||||
"metadata.google.internal",
|
||||
|
|
@ -634,6 +628,12 @@ pub struct WWWAuthenticateParams {
|
|||
}
|
||||
|
||||
impl WWWAuthenticateParams {
|
||||
/// Parse a `WWW-Authenticate` header value, resolving a relative
|
||||
/// `resource_metadata` URL against `base_url`.
|
||||
pub fn parse(header: &str, base_url: &Url) -> Self {
|
||||
AuthorizationManager::extract_www_authenticate_params(header, base_url)
|
||||
}
|
||||
|
||||
/// check if this is an insufficient_scope error
|
||||
pub fn is_insufficient_scope(&self) -> bool {
|
||||
self.error.as_deref() == Some("insufficient_scope")
|
||||
|
|
@ -736,6 +736,11 @@ pub struct AuthorizationRequest {
|
|||
/// OIDC Dynamic Client Registration `application_type` (SEP-837),
|
||||
/// e.g. `"native"` or `"web"`.
|
||||
pub application_type: Option<String>,
|
||||
/// `WWW-Authenticate` header value from a real request's 401 response.
|
||||
/// When set, discovery is seeded from the challenge (its
|
||||
/// `resource_metadata` URL and `scope` hint) instead of probing the
|
||||
/// server — the reactive discovery path.
|
||||
pub challenge: Option<String>,
|
||||
}
|
||||
|
||||
impl AuthorizationRequest {
|
||||
|
|
@ -750,6 +755,7 @@ impl AuthorizationRequest {
|
|||
client_secret: None,
|
||||
client_metadata_url: None,
|
||||
application_type: None,
|
||||
challenge: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -805,6 +811,13 @@ impl AuthorizationRequest {
|
|||
self.application_type = Some(application_type.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Seed discovery from the `WWW-Authenticate` header value of a real
|
||||
/// request's 401 response instead of probing the server.
|
||||
pub fn with_challenge(mut self, www_authenticate: impl Into<String>) -> Self {
|
||||
self.challenge = Some(www_authenticate.into());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
// add type aliases for oauth2 types
|
||||
|
|
@ -1419,6 +1432,52 @@ impl AuthorizationManager {
|
|||
})
|
||||
}
|
||||
|
||||
/// Resolve authorization server metadata starting from the
|
||||
/// `WWW-Authenticate` challenge of a real request's 401 response — the
|
||||
/// reactive discovery path, matching the TypeScript and Python SDKs.
|
||||
///
|
||||
/// Seeds scope selection with the challenge's `scope` hint and prefers
|
||||
/// the challenge's `resource_metadata` URL; falls back to
|
||||
/// [`resolve_metadata`](Self::resolve_metadata) when the challenge is
|
||||
/// `None` or carries no usable metadata pointer.
|
||||
pub async fn resolve_metadata_from_challenge(
|
||||
&self,
|
||||
www_authenticate: Option<&str>,
|
||||
) -> Result<AuthorizationMetadataResolution, AuthError> {
|
||||
let Some(www_authenticate) = www_authenticate else {
|
||||
return self.resolve_metadata().await;
|
||||
};
|
||||
let params = WWWAuthenticateParams::parse(www_authenticate, &self.base_url);
|
||||
|
||||
self.record_challenge_scope(¶ms).await;
|
||||
|
||||
if let Some(resource_metadata_url) = ¶ms.resource_metadata_url
|
||||
&& let Some(metadata) = self
|
||||
.discover_oauth_server_from_resource_metadata_url(resource_metadata_url)
|
||||
.await?
|
||||
{
|
||||
return Ok(AuthorizationMetadataResolution {
|
||||
metadata,
|
||||
source: AuthorizationMetadataSource::ProtectedResourceMetadata,
|
||||
});
|
||||
}
|
||||
|
||||
self.resolve_metadata().await
|
||||
}
|
||||
|
||||
/// Store a challenge's `scope` hint for later scope selection.
|
||||
async fn record_challenge_scope(&self, params: &WWWAuthenticateParams) {
|
||||
let Some(scope) = ¶ms.scope else {
|
||||
return;
|
||||
};
|
||||
let scopes: Vec<String> = scope.split_whitespace().map(str::to_string).collect();
|
||||
if scopes.is_empty() {
|
||||
return;
|
||||
}
|
||||
debug!("WWW-Authenticate challenge contains scope: {scope}");
|
||||
*self.www_auth_scopes.write().await = scopes;
|
||||
}
|
||||
|
||||
fn legacy_authorization_metadata(base_url: &Url) -> AuthorizationMetadata {
|
||||
let endpoint = |path: &str| {
|
||||
let mut url = base_url.clone();
|
||||
|
|
@ -2056,7 +2115,7 @@ impl AuthorizationManager {
|
|||
/// refresh token or the server rejected it, return `AuthorizationRequired`
|
||||
/// so the caller can re-prompt the user. Infrastructure errors (e.g. store
|
||||
/// I/O failures, misconfigured client) are propagated as-is.
|
||||
async fn try_refresh_or_reauth(&self) -> Result<String, AuthError> {
|
||||
pub(crate) async fn try_refresh_or_reauth(&self) -> Result<String, AuthError> {
|
||||
match self.refresh_token().await {
|
||||
Ok(new_creds) => {
|
||||
tracing::info!("Refreshed access token.");
|
||||
|
|
@ -2328,12 +2387,19 @@ impl AuthorizationManager {
|
|||
async fn discover_oauth_server_via_resource_metadata(
|
||||
&self,
|
||||
) -> Result<Option<AuthorizationMetadata>, AuthError> {
|
||||
let Some(resource_metadata_url) = self.discover_resource_metadata_url().await? else {
|
||||
let Some(resource_metadata_url) = self.discover_resource_metadata_url().await else {
|
||||
return Ok(None);
|
||||
};
|
||||
self.discover_oauth_server_from_resource_metadata_url(&resource_metadata_url)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn discover_oauth_server_from_resource_metadata_url(
|
||||
&self,
|
||||
resource_metadata_url: &Url,
|
||||
) -> Result<Option<AuthorizationMetadata>, AuthError> {
|
||||
let Some(resource_metadata) = self
|
||||
.fetch_resource_metadata_from_url(&resource_metadata_url)
|
||||
.fetch_resource_metadata_from_url(resource_metadata_url)
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
|
|
@ -2443,11 +2509,10 @@ impl AuthorizationManager {
|
|||
&& Self::is_same_origin(&root_resource, &path_resource)
|
||||
}
|
||||
|
||||
async fn discover_resource_metadata_url(&self) -> Result<Option<Url>, AuthError> {
|
||||
if let Ok(Some(resource_metadata_url)) =
|
||||
self.fetch_resource_metadata_url(&self.base_url, true).await
|
||||
async fn discover_resource_metadata_url(&self) -> Option<Url> {
|
||||
if let Some(resource_metadata_url) = self.probe_resource_metadata_url(&self.base_url).await
|
||||
{
|
||||
return Ok(Some(resource_metadata_url));
|
||||
return Some(resource_metadata_url);
|
||||
}
|
||||
|
||||
// If the primary URL doesn't use WWW-Authenticate, try oauth-protected-resource discovery.
|
||||
|
|
@ -2459,86 +2524,42 @@ impl AuthorizationManager {
|
|||
discovery_url.set_query(None);
|
||||
discovery_url.set_fragment(None);
|
||||
discovery_url.set_path(&candidate_path);
|
||||
if let Ok(Some(resource_metadata_url)) = self
|
||||
.fetch_resource_metadata_url(&discovery_url, false)
|
||||
.await
|
||||
if let Some(resource_metadata_url) =
|
||||
self.probe_resource_metadata_url(&discovery_url).await
|
||||
{
|
||||
return Ok(Some(resource_metadata_url));
|
||||
return Some(resource_metadata_url);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
None
|
||||
}
|
||||
|
||||
/// Extract the resource metadata url from the WWW-Authenticate header value.
|
||||
/// Probe `url` with a GET, extracting the resource metadata url from a
|
||||
/// 200 (the url itself is the metadata document) or from a 401's
|
||||
/// WWW-Authenticate header value.
|
||||
/// https://www.rfc-editor.org/rfc/rfc9728.html#name-use-of-www-authenticate-for
|
||||
async fn fetch_resource_metadata_url(
|
||||
&self,
|
||||
url: &Url,
|
||||
allow_post_probe: bool,
|
||||
) -> Result<Option<Url>, AuthError> {
|
||||
async fn probe_resource_metadata_url(&self, url: &Url) -> Option<Url> {
|
||||
let response = match self.discovery_get(url).await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
debug!("resource metadata probe failed: {}", e);
|
||||
return Ok(None);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
match response.status() {
|
||||
StatusCode::OK => Ok(Some(url.clone())),
|
||||
StatusCode::UNAUTHORIZED => Ok(self
|
||||
.extract_resource_metadata_url_from_www_authenticate(&response)
|
||||
.await),
|
||||
StatusCode::NOT_FOUND | StatusCode::METHOD_NOT_ALLOWED if allow_post_probe => {
|
||||
self.fetch_resource_metadata_url_with_post_probe(url).await
|
||||
StatusCode::OK => Some(url.clone()),
|
||||
StatusCode::UNAUTHORIZED => {
|
||||
self.extract_resource_metadata_url_from_www_authenticate(&response)
|
||||
.await
|
||||
}
|
||||
status => {
|
||||
debug!("resource metadata probe returned unexpected status: {status}");
|
||||
Ok(None)
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_resource_metadata_url_with_post_probe(
|
||||
&self,
|
||||
url: &Url,
|
||||
) -> Result<Option<Url>, AuthError> {
|
||||
let request = oauth2::http::Request::builder()
|
||||
.method("POST")
|
||||
.uri(url.as_str())
|
||||
.header(HEADER_MCP_PROTOCOL_VERSION, "2024-11-05")
|
||||
.header(CONTENT_TYPE, "application/json")
|
||||
.body(RESOURCE_METADATA_POST_PROBE_BODY.as_bytes().to_vec())
|
||||
.map_err(|error| AuthError::InternalError(error.to_string()))?;
|
||||
let response = match self
|
||||
.http_client
|
||||
.execute(OAuthHttpRequest::new(
|
||||
request,
|
||||
OAuthHttpRedirectPolicy::Stop,
|
||||
))
|
||||
.await
|
||||
{
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
debug!("resource metadata POST probe failed: {}", error);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
|
||||
if response.status() != StatusCode::UNAUTHORIZED {
|
||||
debug!(
|
||||
"resource metadata POST probe returned unexpected status: {}",
|
||||
response.status()
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(self
|
||||
.extract_resource_metadata_url_from_www_authenticate(&response)
|
||||
.await)
|
||||
}
|
||||
|
||||
async fn extract_resource_metadata_url_from_www_authenticate(
|
||||
&self,
|
||||
response: &HttpResponse,
|
||||
|
|
@ -2549,14 +2570,9 @@ impl AuthorizationManager {
|
|||
continue;
|
||||
};
|
||||
let params = Self::extract_www_authenticate_params(value_str, &self.base_url);
|
||||
if let Some(url) = params.resource_metadata_url {
|
||||
if let Some(scope) = ¶ms.scope {
|
||||
debug!("WWW-Authenticate header contains scope: {}", scope);
|
||||
let scopes: Vec<String> =
|
||||
scope.split_whitespace().map(|s| s.to_string()).collect();
|
||||
*self.www_auth_scopes.write().await = scopes;
|
||||
}
|
||||
parsed_url = Some(url);
|
||||
if params.resource_metadata_url.is_some() {
|
||||
self.record_challenge_scope(¶ms).await;
|
||||
parsed_url = params.resource_metadata_url;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
|
@ -3432,7 +3448,7 @@ impl OAuthState {
|
|||
)
|
||||
}
|
||||
|
||||
async fn placeholder(&self) -> Result<Self, AuthError> {
|
||||
async fn placeholder_state(&self) -> Result<Self, AuthError> {
|
||||
let (http_client, refresh_redirect_policy) = self.oauth_http_client_config();
|
||||
Ok(OAuthState::Unauthorized(
|
||||
AuthorizationManager::new_inner(
|
||||
|
|
@ -3548,7 +3564,7 @@ impl OAuthState {
|
|||
&mut self,
|
||||
request: AuthorizationRequest,
|
||||
) -> Result<(), AuthError> {
|
||||
let placeholder = self.placeholder().await?;
|
||||
let placeholder = self.placeholder_state().await?;
|
||||
let old = std::mem::replace(self, placeholder);
|
||||
let OAuthState::Unauthorized(mut manager) = old else {
|
||||
*self = old;
|
||||
|
|
@ -3557,7 +3573,10 @@ impl OAuthState {
|
|||
));
|
||||
};
|
||||
debug!("start discovery");
|
||||
let metadata = match manager.resolve_metadata().await {
|
||||
let resolution = manager
|
||||
.resolve_metadata_from_challenge(request.challenge.as_deref())
|
||||
.await;
|
||||
let metadata = match resolution {
|
||||
Ok(resolution) => resolution.metadata,
|
||||
Err(e) => {
|
||||
*self = OAuthState::Unauthorized(manager);
|
||||
|
|
@ -3580,7 +3599,7 @@ impl OAuthState {
|
|||
|
||||
/// complete authorization
|
||||
pub async fn complete_authorization(&mut self) -> Result<(), AuthError> {
|
||||
let placeholder = self.placeholder().await?;
|
||||
let placeholder = self.placeholder_state().await?;
|
||||
if let OAuthState::Session(session) = std::mem::replace(self, placeholder) {
|
||||
*self = OAuthState::Authorized(session.auth_manager);
|
||||
Ok(())
|
||||
|
|
@ -3590,7 +3609,7 @@ impl OAuthState {
|
|||
}
|
||||
/// covert to authorized http client
|
||||
pub async fn to_authorized_http_client(&mut self) -> Result<(), AuthError> {
|
||||
let placeholder = self.placeholder().await?;
|
||||
let placeholder = self.placeholder_state().await?;
|
||||
if let OAuthState::Authorized(manager) = std::mem::replace(self, placeholder) {
|
||||
*self = OAuthState::AuthorizedHttpClient(AuthorizedHttpClient::new(
|
||||
Arc::new(manager),
|
||||
|
|
@ -3610,7 +3629,7 @@ impl OAuthState {
|
|||
required_scope: &str,
|
||||
redirect_uri: &str,
|
||||
) -> Result<String, AuthError> {
|
||||
let placeholder = self.placeholder().await?;
|
||||
let placeholder = self.placeholder_state().await?;
|
||||
let old = std::mem::replace(self, placeholder);
|
||||
let OAuthState::Authorized(manager) = old else {
|
||||
*self = old;
|
||||
|
|
@ -3742,7 +3761,7 @@ impl OAuthState {
|
|||
&mut self,
|
||||
config: ClientCredentialsConfig,
|
||||
) -> Result<(), AuthError> {
|
||||
let placeholder = self.placeholder().await?;
|
||||
let placeholder = self.placeholder_state().await?;
|
||||
let OAuthState::Unauthorized(mut manager) = std::mem::replace(self, placeholder) else {
|
||||
return Err(AuthError::InternalError(
|
||||
"Client credentials flow requires Unauthorized state".to_string(),
|
||||
|
|
@ -4137,7 +4156,6 @@ mod tests {
|
|||
.body(Vec::new())
|
||||
.unwrap();
|
||||
let client = RecordingOAuthHttpClient::with_responses(vec![
|
||||
empty_response(404),
|
||||
challenge,
|
||||
http_response(
|
||||
200,
|
||||
|
|
@ -4178,7 +4196,6 @@ mod tests {
|
|||
(
|
||||
"https://auth.example.com/tenant1/token",
|
||||
vec![
|
||||
"https://mcp.example.com/mcp",
|
||||
"https://mcp.example.com/mcp",
|
||||
"https://mcp.example.com/custom/metadata/location.json",
|
||||
"https://auth.example.com/.well-known/oauth-authorization-server/tenant1",
|
||||
|
|
@ -4187,14 +4204,12 @@ mod tests {
|
|||
],
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
assert!(
|
||||
client
|
||||
.requests()
|
||||
.iter()
|
||||
.take(2)
|
||||
.map(|request| request.method.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["GET", "POST"]
|
||||
.all(|request| request.method == "GET"),
|
||||
"discovery must not send non-GET requests"
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -4205,7 +4220,6 @@ mod tests {
|
|||
empty_response(404),
|
||||
empty_response(404),
|
||||
empty_response(404),
|
||||
empty_response(404),
|
||||
]);
|
||||
let manager = AuthorizationManager::new_with_oauth_http_client(
|
||||
"https://legacy.example.com/",
|
||||
|
|
@ -4234,7 +4248,6 @@ mod tests {
|
|||
"https://legacy.example.com/token",
|
||||
Some("https://legacy.example.com/register"),
|
||||
vec![
|
||||
"https://legacy.example.com/",
|
||||
"https://legacy.example.com/",
|
||||
"https://legacy.example.com/.well-known/oauth-protected-resource",
|
||||
"https://legacy.example.com/.well-known/oauth-authorization-server",
|
||||
|
|
@ -4293,6 +4306,85 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_metadata_from_challenge_uses_challenge_pointer_and_scope() {
|
||||
let client = RecordingOAuthHttpClient::with_responses(vec![
|
||||
http_response(
|
||||
200,
|
||||
serde_json::json!({
|
||||
"resource": "https://mcp.example.com/mcp",
|
||||
"authorization_servers": ["https://auth.example.com"]
|
||||
}),
|
||||
),
|
||||
http_response(
|
||||
200,
|
||||
serde_json::json!({
|
||||
"issuer": "https://auth.example.com",
|
||||
"authorization_endpoint": "https://auth.example.com/authorize",
|
||||
"token_endpoint": "https://auth.example.com/token"
|
||||
}),
|
||||
),
|
||||
]);
|
||||
let manager = AuthorizationManager::new_with_oauth_http_client(
|
||||
"https://mcp.example.com/mcp",
|
||||
Arc::new(client.clone()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let resolution = manager
|
||||
.resolve_metadata_from_challenge(Some(
|
||||
r#"Bearer resource_metadata="https://mcp.example.com/custom/prm.json", scope="mcp:read mcp:write""#,
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
resolution.source,
|
||||
resolution.metadata.token_endpoint.as_str(),
|
||||
manager.select_scopes(None, &[]),
|
||||
client.requests().first().map(|request| request.uri.clone()),
|
||||
),
|
||||
(
|
||||
AuthorizationMetadataSource::ProtectedResourceMetadata,
|
||||
"https://auth.example.com/token",
|
||||
vec!["mcp:read".to_string(), "mcp:write".to_string()],
|
||||
// discovery starts at the challenge's pointer: no probing
|
||||
Some("https://mcp.example.com/custom/prm.json".to_string()),
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_metadata_from_challenge_falls_back_without_metadata_pointer() {
|
||||
let client = RecordingOAuthHttpClient::with_responses(vec![
|
||||
empty_response(404),
|
||||
empty_response(404),
|
||||
empty_response(404),
|
||||
empty_response(404),
|
||||
]);
|
||||
let manager = AuthorizationManager::new_with_oauth_http_client(
|
||||
"https://legacy.example.com/",
|
||||
Arc::new(client),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let resolution = manager
|
||||
.resolve_metadata_from_challenge(Some(r#"Bearer scope="mcp:basic""#))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
(resolution.source, manager.select_scopes(None, &[]),),
|
||||
(
|
||||
AuthorizationMetadataSource::LegacyEndpointFallback,
|
||||
vec!["mcp:basic".to_string()],
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_metadata_reports_authorization_server_metadata() {
|
||||
let client = RecordingOAuthHttpClient::with_responses(vec![
|
||||
|
|
|
|||
|
|
@ -1,11 +1,70 @@
|
|||
use std::collections::HashMap;
|
||||
|
||||
use http::{HeaderName, HeaderValue};
|
||||
use tracing::debug;
|
||||
|
||||
use crate::transport::{
|
||||
auth::AuthClient,
|
||||
auth::{AuthClient, AuthError},
|
||||
streamable_http_client::{StreamableHttpClient, StreamableHttpError},
|
||||
};
|
||||
|
||||
impl<C> AuthClient<C>
|
||||
where
|
||||
C: StreamableHttpClient + Send + Sync,
|
||||
{
|
||||
/// Run `call` with a token when one is available, reacting to the
|
||||
/// server's auth verdict instead of requiring credentials up front:
|
||||
///
|
||||
/// - no usable credentials → the request goes out unauthenticated, and a
|
||||
/// 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.
|
||||
async fn call_reacting_to_challenges<T, F, Fut>(
|
||||
&self,
|
||||
auth_token: Option<String>,
|
||||
call: F,
|
||||
) -> Result<T, StreamableHttpError<C::Error>>
|
||||
where
|
||||
F: Fn(Option<String>) -> Fut,
|
||||
Fut: Future<Output = Result<T, StreamableHttpError<C::Error>>>,
|
||||
{
|
||||
// Missing credentials are not an error in the reactive model: the
|
||||
// request goes out unauthenticated and the server's 401 challenge
|
||||
// drives authorization.
|
||||
let auth_token = match auth_token {
|
||||
None => match self.get_access_token().await {
|
||||
Ok(token) => Some(token),
|
||||
Err(AuthError::AuthorizationRequired) => None,
|
||||
Err(error) => return Err(error.into()),
|
||||
},
|
||||
token => token,
|
||||
};
|
||||
match call(auth_token.clone()).await {
|
||||
Err(StreamableHttpError::AuthRequired(challenge)) => {
|
||||
// One silent recovery attempt: refresh the rejected token and
|
||||
// retry only when the refresh actually produced a new one.
|
||||
let Some(sent_token) = auth_token else {
|
||||
return Err(StreamableHttpError::AuthRequired(challenge));
|
||||
};
|
||||
let refreshed = {
|
||||
let manager = self.auth_manager.lock().await;
|
||||
manager.try_refresh_or_reauth().await
|
||||
};
|
||||
match refreshed {
|
||||
Ok(fresh_token) if fresh_token != sent_token => call(Some(fresh_token)).await,
|
||||
Ok(_) => Err(StreamableHttpError::AuthRequired(challenge)),
|
||||
Err(error) => {
|
||||
debug!("token refresh after server rejection failed: {error}");
|
||||
Err(StreamableHttpError::AuthRequired(challenge))
|
||||
}
|
||||
}
|
||||
}
|
||||
result => result,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> StreamableHttpClient for AuthClient<C>
|
||||
where
|
||||
C: StreamableHttpClient + Send + Sync,
|
||||
|
|
@ -16,16 +75,21 @@ where
|
|||
&self,
|
||||
uri: std::sync::Arc<str>,
|
||||
session_id: std::sync::Arc<str>,
|
||||
mut auth_token: Option<String>,
|
||||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<(), crate::transport::streamable_http_client::StreamableHttpError<Self::Error>>
|
||||
{
|
||||
if auth_token.is_none() {
|
||||
auth_token = Some(self.get_access_token().await?);
|
||||
}
|
||||
self.http_client
|
||||
.delete_session(uri, session_id, auth_token, custom_headers)
|
||||
.await
|
||||
self.call_reacting_to_challenges(auth_token, |token| {
|
||||
let uri = uri.clone();
|
||||
let session_id = session_id.clone();
|
||||
let custom_headers = custom_headers.clone();
|
||||
async move {
|
||||
self.http_client
|
||||
.delete_session(uri, session_id, token, custom_headers)
|
||||
.await
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_stream(
|
||||
|
|
@ -33,18 +97,24 @@ where
|
|||
uri: std::sync::Arc<str>,
|
||||
session_id: Option<std::sync::Arc<str>>,
|
||||
last_event_id: Option<String>,
|
||||
mut auth_token: Option<String>,
|
||||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<
|
||||
futures::stream::BoxStream<'static, Result<sse_stream::Sse, sse_stream::Error>>,
|
||||
crate::transport::streamable_http_client::StreamableHttpError<Self::Error>,
|
||||
> {
|
||||
if auth_token.is_none() {
|
||||
auth_token = Some(self.get_access_token().await?);
|
||||
}
|
||||
self.http_client
|
||||
.get_stream(uri, session_id, last_event_id, auth_token, custom_headers)
|
||||
.await
|
||||
self.call_reacting_to_challenges(auth_token, |token| {
|
||||
let uri = uri.clone();
|
||||
let session_id = session_id.clone();
|
||||
let last_event_id = last_event_id.clone();
|
||||
let custom_headers = custom_headers.clone();
|
||||
async move {
|
||||
self.http_client
|
||||
.get_stream(uri, session_id, last_event_id, token, custom_headers)
|
||||
.await
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_stream_with_max_sse_event_size(
|
||||
|
|
@ -52,26 +122,32 @@ where
|
|||
uri: std::sync::Arc<str>,
|
||||
session_id: Option<std::sync::Arc<str>>,
|
||||
last_event_id: Option<String>,
|
||||
mut auth_token: Option<String>,
|
||||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
max_sse_event_size: usize,
|
||||
) -> Result<
|
||||
futures::stream::BoxStream<'static, Result<sse_stream::Sse, sse_stream::Error>>,
|
||||
crate::transport::streamable_http_client::StreamableHttpError<Self::Error>,
|
||||
> {
|
||||
if auth_token.is_none() {
|
||||
auth_token = Some(self.get_access_token().await?);
|
||||
}
|
||||
self.http_client
|
||||
.get_stream_with_max_sse_event_size(
|
||||
uri,
|
||||
session_id,
|
||||
last_event_id,
|
||||
auth_token,
|
||||
custom_headers,
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
self.call_reacting_to_challenges(auth_token, |token| {
|
||||
let uri = uri.clone();
|
||||
let session_id = session_id.clone();
|
||||
let last_event_id = last_event_id.clone();
|
||||
let custom_headers = custom_headers.clone();
|
||||
async move {
|
||||
self.http_client
|
||||
.get_stream_with_max_sse_event_size(
|
||||
uri,
|
||||
session_id,
|
||||
last_event_id,
|
||||
token,
|
||||
custom_headers,
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn post_message(
|
||||
|
|
@ -79,18 +155,24 @@ where
|
|||
uri: std::sync::Arc<str>,
|
||||
message: crate::model::ClientJsonRpcMessage,
|
||||
session_id: Option<std::sync::Arc<str>>,
|
||||
mut auth_token: Option<String>,
|
||||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<
|
||||
crate::transport::streamable_http_client::StreamableHttpPostResponse,
|
||||
StreamableHttpError<Self::Error>,
|
||||
> {
|
||||
if auth_token.is_none() {
|
||||
auth_token = Some(self.get_access_token().await?);
|
||||
}
|
||||
self.http_client
|
||||
.post_message(uri, message, session_id, auth_token, custom_headers)
|
||||
.await
|
||||
self.call_reacting_to_challenges(auth_token, |token| {
|
||||
let uri = uri.clone();
|
||||
let message = message.clone();
|
||||
let session_id = session_id.clone();
|
||||
let custom_headers = custom_headers.clone();
|
||||
async move {
|
||||
self.http_client
|
||||
.post_message(uri, message, session_id, token, custom_headers)
|
||||
.await
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn post_message_with_max_sse_event_size(
|
||||
|
|
@ -98,25 +180,31 @@ where
|
|||
uri: std::sync::Arc<str>,
|
||||
message: crate::model::ClientJsonRpcMessage,
|
||||
session_id: Option<std::sync::Arc<str>>,
|
||||
mut auth_token: Option<String>,
|
||||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
max_sse_event_size: usize,
|
||||
) -> Result<
|
||||
crate::transport::streamable_http_client::StreamableHttpPostResponse,
|
||||
StreamableHttpError<Self::Error>,
|
||||
> {
|
||||
if auth_token.is_none() {
|
||||
auth_token = Some(self.get_access_token().await?);
|
||||
}
|
||||
self.http_client
|
||||
.post_message_with_max_sse_event_size(
|
||||
uri,
|
||||
message,
|
||||
session_id,
|
||||
auth_token,
|
||||
custom_headers,
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
self.call_reacting_to_challenges(auth_token, |token| {
|
||||
let uri = uri.clone();
|
||||
let message = message.clone();
|
||||
let session_id = session_id.clone();
|
||||
let custom_headers = custom_headers.clone();
|
||||
async move {
|
||||
self.http_client
|
||||
.post_message_with_max_sse_event_size(
|
||||
uri,
|
||||
message,
|
||||
session_id,
|
||||
token,
|
||||
custom_headers,
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -124,7 +124,8 @@ fn negotiate_version_headers(
|
|||
(version, headers)
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
#[derive(Debug, Error)]
|
||||
#[error("authorization required: {www_authenticate_header}")]
|
||||
#[non_exhaustive]
|
||||
pub struct AuthRequiredError {
|
||||
pub www_authenticate_header: String,
|
||||
|
|
@ -139,7 +140,8 @@ impl AuthRequiredError {
|
|||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
#[derive(Debug, Error)]
|
||||
#[error("insufficient scope: {www_authenticate_header}")]
|
||||
#[non_exhaustive]
|
||||
pub struct InsufficientScopeError {
|
||||
pub www_authenticate_header: String,
|
||||
|
|
@ -197,15 +199,31 @@ pub enum StreamableHttpError<E: std::error::Error + Send + Sync + 'static> {
|
|||
#[error("Auth error: {0}")]
|
||||
Auth(#[from] crate::transport::auth::AuthError),
|
||||
#[error("Auth required")]
|
||||
AuthRequired(AuthRequiredError),
|
||||
AuthRequired(#[source] AuthRequiredError),
|
||||
#[error("Insufficient scope")]
|
||||
InsufficientScope(InsufficientScopeError),
|
||||
InsufficientScope(#[source] InsufficientScopeError),
|
||||
#[error("Header name '{0}' is reserved and conflicts with default headers")]
|
||||
ReservedHeaderConflict(String),
|
||||
#[error("Session expired (HTTP 404)")]
|
||||
SessionExpired,
|
||||
}
|
||||
|
||||
impl<E: std::error::Error + Send + Sync + 'static> StreamableHttpError<E> {
|
||||
/// The `WWW-Authenticate` challenge carried by this error, when the
|
||||
/// server answered 401 ([`AuthRequired`](Self::AuthRequired)) or 403
|
||||
/// ([`InsufficientScope`](Self::InsufficientScope)). Feed it to
|
||||
/// [`AuthorizationRequest::with_challenge`](crate::transport::auth::AuthorizationRequest::with_challenge)
|
||||
/// to authorize reactively.
|
||||
#[cfg(feature = "auth")]
|
||||
pub fn auth_challenge(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::AuthRequired(error) => Some(&error.www_authenticate_header),
|
||||
Self::InsufficientScope(error) => Some(&error.www_authenticate_header),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Error)]
|
||||
#[non_exhaustive]
|
||||
pub enum StreamableHttpProtocolError {
|
||||
|
|
|
|||
|
|
@ -8,8 +8,9 @@ use axum::{
|
|||
routing::get,
|
||||
};
|
||||
use rmcp::{
|
||||
ServiceExt,
|
||||
RoleClient, ServiceExt,
|
||||
model::ClientInfo,
|
||||
service::RunningService,
|
||||
transport::{
|
||||
StreamableHttpClientTransport,
|
||||
auth::{AuthClient, AuthorizationRequest, OAuthState},
|
||||
|
|
@ -55,6 +56,109 @@ async fn callback_handler(
|
|||
Html(CALLBACK_HTML.to_string())
|
||||
}
|
||||
|
||||
enum ConnectOutcome {
|
||||
/// The server accepted the unauthenticated connection.
|
||||
Connected(RunningService<RoleClient, ClientInfo>),
|
||||
/// The server answered 401; authorize with this `WWW-Authenticate`
|
||||
/// challenge and reconnect.
|
||||
AuthRequired(String),
|
||||
}
|
||||
|
||||
/// Attempt the real connection unauthenticated — the reactive discovery
|
||||
/// trigger (matching the TypeScript and Python SDKs). The server's 401
|
||||
/// challenge, not a probe, tells us whether and how to authorize.
|
||||
async fn try_connect(http_client: reqwest::Client, server_url: &str) -> Result<ConnectOutcome> {
|
||||
let transport = StreamableHttpClientTransport::with_client(
|
||||
http_client,
|
||||
StreamableHttpClientTransportConfig::with_uri(server_url),
|
||||
);
|
||||
match ClientInfo::default().serve(transport).await {
|
||||
Ok(client) => Ok(ConnectOutcome::Connected(client)),
|
||||
Err(error) => match error.auth_challenge() {
|
||||
Some(challenge) => Ok(ConnectOutcome::AuthRequired(challenge.to_string())),
|
||||
None => Err(error.into()),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Run the browser OAuth flow seeded by the server's challenge, then
|
||||
/// reconnect with the authorized transport.
|
||||
async fn authorize_and_connect(
|
||||
challenge: String,
|
||||
oauth_http_client: reqwest::Client,
|
||||
server_url: &str,
|
||||
client_metadata_url: &str,
|
||||
code_receiver: oneshot::Receiver<CallbackParams>,
|
||||
output: &mut BufWriter<tokio::io::Stdout>,
|
||||
) -> Result<RunningService<RoleClient, ClientInfo>> {
|
||||
tracing::info!("Server requires authorization: {challenge}");
|
||||
|
||||
// initialize oauth state machine
|
||||
let mut oauth_state = OAuthState::new(server_url, Some(oauth_http_client))
|
||||
.await
|
||||
.context("Failed to initialize oauth state machine")?;
|
||||
// Seed discovery from the server's challenge, and use CIMD (SEP-991)
|
||||
// with client metadata URL. Passing no scopes lets the SDK auto-select
|
||||
// from the challenge's scope hint, Protected Resource Metadata, or AS
|
||||
// metadata.
|
||||
oauth_state
|
||||
.start_authorization(
|
||||
AuthorizationRequest::new(MCP_REDIRECT_URI)
|
||||
.with_client_name("Test MCP Client")
|
||||
.with_client_metadata_url(client_metadata_url)
|
||||
.with_challenge(challenge),
|
||||
)
|
||||
.await
|
||||
.context("Failed to start authorization")?;
|
||||
|
||||
// Output authorization URL to user
|
||||
output
|
||||
.write_all(b"Please open the following URL in your browser to authorize:\n\n")
|
||||
.await?;
|
||||
output
|
||||
.write_all(oauth_state.get_authorization_url().await?.as_bytes())
|
||||
.await?;
|
||||
output
|
||||
.write_all(b"\n\nWaiting for browser callback, please do not close this window...\n")
|
||||
.await?;
|
||||
output.flush().await?;
|
||||
|
||||
// Wait for authorization code
|
||||
tracing::info!("Waiting for authorization code...");
|
||||
let CallbackParams {
|
||||
code: auth_code,
|
||||
state: csrf_token,
|
||||
iss,
|
||||
} = code_receiver
|
||||
.await
|
||||
.context("Failed to get authorization code")?;
|
||||
tracing::info!("Received authorization code: {}", auth_code);
|
||||
// Exchange code for access token
|
||||
tracing::info!("Exchanging authorization code for access token...");
|
||||
oauth_state
|
||||
.handle_callback_with_issuer(&auth_code, &csrf_token, iss.as_deref())
|
||||
.await
|
||||
.context("Failed to handle callback")?;
|
||||
tracing::info!("Successfully obtained access token");
|
||||
|
||||
output
|
||||
.write_all(b"\nAuthorization successful! Access token obtained.\n\n")
|
||||
.await?;
|
||||
output.flush().await?;
|
||||
|
||||
// Reconnect with the authorized transport
|
||||
tracing::info!("Establishing authorized connection to MCP server...");
|
||||
let am = oauth_state
|
||||
.into_authorization_manager()
|
||||
.ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?;
|
||||
let auth_client = AuthClient::new(reqwest::Client::default(), am);
|
||||
let transport = StreamableHttpClientTransport::with_client(
|
||||
auth_client,
|
||||
StreamableHttpClientTransportConfig::with_uri(server_url),
|
||||
);
|
||||
Ok(ClientInfo::default().serve(transport).await?)
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
// Initialize logging
|
||||
|
|
@ -123,73 +227,32 @@ async fn main() -> Result<()> {
|
|||
.build()
|
||||
.context("Failed to build OAuth HTTP client")?;
|
||||
|
||||
// initialize oauth state machine
|
||||
let mut oauth_state = OAuthState::new(&server_url, Some(oauth_http_client))
|
||||
.await
|
||||
.context("Failed to initialize oauth state machine")?;
|
||||
// use CIMD (SEP-991) with client metadata URL.
|
||||
// passing no scopes lets the SDK auto-select from the server's
|
||||
// WWW-Authenticate header, Protected Resource Metadata, or AS metadata.
|
||||
oauth_state
|
||||
.start_authorization(
|
||||
AuthorizationRequest::new(MCP_REDIRECT_URI)
|
||||
.with_client_name("Test MCP Client")
|
||||
.with_client_metadata_url(&client_metadata_url),
|
||||
)
|
||||
.await
|
||||
.context("Failed to start authorization")?;
|
||||
|
||||
// Output authorization URL to user
|
||||
let mut output = BufWriter::new(tokio::io::stdout());
|
||||
output.write_all(b"\n=== MCP OAuth Client ===\n\n").await?;
|
||||
output
|
||||
.write_all(b"Please open the following URL in your browser to authorize:\n\n")
|
||||
.await?;
|
||||
output
|
||||
.write_all(oauth_state.get_authorization_url().await?.as_bytes())
|
||||
.await?;
|
||||
output
|
||||
.write_all(b"\n\nWaiting for browser callback, please do not close this window...\n")
|
||||
.await?;
|
||||
output.flush().await?;
|
||||
|
||||
// Wait for authorization code
|
||||
tracing::info!("Waiting for authorization code...");
|
||||
let CallbackParams {
|
||||
code: auth_code,
|
||||
state: csrf_token,
|
||||
iss,
|
||||
} = code_receiver
|
||||
.await
|
||||
.context("Failed to get authorization code")?;
|
||||
tracing::info!("Received authorization code: {}", auth_code);
|
||||
// Exchange code for access token
|
||||
tracing::info!("Exchanging authorization code for access token...");
|
||||
oauth_state
|
||||
.handle_callback_with_issuer(&auth_code, &csrf_token, iss.as_deref())
|
||||
.await
|
||||
.context("Failed to handle callback")?;
|
||||
tracing::info!("Successfully obtained access token");
|
||||
|
||||
output
|
||||
.write_all(b"\nAuthorization successful! Access token obtained.\n\n")
|
||||
.await?;
|
||||
output.flush().await?;
|
||||
|
||||
// Create authorized transport, this transport is authorized by the oauth state machine
|
||||
tracing::info!("Establishing authorized connection to MCP server...");
|
||||
let am = oauth_state
|
||||
.into_authorization_manager()
|
||||
.ok_or_else(|| anyhow::anyhow!("Failed to get authorization manager"))?;
|
||||
let client = AuthClient::new(reqwest::Client::default(), am);
|
||||
let transport = StreamableHttpClientTransport::with_client(
|
||||
client,
|
||||
StreamableHttpClientTransportConfig::with_uri(server_url.as_str()),
|
||||
);
|
||||
|
||||
// Create client and connect to MCP server
|
||||
let client_service = ClientInfo::default();
|
||||
let client = client_service.serve(transport).await?;
|
||||
// Reactive discovery: attempt the real connection first. The server's
|
||||
// 401 challenge — not a probe — tells us whether and how to authorize.
|
||||
// The transport gets a default client: `oauth_http_client`'s request
|
||||
// timeout would cut long-lived SSE streams short.
|
||||
tracing::info!("Attempting connection to MCP server...");
|
||||
let client = match try_connect(reqwest::Client::default(), &server_url).await? {
|
||||
ConnectOutcome::Connected(client) => {
|
||||
tracing::info!("Server accepted the connection without authorization");
|
||||
client
|
||||
}
|
||||
ConnectOutcome::AuthRequired(challenge) => {
|
||||
authorize_and_connect(
|
||||
challenge,
|
||||
oauth_http_client,
|
||||
&server_url,
|
||||
&client_metadata_url,
|
||||
code_receiver,
|
||||
&mut output,
|
||||
)
|
||||
.await?
|
||||
}
|
||||
};
|
||||
tracing::info!("Successfully connected to MCP server");
|
||||
|
||||
// Test API requests
|
||||
|
|
|
|||
Loading…
Reference in a new issue