diff --git a/crates/aether-auth/Cargo.toml b/crates/aether-auth/Cargo.toml index 88d2796b2..85bdadace 100644 --- a/crates/aether-auth/Cargo.toml +++ b/crates/aether-auth/Cargo.toml @@ -21,7 +21,7 @@ keyring = [ "dep:windows-native-keyring-store", "dep:dbus-secret-service-keyring-store", ] -mcp = ["dep:rmcp", "dep:url"] +mcp = ["dep:rmcp"] [dependencies] age = { workspace = true } @@ -41,7 +41,7 @@ keyring-core = { workspace = true, optional = true } rmcp = { workspace = true, optional = true, features = ["auth", "client"] } oauth2 = { workspace = true } reqwest = { workspace = true } -url = { workspace = true, optional = true } +url = { workspace = true } [target.'cfg(target_os = "macos")'.dependencies] apple-native-keyring-store = { workspace = true, optional = true } diff --git a/crates/aether-auth/README.md b/crates/aether-auth/README.md index 674df453b..1456ab942 100644 --- a/crates/aether-auth/README.md +++ b/crates/aether-auth/README.md @@ -52,7 +52,7 @@ let loaded = store.load_credential("anthropic").await?; # } ``` -For MCP servers that require OAuth, the `mcp` feature provides `perform_oauth_flow`, which orchestrates the full authorization-code flow (browser launch, callback capture, token exchange, credential storage) and `create_auth_manager_from_store`, which builds an `rmcp::transport::auth::AuthorizationManager` from stored credentials with automatic token refresh. +For MCP servers that require OAuth, the `mcp` feature integrates Aether's browser and callback UI with `rmcp`'s authorization state machine. ## Feature Flags diff --git a/crates/aether-auth/src/browser.rs b/crates/aether-auth/src/browser.rs index 0c77e4c34..59f9e6280 100644 --- a/crates/aether-auth/src/browser.rs +++ b/crates/aether-auth/src/browser.rs @@ -1,5 +1,5 @@ use crate::error::OAuthError; -use crate::handler::{OAuthCallback, OAuthHandler}; +use crate::handler::OAuthHandler; use futures::future::BoxFuture; use std::process::Command; use std::time::Duration; @@ -24,9 +24,6 @@ impl BrowserOAuthHandler { } /// Create a handler bound to a specific port with a custom redirect URI. - /// - /// Use this when the OAuth provider has a fixed redirect URI registered - /// (e.g. `http://localhost:1455/auth/callback` for Codex). pub fn with_redirect_uri(redirect_uri: impl Into, port: u16) -> Result { let std_listener = std::net::TcpListener::bind(format!("127.0.0.1:{port}"))?; std_listener.set_nonblocking(true)?; @@ -40,24 +37,19 @@ impl OAuthHandler for BrowserOAuthHandler { &self.redirect_uri } - fn authorize(&self, auth_url: &str) -> BoxFuture<'_, Result> { + fn authorize(&self, auth_url: &str) -> BoxFuture<'_, Result> { let auth_url = auth_url.to_string(); Box::pin(async move { - if let Err(e) = open_browser(&auth_url) { - tracing::warn!("Failed to open browser: {e}. Open manually: {auth_url}"); + if let Err(error) = open_browser(&auth_url) { + tracing::warn!("Failed to open browser: {error}"); } - accept_oauth_callback(&self.listener).await }) } } -/// Accept an OAuth callback on an already-bound listener. -/// -/// Loops until it receives a valid OAuth callback, skipping stale connections -/// left over from previous flows (e.g. browser favicon requests, preconnect -/// sockets, or closed connections). -pub async fn accept_oauth_callback(listener: &TcpListener) -> Result { +/// Accept an OAuth callback and return its absolute URL. +pub async fn accept_oauth_callback(listener: &TcpListener) -> Result { loop { let (mut socket, _) = listener.accept().await?; let request_line = { @@ -69,31 +61,24 @@ pub async fn accept_oauth_callback(listener: &TcpListener) -> Result { - let _ = socket.write_all(create_success_response().as_bytes()).await; - return Ok(callback); - } - Err(e) => { - if request_line.contains('?') { - return Err(e); - } + match callback_url(&request_line) { + Ok(callback_url) => { + let _ = socket.write_all(success_response().as_bytes()).await; + return Ok(callback_url); } + Err(error) if request_line.contains('?') => return Err(error), + Err(_) => {} } } } -/// Start a local callback server to capture the OAuth authorization code and state -/// -/// Listens on the specified port and waits for the OAuth redirect. -/// Returns the authorization code and state (CSRF token) from the callback URL. -pub async fn wait_for_callback(port: u16) -> Result { - let addr = format!("127.0.0.1:{port}"); - let listener = TcpListener::bind(&addr).await?; +/// Start a local callback server and return the OAuth callback URL. +pub async fn wait_for_callback(port: u16) -> Result { + let listener = TcpListener::bind(format!("127.0.0.1:{port}")).await?; accept_oauth_callback(&listener).await } -/// Open a URL in the default browser +/// Open a URL in the default browser. pub fn open_browser(url: &str) -> Result<(), OAuthError> { #[cfg(target_os = "macos")] { @@ -113,83 +98,24 @@ pub fn open_browser(url: &str) -> Result<(), OAuthError> { Ok(()) } -/// Parse the authorization code and state from the HTTP request line -fn parse_callback_from_request(request_line: &str) -> Result { - // Request format: GET /oauth2callback?code=XXX&state=YYY HTTP/1.1 - let parts: Vec<&str> = request_line.split_whitespace().collect(); - if parts.len() < 2 { - return Err(OAuthError::InvalidCallback("Invalid HTTP request format".to_string())); - } - - let path = parts[1]; - let query_start = - path.find('?').ok_or_else(|| OAuthError::InvalidCallback("No query parameters in callback".to_string()))?; - - let query = &path[query_start + 1..]; - - // Check for error in callback - for param in query.split('&') { - if let Some((key, value)) = param.split_once('=') - && key == "error" - { - let error_desc = query - .split('&') - .find_map(|p| { - p.split_once('=').filter(|(k, _)| *k == "error_description").map(|(_, v)| urlencoding_decode(v)) - }) - .unwrap_or_else(|| value.to_string()); - return Err(OAuthError::InvalidCallback(format!("OAuth error: {error_desc}"))); - } +fn callback_url(request_line: &str) -> Result { + let path = request_line + .split_whitespace() + .nth(1) + .ok_or_else(|| OAuthError::InvalidCallback("Invalid HTTP request format".to_string()))?; + let url = url::Url::parse(&format!("http://localhost{path}")) + .map_err(|error| OAuthError::InvalidCallback(format!("Invalid callback URL: {error}")))?; + if url.query().is_none() { + return Err(OAuthError::InvalidCallback("No query parameters in callback".to_string())); } - - // Extract code and state - let mut code = None; - let mut state = None; - - for param in query.split('&') { - if let Some((key, value)) = param.split_once('=') { - match key { - "code" => code = Some(urlencoding_decode(value)), - "state" => state = Some(urlencoding_decode(value)), - _ => {} - } - } - } - - let code = code.ok_or_else(|| OAuthError::InvalidCallback("No authorization code in callback".into()))?; - let state = state.ok_or_else(|| OAuthError::InvalidCallback("No state parameter in callback".into()))?; - - Ok(OAuthCallback { code, state }) -} - -/// Simple URL decoding (handles %XX escapes) -fn urlencoding_decode(s: &str) -> String { - let mut result = String::with_capacity(s.len()); - let mut chars = s.chars().peekable(); - - while let Some(c) = chars.next() { - if c == '%' { - let hex: String = chars.by_ref().take(2).collect(); - if let Ok(byte) = u8::from_str_radix(&hex, 16) { - result.push(byte as char); - } else { - result.push('%'); - result.push_str(&hex); - } - } else if c == '+' { - result.push(' '); - } else { - result.push(c); - } + if url.query_pairs().any(|(key, _)| key == "error") { + return Err(OAuthError::InvalidCallback("OAuth authorization failed".to_string())); } - - result + Ok(url.to_string()) } -/// Create an HTML success response -fn create_success_response() -> String { +fn success_response() -> String { let body = include_str!("oauth_success.html"); - format!( "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", body.len(), @@ -202,97 +128,45 @@ mod tests { use super::*; #[test] - fn parse_callback_from_valid_request() { - let request = "GET /oauth2callback?code=4%2F0AYWS-abc123&state=verifier HTTP/1.1\r\n"; - let callback = parse_callback_from_request(request).unwrap(); - assert_eq!(callback.code, "4/0AYWS-abc123"); - assert_eq!(callback.state, "verifier"); - } - - #[test] - fn parse_callback_handles_plus_encoding() { - let request = "GET /oauth2callback?code=hello+world&state=test+state HTTP/1.1\r\n"; - let callback = parse_callback_from_request(request).unwrap(); - assert_eq!(callback.code, "hello world"); - assert_eq!(callback.state, "test state"); - } - - #[test] - fn parse_callback_returns_error_for_oauth_error() { - let request = "GET /oauth2callback?error=access_denied&error_description=User+denied HTTP/1.1\r\n"; - let result = parse_callback_from_request(request); - assert!(result.is_err()); - let err = result.unwrap_err().to_string(); - assert!(err.contains("User denied")); - } - - #[test] - fn parse_callback_returns_error_for_missing_code() { - let request = "GET /oauth2callback?state=verifier HTTP/1.1\r\n"; - let result = parse_callback_from_request(request); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("No authorization code")); - } - - #[test] - fn parse_callback_returns_error_for_missing_state() { - let request = "GET /oauth2callback?code=abc123 HTTP/1.1\r\n"; - let result = parse_callback_from_request(request); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("No state parameter")); - } - - #[tokio::test] - async fn with_redirect_uri_binds_to_specified_port() { - let handler = BrowserOAuthHandler::with_redirect_uri("http://localhost:9999/callback", 0).unwrap(); - assert_eq!(handler.redirect_uri(), "http://localhost:9999/callback"); - } - - #[test] - fn urlencoding_decode_handles_percent() { - assert_eq!(urlencoding_decode("hello%20world"), "hello world"); - assert_eq!(urlencoding_decode("a%2Fb"), "a/b"); + fn callback_request_preserves_encoded_parameters() { + let callback = callback_url( + "GET /oauth2callback?code=caf%C3%A9&state=test+state&iss=https%3A%2F%2Fauth.example.com HTTP/1.1\r\n", + ) + .unwrap(); + let url = url::Url::parse(&callback).unwrap(); + let params = url.query_pairs().collect::>(); + assert_eq!(params.get("code").map(AsRef::as_ref), Some("café")); + assert_eq!(params.get("state").map(AsRef::as_ref), Some("test state")); + assert_eq!(params.get("iss").map(AsRef::as_ref), Some("https://auth.example.com")); } #[test] - fn urlencoding_decode_handles_plus() { - assert_eq!(urlencoding_decode("hello+world"), "hello world"); + fn callback_error_is_sanitized() { + let error = + callback_url("GET /oauth2callback?error=access_denied&error_description=attacker+controlled HTTP/1.1\r\n") + .unwrap_err() + .to_string(); + assert!(error.contains("OAuth authorization failed")); + assert!(!error.contains("attacker")); } #[tokio::test] - async fn accept_oauth_callback_skips_stale_favicon_request() { - use tokio::io::AsyncWriteExt; - + async fn callback_listener_skips_stale_requests() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); let handle = tokio::spawn(async move { accept_oauth_callback(&listener).await }); - let mut stale = tokio::net::TcpStream::connect(format!("127.0.0.1:{port}")).await.unwrap(); + let mut stale = tokio::net::TcpStream::connect(("127.0.0.1", port)).await.unwrap(); stale.write_all(b"GET /favicon.ico HTTP/1.1\r\n").await.unwrap(); + let mut callback = tokio::net::TcpStream::connect(("127.0.0.1", port)).await.unwrap(); + callback.write_all(b"GET /?code=abc&state=xyz HTTP/1.1\r\n").await.unwrap(); - let mut real = tokio::net::TcpStream::connect(format!("127.0.0.1:{port}")).await.unwrap(); - real.write_all(b"GET /oauth2callback?code=abc&state=xyz HTTP/1.1\r\n").await.unwrap(); - - let callback = handle.await.unwrap().unwrap(); - assert_eq!(callback.code, "abc"); - assert_eq!(callback.state, "xyz"); + assert!(handle.await.unwrap().unwrap().contains("code=abc&state=xyz")); } #[tokio::test] - async fn accept_oauth_callback_skips_closed_connection() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let handle = tokio::spawn(async move { accept_oauth_callback(&listener).await }); - - { - let _stream = tokio::net::TcpStream::connect(format!("127.0.0.1:{port}")).await.unwrap(); - } - - let mut real = tokio::net::TcpStream::connect(format!("127.0.0.1:{port}")).await.unwrap(); - real.write_all(b"GET /oauth2callback?code=def&state=uvw HTTP/1.1\r\n").await.unwrap(); - - let callback = handle.await.unwrap().unwrap(); - assert_eq!(callback.code, "def"); - assert_eq!(callback.state, "uvw"); + async fn custom_redirect_uri_is_retained() { + let handler = BrowserOAuthHandler::with_redirect_uri("http://localhost:9999/callback", 0).unwrap(); + assert_eq!(handler.redirect_uri(), "http://localhost:9999/callback"); } } diff --git a/crates/aether-auth/src/handler.rs b/crates/aether-auth/src/handler.rs index 67eb85b11..a325ad7fe 100644 --- a/crates/aether-auth/src/handler.rs +++ b/crates/aether-auth/src/handler.rs @@ -1,23 +1,11 @@ use crate::error::OAuthError; use futures::future::BoxFuture; -/// OAuth callback data containing both the authorization code and state (CSRF token) -#[derive(Debug, Clone)] -pub struct OAuthCallback { - pub code: String, - pub state: String, -} - -/// Trait that consuming applications implement to handle OAuth UI/UX. -/// -/// Uses `BoxFuture` instead of `async fn` to support `dyn OAuthHandler` -/// (required for `Arc` in `McpManager`). +/// UI boundary for interactive OAuth authorization. pub trait OAuthHandler: Send + Sync { - /// The redirect URI the OAuth provider should send the user back to, - /// e.g. `http://127.0.0.1:/oauth2callback`. + /// The redirect URI the OAuth provider should send the user back to. fn redirect_uri(&self) -> &str; - /// Called when user needs to authorize. App should open browser to `auth_url` - /// and return the authorization code and state (CSRF token) from the callback. - fn authorize(&self, auth_url: &str) -> BoxFuture<'_, Result>; + /// Present `auth_url` and return the absolute callback URL. + fn authorize(&self, auth_url: &str) -> BoxFuture<'_, Result>; } diff --git a/crates/aether-auth/src/lib.rs b/crates/aether-auth/src/lib.rs index 8acd7f1ee..313c33a06 100644 --- a/crates/aether-auth/src/lib.rs +++ b/crates/aether-auth/src/lib.rs @@ -12,7 +12,7 @@ pub use credential::{OAuthCredential, OAuthCredentialStorage, oauth_http_client} pub use encrypted_file::EncryptedFileOAuthCredentialStorage; pub use error::OAuthError; pub use fake::FakeOAuthCredentialStore; -pub use handler::{OAuthCallback, OAuthHandler}; +pub use handler::OAuthHandler; #[cfg(feature = "keyring")] pub mod keyring; @@ -22,4 +22,6 @@ pub use keyring::OsKeyringStore; #[cfg(feature = "mcp")] pub mod mcp; #[cfg(feature = "mcp")] -pub use mcp::{McpCredentialStore, create_auth_manager_from_store, perform_oauth_flow}; +pub use mcp::{ + McpCredentialStore, OAuthClientRegistration, OAuthFlowOptions, create_auth_manager_from_store, perform_oauth_flow, +}; diff --git a/crates/aether-auth/src/mcp/credential_store.rs b/crates/aether-auth/src/mcp/credential_store.rs index 9e335fbc4..b618a540c 100644 --- a/crates/aether-auth/src/mcp/credential_store.rs +++ b/crates/aether-auth/src/mcp/credential_store.rs @@ -7,6 +7,9 @@ use crate::{OAuthCredentialStorage, OAuthError}; /// Per-server adapter that binds an [`OAuthCredentialStorage`] to a single MCP server id /// and implements `rmcp::transport::auth::CredentialStore` by persisting rmcp's /// [`StoredCredentials`] verbatim as a namespaced, opaque JSON secret. +/// +/// rmcp owns the credential schema (tokens, granted scopes, issuer binding, +/// `token_received_at`), so nothing is translated in either direction. #[derive(Clone)] pub struct McpCredentialStore { server_id: String, @@ -26,21 +29,14 @@ impl McpCredentialStore { self } - fn secret_key(&self) -> String { - format!("mcp:{}", self.server_id) - } -} - -#[async_trait] -impl CredentialStore for McpCredentialStore { - async fn load(&self) -> Result, AuthError> { + /// Load and filter persisted credentials. + pub async fn load_stored(&self) -> Result, OAuthError> { self.store .load(&self.secret_key()) - .await - .map_err(|error| internal_err(&error))? + .await? .map(serde_json::from_value::) .transpose() - .map_err(|error| AuthError::InternalError(format!("invalid MCP credential: {error}"))) + .map_err(|error| OAuthError::CredentialStore(format!("invalid MCP credential: {error}"))) .map(|credentials| { credentials.filter(|credential| { self.expected_client_id.as_deref().is_none_or(|expected| credential.client_id == expected) @@ -48,6 +44,17 @@ impl CredentialStore for McpCredentialStore { }) } + fn secret_key(&self) -> String { + format!("mcp:{}", self.server_id) + } +} + +#[async_trait] +impl CredentialStore for McpCredentialStore { + async fn load(&self) -> Result, AuthError> { + self.load_stored().await.map_err(|e| internal_err(&e)) + } + async fn save(&self, credentials: StoredCredentials) -> Result<(), AuthError> { let value = serde_json::to_value(&credentials) .map_err(|e| AuthError::InternalError(format!("failed to serialize credentials: {e}")))?; diff --git a/crates/aether-auth/src/mcp/integration.rs b/crates/aether-auth/src/mcp/integration.rs index d1f11dc2e..b3fb7b349 100644 --- a/crates/aether-auth/src/mcp/integration.rs +++ b/crates/aether-auth/src/mcp/integration.rs @@ -1,122 +1,102 @@ use super::credential_store::McpCredentialStore; use crate::{OAuthCredentialStorage, OAuthError, OAuthHandler}; -use rmcp::transport::auth::{AuthClient, AuthError, AuthorizationManager, OAuthClientConfig}; +use rmcp::transport::auth::{ + AuthClient, AuthError, AuthorizationManager, AuthorizationRequest, AuthorizationSession, OAuthClientConfig, +}; use std::sync::Arc; const OAUTH_CLIENT_NAME: &str = "Aether MCP Client"; -/// Returns `Ok(Some(manager))` if credentials were found and initialized successfully, -/// `Ok(None)` if no stored credentials exist, or `Err` on failure. -/// -/// When `expected_client_id` is given, stored credentials issued to a different client -/// are ignored, forcing a fresh authorization flow. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub enum OAuthClientRegistration { + PreRegistered(String), + ClientMetadata(String), + #[default] + Dynamic, +} + +impl OAuthClientRegistration { + pub fn pre_registered_client_id(&self) -> Option<&str> { + match self { + Self::PreRegistered(client_id) => Some(client_id), + Self::ClientMetadata(_) | Self::Dynamic => None, + } + } +} + +#[derive(Debug, Clone, Default)] +pub struct OAuthFlowOptions { + pub client_registration: OAuthClientRegistration, + pub challenge: Option, +} + pub async fn create_auth_manager_from_store( server_id: &str, base_url: &str, expected_client_id: Option<&str>, + redirect_uri: &str, store: Arc, ) -> Result, OAuthError> { let credential_store = McpCredentialStore::new(store, server_id.to_string()).with_expected_client_id(expected_client_id); + let Some(stored) = credential_store.load_stored().await? else { return Ok(None) }; + let manager = { + let mut manager = AuthorizationManager::new(base_url).await.map_err(rmcp_err("OAuth init failed"))?; + manager.set_credential_store(credential_store); - let mut auth_manager = AuthorizationManager::new(base_url).await.map_err(rmcp_err("OAuth init failed"))?; - auth_manager.set_credential_store(credential_store); + if !manager.initialize_from_store().await.map_err(rmcp_err("credential store load failed"))? { + return Ok(None); + } - if auth_manager.initialize_from_store().await.map_err(rmcp_err("credential store load failed"))? { - Ok(Some(auth_manager)) - } else { - Ok(None) - } + manager + .configure_client(OAuthClientConfig::new(stored.client_id, redirect_uri)) + .map_err(rmcp_err("OAuth client restoration failed"))?; + + manager + }; + + Ok(Some(manager)) } -/// Run a full OAuth authorization flow against an MCP server. -/// -/// Discovers the server's authorization metadata, obtains a client (the pre-registered -/// `client_id` when given, dynamic registration otherwise), opens the browser (via the -/// handler), handles the callback, and returns an `AuthClient` ready for authenticated -/// HTTP transport. When `store` is `Some`, the resulting tokens are persisted via the -/// store; when `None`, tokens live only for the lifetime of the returned `AuthClient`. +/// Run an interactive MCP OAuth flow using rmcp's state machine. pub async fn perform_oauth_flow( server_id: &str, base_url: &str, handler: &dyn OAuthHandler, - client_id: Option<&str>, + options: OAuthFlowOptions, store: Option>, ) -> Result, OAuthError> { let mut manager = AuthorizationManager::new(base_url).await.map_err(rmcp_err("OAuth init failed"))?; - if let Some(store) = store { manager.set_credential_store( - McpCredentialStore::new(store, server_id.to_string()).with_expected_client_id(client_id), + McpCredentialStore::new(store, server_id.to_string()) + .with_expected_client_id(options.client_registration.pre_registered_client_id()), ); } - let metadata = manager.resolve_metadata().await.map_err(rmcp_err("OAuth metadata discovery failed"))?.metadata; - manager.set_metadata(metadata); - - let scopes = manager.select_scopes(None, &[]); - let scope_refs = scopes.iter().map(String::as_str).collect::>(); - let client_config = match client_id { - Some(client_id) => OAuthClientConfig::new(client_id, handler.redirect_uri()), - None => manager - .register_client(OAUTH_CLIENT_NAME, handler.redirect_uri(), &scope_refs) - .await - .map_err(rmcp_err("dynamic client registration failed"))?, - }; - manager.configure_client(client_config).map_err(rmcp_err("OAuth client configuration failed"))?; - - let auth_url = - manager.get_authorization_url(&scope_refs).await.map_err(rmcp_err("get_authorization_url failed"))?; - - // Some authorization servers (e.g. Sentry) bake `resource` into their - // authorization_endpoint metadata. rmcp then adds its own `resource` param, - // producing a duplicate that the server rejects with "invalid_target". - let callback = handler.authorize(&dedupe_query_params(&auth_url)).await?; - - manager - .exchange_code_for_token(&callback.code, &callback.state) + let resolution = manager + .resolve_metadata_from_challenge(options.challenge.as_deref()) .await - .map_err(rmcp_err("token exchange failed"))?; - - Ok(AuthClient::new(reqwest::Client::default(), manager)) + .map_err(rmcp_err("OAuth metadata discovery failed"))?; + manager.set_metadata(resolution.metadata); + + let request = AuthorizationRequest::new(handler.redirect_uri()).with_client_name(OAUTH_CLIENT_NAME); + let session = AuthorizationSession::new( + manager, + match options.client_registration { + OAuthClientRegistration::PreRegistered(client_id) => request.with_preregistered_client(client_id), + OAuthClientRegistration::ClientMetadata(url) => request.with_client_metadata_url(url), + OAuthClientRegistration::Dynamic => request, + }, + ) + .await + .map_err(|(_, error)| rmcp_err("OAuth authorization failed")(error))?; + + let callback_url = handler.authorize(session.get_authorization_url()).await?; + session.handle_callback_url(&callback_url).await.map_err(rmcp_err("token exchange failed"))?; + Ok(AuthClient::new(reqwest::Client::default(), session.auth_manager)) } fn rmcp_err(context: &'static str) -> impl Fn(AuthError) -> OAuthError { - move |e| OAuthError::Rmcp(format!("{context}: {e}")) -} - -fn dedupe_query_params(url_str: &str) -> String { - let Ok(mut url) = url::Url::parse(url_str) else { - return url_str.to_string(); - }; - let pairs: std::collections::HashMap = - url.query_pairs().map(|(k, v)| (k.into_owned(), v.into_owned())).collect(); - url.query_pairs_mut().clear().extend_pairs(&pairs); - url.to_string() -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn dedupes_duplicate_resource_param() { - let input = "https://example.com/oauth/authorize?resource=https%3A%2F%2Fa%2Fb&response_type=code&client_id=x&resource=https%3A%2F%2Fa%2Fb"; - let out = dedupe_query_params(input); - let url = url::Url::parse(&out).unwrap(); - let resources: Vec<_> = url.query_pairs().filter(|(k, _)| k == "resource").collect(); - assert_eq!(resources.len(), 1); - assert_eq!(resources[0].1, "https://a/b"); - assert!(url.query_pairs().any(|(k, _)| k == "response_type")); - assert!(url.query_pairs().any(|(k, _)| k == "client_id")); - } - - #[test] - fn preserves_unique_params() { - let input = "https://example.com/?resource=x&other=y"; - let out = dedupe_query_params(input); - let url = url::Url::parse(&out).unwrap(); - let pairs: Vec<_> = url.query_pairs().collect(); - assert_eq!(pairs.len(), 2); - } + move |error| OAuthError::Rmcp(format!("{context}: {error}")) } diff --git a/crates/aether-auth/src/mcp/mod.rs b/crates/aether-auth/src/mcp/mod.rs index 6c756cf73..a22130f16 100644 --- a/crates/aether-auth/src/mcp/mod.rs +++ b/crates/aether-auth/src/mcp/mod.rs @@ -2,4 +2,4 @@ mod credential_store; mod integration; pub use credential_store::McpCredentialStore; -pub use integration::{create_auth_manager_from_store, perform_oauth_flow}; +pub use integration::{OAuthClientRegistration, OAuthFlowOptions, create_auth_manager_from_store, perform_oauth_flow}; diff --git a/crates/aether-auth/src/oauth.md b/crates/aether-auth/src/oauth.md index 62067e80b..194709e45 100644 --- a/crates/aether-auth/src/oauth.md +++ b/crates/aether-auth/src/oauth.md @@ -2,7 +2,7 @@ OAuth 2.0 authentication primitives for the Aether agent framework. # Architecture -- [`OAuthHandler`] -- Trait implemented by consuming applications to handle OAuth UI/UX. The handler opens a browser and waits for the authorization code on a local port. +- [`OAuthHandler`] -- Trait implemented by consuming applications to handle OAuth UI/UX. The handler opens a browser and returns the absolute callback URL. - [`BrowserOAuthHandler`] -- Default implementation that opens the system browser and listens on a dynamic local port. - [`OAuthCredentialStorage`] -- Trait for persisting OAuth credentials keyed by provider ID, MCP server ID, or another credential key. - [`OsKeyringStore`] -- OS-keychain-backed [`OAuthCredentialStorage`] (macOS Keychain, Windows Credential Manager, Linux/FreeBSD Secret Service). Available under the `keyring` feature. @@ -12,8 +12,8 @@ OAuth 2.0 authentication primitives for the Aether agent framework. Behind the `mcp` feature: - [`McpCredentialStore`] -- Per-server adapter that binds an [`OAuthCredentialStorage`] to one MCP server ID and implements `rmcp::transport::auth::CredentialStore`. -- [`perform_oauth_flow`] -- Orchestrates the full MCP authorization code flow: browser launch, callback capture, token exchange, and credential storage. -- [`create_auth_manager_from_store`] -- Build an `AuthorizationManager` from stored credentials, handling automatic token refresh. +- [`perform_oauth_flow`] -- Integrates Aether's browser/callback UI with `rmcp`'s MCP OAuth state machine. `rmcp` provides discovery, registration selection (pre-registered client, CIMD, then DCR), PKCE, resource indicators, scope selection and unioning, refresh, and issuer validation. +- [`create_auth_manager_from_store`] -- Build an issuer-bound `AuthorizationManager` from stored credentials, handling automatic token refresh and rejecting credentials minted by a different authorization server. # Errors diff --git a/crates/aether-auth/tests/mcp_oauth.rs b/crates/aether-auth/tests/mcp_oauth.rs index 42d115b33..7622900e3 100644 --- a/crates/aether-auth/tests/mcp_oauth.rs +++ b/crates/aether-auth/tests/mcp_oauth.rs @@ -1,8 +1,8 @@ #![cfg(feature = "mcp")] use aether_auth::{ - FakeOAuthCredentialStore, OAuthCallback, OAuthCredentialStorage, OAuthError, OAuthHandler, - create_auth_manager_from_store, perform_oauth_flow, + FakeOAuthCredentialStore, OAuthClientRegistration, OAuthCredentialStorage, OAuthError, OAuthFlowOptions, + OAuthHandler, create_auth_manager_from_store, perform_oauth_flow, }; use futures::future::BoxFuture; use rmcp::transport::auth::StoredCredentials; @@ -11,22 +11,28 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; #[tokio::test] async fn configured_client_id_ignores_credentials_for_another_client() { - let stored = StoredCredentials::new("old-client".to_string(), None, Vec::new(), None); + let stored = StoredCredentials::new("old-client".to_string(), None, Vec::new(), None) + .with_issuer(Some("https://old.example".to_string())); let store = Arc::new(FakeOAuthCredentialStore::new().with_value("mcp:slack", serde_json::to_value(stored).unwrap())); - let manager = - create_auth_manager_from_store("slack", "https://mcp.slack.com/mcp", Some("configured-client"), store.clone()) - .await - .unwrap(); - - assert!(manager.is_none()); - let preserved = store.load("mcp:slack").await.unwrap().expect("mismatched credential should be preserved"); - assert_eq!(preserved["client_id"], "old-client"); + let restored = create_auth_manager_from_store( + "slack", + "https://mcp.slack.com/mcp", + Some("configured-client"), + "http://localhost:3118/", + store.clone(), + ) + .await + .unwrap(); + + assert!(restored.is_none()); + assert_eq!(store.load("mcp:slack").await.unwrap().unwrap()["client_id"], "old-client"); } struct FakeHandler { auth_url: Arc>>, + issuer: Option, } impl OAuthHandler for FakeHandler { @@ -34,14 +40,28 @@ impl OAuthHandler for FakeHandler { "http://localhost:3118/" } - fn authorize(&self, auth_url: &str) -> BoxFuture<'_, Result> { + fn authorize(&self, auth_url: &str) -> BoxFuture<'_, Result> { *self.auth_url.lock().unwrap() = Some(auth_url.to_string()); let state = url::Url::parse(auth_url).unwrap().query_pairs().find(|(name, _)| name == "state").unwrap().1.into_owned(); - Box::pin(async move { Ok(OAuthCallback { code: "test-code".to_string(), state }) }) + let issuer = self.issuer.clone(); + Box::pin(async move { + let mut callback = url::Url::parse("http://localhost:3118/").unwrap(); + callback.query_pairs_mut().append_pair("code", "test-code").append_pair("state", &state); + if let Some(issuer) = issuer { + callback.query_pairs_mut().append_pair("iss", &issuer); + } + Ok(callback.to_string()) + }) } } +#[derive(Debug, Clone, Copy, Default)] +struct OAuthServerOptions { + supports_cimd: bool, + requires_response_issuer: bool, +} + struct OAuthServer { base_url: String, requests: Arc>>, @@ -49,7 +69,8 @@ struct OAuthServer { } impl OAuthServer { - async fn bind() -> Self { + async fn bind(options: OAuthServerOptions) -> Self { + let OAuthServerOptions { supports_cimd, requires_response_issuer } = options; let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let origin = format!("http://{}", listener.local_addr().unwrap()); let base_url = format!("{origin}/mcp"); @@ -60,10 +81,8 @@ impl OAuthServer { let mut buffer = vec![0; 8192]; let read = stream.read(&mut buffer).await.unwrap(); let request = String::from_utf8_lossy(&buffer[..read]); - let request_line = request.lines().next().unwrap().to_string(); - captured_requests.lock().unwrap().push(request_line.clone()); - let path = request_line.split_whitespace().nth(1).unwrap(); - + captured_requests.lock().unwrap().push(request.to_string()); + let path = request.lines().next().unwrap().split_whitespace().nth(1).unwrap(); let (status, headers, body) = if path == "/mcp" { ( "401 Unauthorized", @@ -79,18 +98,14 @@ impl OAuthServer { "authorization_servers": [&origin] }) } else if path == "/token" { - serde_json::json!({ - "access_token": "access-token", - "token_type": "Bearer", - "expires_in": 3600 - }) + serde_json::json!({"access_token": "access-token", "token_type": "Bearer", "expires_in": 3600}) } else if path == "/register" { serde_json::json!({ "client_id": "registered-client", "redirect_uris": ["http://localhost:3118/"] }) } else { - serde_json::json!({ + let mut metadata = serde_json::json!({ "issuer": origin, "authorization_endpoint": format!("{origin}/authorize"), "token_endpoint": format!("{origin}/token"), @@ -98,7 +113,14 @@ impl OAuthServer { "response_types_supported": ["code"], "code_challenge_methods_supported": ["S256"], "scopes_supported": ["openid"] - }) + }); + if supports_cimd { + metadata["client_id_metadata_document_supported"] = serde_json::Value::Bool(true); + } + if requires_response_issuer { + metadata["authorization_response_iss_parameter_supported"] = serde_json::Value::Bool(true); + } + metadata } .to_string(); ("200 OK", String::new(), body) @@ -122,30 +144,86 @@ impl Drop for OAuthServer { } #[tokio::test] -async fn no_client_id_uses_dynamic_registration() { - let server = OAuthServer::bind().await; +async fn configured_cimd_is_forwarded_to_rmcp() { + let server = OAuthServer::bind(OAuthServerOptions { supports_cimd: true, ..OAuthServerOptions::default() }).await; let auth_url = Arc::new(Mutex::new(None)); - let handler = FakeHandler { auth_url: Arc::clone(&auth_url) }; + let handler = FakeHandler { auth_url: Arc::clone(&auth_url), issuer: None }; + let metadata_url = "https://aether-agent.io/oauth/client-metadata.json"; + + perform_oauth_flow( + "slack", + &server.base_url, + &handler, + OAuthFlowOptions { + client_registration: OAuthClientRegistration::ClientMetadata(metadata_url.to_string()), + ..OAuthFlowOptions::default() + }, + None, + ) + .await + .unwrap(); + + let auth_url = url::Url::parse(auth_url.lock().unwrap().as_ref().unwrap()).unwrap(); + assert!(auth_url.query_pairs().any(|(name, value)| name == "client_id" && value == metadata_url)); + assert!(!server.requests.lock().unwrap().iter().any(|request| request.starts_with("POST /register"))); +} - perform_oauth_flow("slack", &server.base_url, &handler, None, None).await.unwrap(); +#[tokio::test] +async fn dynamic_registration_is_explicitly_selected() { + let server = OAuthServer::bind(OAuthServerOptions::default()).await; + let auth_url = Arc::new(Mutex::new(None)); + let handler = FakeHandler { auth_url: Arc::clone(&auth_url), issuer: None }; - let auth_url = auth_url.lock().unwrap().clone().unwrap(); - let auth_url = url::Url::parse(&auth_url).unwrap(); + perform_oauth_flow("slack", &server.base_url, &handler, OAuthFlowOptions::default(), None).await.unwrap(); + + let auth_url = url::Url::parse(auth_url.lock().unwrap().as_ref().unwrap()).unwrap(); assert!(auth_url.query_pairs().any(|(name, value)| name == "client_id" && value == "registered-client")); assert!(server.requests.lock().unwrap().iter().any(|request| request.starts_with("POST /register"))); } #[tokio::test] -async fn configured_client_id_skips_dynamic_registration() { - let server = OAuthServer::bind().await; +async fn configured_client_id_is_forwarded_to_rmcp() { + let server = OAuthServer::bind(OAuthServerOptions::default()).await; let auth_url = Arc::new(Mutex::new(None)); - let handler = FakeHandler { auth_url: Arc::clone(&auth_url) }; - - perform_oauth_flow("slack", &server.base_url, &handler, Some("static-client"), None).await.unwrap(); - - let auth_url = auth_url.lock().unwrap().clone().unwrap(); - let auth_url = url::Url::parse(&auth_url).unwrap(); + let handler = FakeHandler { auth_url: Arc::clone(&auth_url), issuer: None }; + + perform_oauth_flow( + "slack", + &server.base_url, + &handler, + OAuthFlowOptions { + client_registration: OAuthClientRegistration::PreRegistered("static-client".to_string()), + ..OAuthFlowOptions::default() + }, + None, + ) + .await + .unwrap(); + + let auth_url = url::Url::parse(auth_url.lock().unwrap().as_ref().unwrap()).unwrap(); assert!(auth_url.query_pairs().any(|(name, value)| name == "client_id" && value == "static-client")); - assert!(auth_url.query_pairs().any(|(name, value)| name == "redirect_uri" && value == "http://localhost:3118/")); assert!(!server.requests.lock().unwrap().iter().any(|request| request.starts_with("POST /register"))); } + +#[tokio::test] +async fn callback_url_is_delegated_to_rmcp_for_issuer_validation() { + let server = + OAuthServer::bind(OAuthServerOptions { requires_response_issuer: true, ..OAuthServerOptions::default() }).await; + let issuer = server.base_url.trim_end_matches("/mcp").to_string(); + let handler = FakeHandler { auth_url: Arc::new(Mutex::new(None)), issuer: Some(issuer) }; + + perform_oauth_flow( + "slack", + &server.base_url, + &handler, + OAuthFlowOptions { + client_registration: OAuthClientRegistration::PreRegistered("static-client".to_string()), + ..OAuthFlowOptions::default() + }, + None, + ) + .await + .unwrap(); + + assert!(server.requests.lock().unwrap().iter().any(|request| request.starts_with("POST /token"))); +} diff --git a/crates/aether-core/tests/mcp/config_parser_tests.rs b/crates/aether-core/tests/mcp/config_parser_tests.rs index 44751d316..1f92c9660 100644 --- a/crates/aether-core/tests/mcp/config_parser_tests.rs +++ b/crates/aether-core/tests/mcp/config_parser_tests.rs @@ -1,6 +1,7 @@ use mcp_utils::client::{McpConfig, McpHttpConfig, McpServer, McpTransport, ParseError}; use std::collections::HashMap; use std::env; +use std::num::NonZeroU16; use utils::variables::Vars; async fn parse_servers(json: &str) -> Result, ParseError> { @@ -80,8 +81,8 @@ async fn test_parse_http_oauth_config() { let server = parse_one(&json).await; match server.transport { McpTransport::Http(McpHttpConfig { oauth: Some(oauth), .. }) => { - assert_eq!(oauth.client_id, "1601185624273.8899143856786"); - assert_eq!(oauth.callback_port.get(), 3118); + assert_eq!(oauth.client_id.as_deref(), Some("1601185624273.8899143856786")); + assert_eq!(oauth.callback_port.map(NonZeroU16::get), Some(3118)); } other => panic!("Expected HTTP OAuth config, got {other:?}"), } diff --git a/crates/aether-core/tests/mcp/oauth_tests.rs b/crates/aether-core/tests/mcp/oauth_tests.rs index a9bbc8bfb..b588efc6b 100644 --- a/crates/aether-core/tests/mcp/oauth_tests.rs +++ b/crates/aether-core/tests/mcp/oauth_tests.rs @@ -1,4 +1,4 @@ -use aether_auth::{OAuthCallback, OAuthError, OAuthHandler, accept_oauth_callback}; +use aether_auth::{OAuthError, OAuthHandler, accept_oauth_callback}; use aether_core::mcp::mcp; use aether_core::testing::{FakeMcpServer, fake_mcp}; use futures::future::BoxFuture; @@ -95,31 +95,6 @@ fn test_manager_with_home(with_oauth: bool) -> (tempfile::TempDir, McpManager) { (home, manager) } -struct FakeOAuthHandler { - callback: OAuthCallback, - redirect_uri: String, -} - -impl FakeOAuthHandler { - fn new(code: &str, state: &str) -> Self { - Self { - callback: OAuthCallback { code: code.to_string(), state: state.to_string() }, - redirect_uri: "http://127.0.0.1:0/oauth2callback".to_string(), - } - } -} - -impl OAuthHandler for FakeOAuthHandler { - fn redirect_uri(&self) -> &str { - &self.redirect_uri - } - - fn authorize(&self, _auth_url: &str) -> BoxFuture<'_, Result> { - let callback = self.callback.clone(); - Box::pin(async move { Ok(callback) }) - } -} - struct CancellingOAuthHandler; impl OAuthHandler for CancellingOAuthHandler { @@ -127,7 +102,7 @@ impl OAuthHandler for CancellingOAuthHandler { "http://127.0.0.1:0/oauth2callback" } - fn authorize(&self, _auth_url: &str) -> BoxFuture<'_, Result> { + fn authorize(&self, _auth_url: &str) -> BoxFuture<'_, Result> { Box::pin(async { Err(OAuthError::UserCancelled) }) } } @@ -136,15 +111,6 @@ fn fake_oauth_handler_factory() -> OAuthHandlerFactory { Arc::new(|_ctx| Ok(Arc::new(CancellingOAuthHandler))) } -#[tokio::test] -async fn fake_oauth_handler_returns_configured_callback() { - let handler = FakeOAuthHandler::new("test_code", "test_state"); - let result = handler.authorize("https://example.com/auth").await; - let callback = result.unwrap(); - assert_eq!(callback.code, "test_code"); - assert_eq!(callback.state, "test_state"); -} - #[tokio::test] async fn cancelling_handler_returns_user_cancelled() { let handler = CancellingOAuthHandler; @@ -154,7 +120,7 @@ async fn cancelling_handler_returns_user_cancelled() { #[tokio::test] async fn builder_with_oauth_handler_factory_spawns_successfully() { - let handler = Arc::new(FakeOAuthHandler::new("code", "state")); + let handler = Arc::new(CancellingOAuthHandler); let mut spawn = mcp("/workspace") .with_oauth_handler_factory(Arc::new(move |_ctx| Ok(handler.clone()))) @@ -182,7 +148,9 @@ async fn real_401_challenge_is_classified_as_needs_oauth() { ) .await; let mut manager = test_manager(true); + manager.add_mcps(vec![endpoint.server("protected")]).await.unwrap(); + let status = &manager.server_statuses()[0]; assert!(matches!(status.status, McpServerStatus::NeedsOAuth)); assert_eq!(status.auth_capability, McpServerAuthCapability::OAuth); @@ -236,17 +204,11 @@ async fn accept_oauth_callback_parses_code_and_state() { let _response = client.get(&callback_url).send().await.expect("Failed to send callback request"); let result = handle.await.unwrap(); - let callback = result.unwrap(); - assert_eq!(callback.code, "abc123"); - assert_eq!(callback.state, "csrf_token"); -} - -#[tokio::test] -async fn oauth_handler_is_dyn_compatible() { - let handler: Arc = Arc::new(FakeOAuthHandler::new("code", "state")); - let result = handler.authorize("https://example.com").await; - assert!(result.is_ok()); - assert_eq!(result.unwrap().code, "code"); + let callback_url = result.unwrap(); + let callback = reqwest::Url::parse(&callback_url).unwrap(); + let params = callback.query_pairs().collect::>(); + assert_eq!(params.get("code").map(AsRef::as_ref), Some("abc123")); + assert_eq!(params.get("state").map(AsRef::as_ref), Some("csrf_token")); } #[tokio::test] diff --git a/crates/llm/src/providers/codex/oauth.rs b/crates/llm/src/providers/codex/oauth.rs index ed3401d25..d99976da8 100644 --- a/crates/llm/src/providers/codex/oauth.rs +++ b/crates/llm/src/providers/codex/oauth.rs @@ -42,11 +42,22 @@ pub async fn perform_codex_oauth_flow(store: &dyn OAuthCredentialStorage) -> Res // Port 1455 is hardcoded because the Codex API has a fixed redirect URI // (http://localhost:1455/auth/callback) registered with OpenAI's OAuth server. let handler = BrowserOAuthHandler::with_redirect_uri(REDIRECT_URI, 1455)?; - let callback = handler.authorize(auth_url.as_str()).await?; - - if callback.state != state { + let callback_url = handler.authorize(auth_url.as_str()).await?; + let callback = Url::parse(&callback_url) + .map_err(|error| OAuthError::InvalidCallback(format!("Invalid callback URL: {error}")))?; + let mut code = None; + let mut callback_state = None; + for (name, value) in callback.query_pairs() { + match name.as_ref() { + "code" => code = Some(value.into_owned()), + "state" => callback_state = Some(value.into_owned()), + _ => {} + } + } + if callback_state.as_deref() != Some(&state) { return Err(OAuthError::StateMismatch.into()); } + let code = code.ok_or_else(|| OAuthError::InvalidCallback("No authorization code in callback".to_string()))?; let oauth_client = BasicClient::new(ClientId::new(CLIENT_ID.to_string())) .set_auth_uri( @@ -65,7 +76,7 @@ pub async fn perform_codex_oauth_flow(store: &dyn OAuthCredentialStorage) -> Res let http_client = oauth_http_client()?; let token_response = oauth_client - .exchange_code(AuthorizationCode::new(callback.code)) + .exchange_code(AuthorizationCode::new(code)) .set_pkce_verifier(pkce_verifier) .request_async(&http_client) .await diff --git a/crates/mcp-utils/src/client/config.rs b/crates/mcp-utils/src/client/config.rs index 11ee54792..c4ff1704d 100644 --- a/crates/mcp-utils/src/client/config.rs +++ b/crates/mcp-utils/src/client/config.rs @@ -1,3 +1,4 @@ +use aether_auth::OAuthClientRegistration; use futures::future::BoxFuture; use rmcp::{RoleServer, service::DynService, transport::streamable_http_client::StreamableHttpClientTransportConfig}; use schemars::JsonSchema; @@ -48,11 +49,18 @@ pub struct StdioServerConfig { pub proxy: ToolExposure, } -#[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)] +pub const AETHER_OAUTH_CLIENT_METADATA_URL: &str = "https://aether-agent.io/oauth/client-metadata.json"; +pub const AETHER_OAUTH_CALLBACK_PORT: NonZeroU16 = NonZeroU16::new(3118).unwrap(); + +#[derive(Debug, Clone, Default, Deserialize, Serialize, JsonSchema, PartialEq)] #[serde(rename_all = "camelCase", deny_unknown_fields)] pub struct McpOAuthConfig { - pub client_id: String, - pub callback_port: NonZeroU16, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub client_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub client_metadata_url: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub callback_port: Option, } #[derive(Debug, Clone, Deserialize, Serialize, JsonSchema, PartialEq)] @@ -159,13 +167,41 @@ pub struct McpHttpConfig { pub oauth: Option, } -impl McpHttpConfig { - pub fn oauth_client_id(&self) -> Option<&str> { - self.oauth.as_ref().map(|oauth| oauth.client_id.as_str()) +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ResolvedOAuth { + pub client_registration: OAuthClientRegistration, + pub callback_port: NonZeroU16, +} + +impl ResolvedOAuth { + pub fn redirect_uri(&self) -> String { + loopback_redirect_uri(self.callback_port.get()) } +} - pub fn callback_port(&self) -> Option { - self.oauth.as_ref().map(|oauth| oauth.callback_port) +pub fn loopback_redirect_uri(port: u16) -> String { + format!("http://localhost:{port}/") +} + +impl McpHttpConfig { + pub fn resolved_oauth(&self) -> Option { + if self.transport.auth_header.is_some() { + return None; + } + let oauth = self.oauth.as_ref(); + let client_registration = match oauth { + Some(McpOAuthConfig { client_id: Some(client_id), .. }) => { + OAuthClientRegistration::PreRegistered(client_id.clone()) + } + Some(McpOAuthConfig { client_metadata_url: Some(url), .. }) => { + OAuthClientRegistration::ClientMetadata(url.clone()) + } + _ => OAuthClientRegistration::ClientMetadata(AETHER_OAUTH_CLIENT_METADATA_URL.to_string()), + }; + Some(ResolvedOAuth { + client_registration, + callback_port: oauth.and_then(|oauth| oauth.callback_port).unwrap_or(AETHER_OAUTH_CALLBACK_PORT), + }) } } @@ -439,7 +475,14 @@ impl McpServerConfig { let oauth = oauth .map(|oauth| -> Result { - Ok(McpOAuthConfig { client_id: vars.expand(&oauth.client_id)?, ..oauth }) + Ok(McpOAuthConfig { + client_id: oauth.client_id.map(|value| vars.expand(&value)).transpose()?, + client_metadata_url: oauth + .client_metadata_url + .map(|value| vars.expand(&value)) + .transpose()?, + callback_port: oauth.callback_port, + }) }) .transpose()?; diff --git a/crates/mcp-utils/src/client/connection.rs b/crates/mcp-utils/src/client/connection.rs index 205f82f94..faed2d8ac 100644 --- a/crates/mcp-utils/src/client/connection.rs +++ b/crates/mcp-utils/src/client/connection.rs @@ -79,7 +79,7 @@ pub struct McpConnectAttempt { pub enum McpConnectOutcome { Connected { conn: McpServerConnection, reauth_config: Option }, - NeedsOAuth { config: McpHttpConfig, error: McpError }, + NeedsOAuth { config: McpHttpConfig, challenge: Option, error: McpError }, Failed { error: McpError }, } @@ -149,15 +149,23 @@ pub(super) async fn connect_server(server: McpServer, ctx: &ConnectConfig) -> Mc McpConnectAttempt { name, outcome: outcome.with_reauth(reauth_config) } } -pub async fn authenticate_http(name: String, config: McpHttpConfig, ctx: Arc) -> McpConnectAttempt { +pub async fn authenticate_http( + name: String, + config: McpHttpConfig, + challenge: Option, + ctx: Arc, +) -> McpConnectAttempt { let outcome = match async { let factory = ctx .oauth_handler_factory .as_ref() .ok_or_else(|| McpError::ConnectionFailed(format!("No OAuth handler factory available for '{name}'")))?; + let oauth = config + .resolved_oauth() + .ok_or_else(|| McpError::ConnectionFailed(format!("OAuth is not available for '{name}'")))?; let handler = factory(OAuthHandlerContext { server_name: name.clone(), - callback_port: config.callback_port(), + callback_port: Some(oauth.callback_port), tx: ctx.event_sender.clone(), })?; @@ -165,7 +173,7 @@ pub async fn authenticate_http(name: String, config: McpHttpConfig, ctx: Arc>, ) -> McpConnectOutcome { let conn_err = |e| McpError::ConnectionFailed(format!("HTTP MCP server {name}: {e}")); - let stored_auth_manager = if let Some(store) = oauth_credential_store - && config.transport.auth_header.is_none() - { - match create_auth_manager_from_store(name, &config.transport.uri, config.oauth_client_id(), Arc::clone(store)) - .await + let oauth = config.resolved_oauth(); + let restored = if let (Some(store), Some(oauth)) = (oauth_credential_store, oauth.as_ref()) { + match create_auth_manager_from_store( + name, + &config.transport.uri, + oauth.client_registration.pre_registered_client_id(), + &oauth.redirect_uri(), + Arc::clone(store), + ) + .await { Ok(manager) => manager, Err(e) => { @@ -280,9 +293,9 @@ async fn connect_http( } else { None }; - let result = if let Some(auth_manager) = stored_auth_manager { + let result = if let Some(manager) = restored { tracing::debug!("Using OAuth for server '{name}'"); - let auth_client = AuthClient::new(reqwest::Client::default(), auth_manager); + let auth_client = AuthClient::new(reqwest::Client::default(), manager); let transport = StreamableHttpClientTransport::with_client(auth_client, config.transport.clone()); serve_client_with_lifecycle(mcp_client, transport, client_lifecycle_mode()).await } else { @@ -295,11 +308,12 @@ async fn connect_http( McpConnectOutcome::Connected { conn: McpServerConnection::from_parts(client, None), reauth_config: None } } Err(error) => { - let authorization_required = error.is_authorization_required() || error.auth_challenge().is_some(); + let challenge = error.auth_challenge().map(str::to_string); + let authorization_required = error.is_authorization_required(); let error = conn_err(error); tracing::warn!("Failed to connect to MCP server '{name}': {error}"); - if oauth_handler_factory.is_some() && config.transport.auth_header.is_none() && authorization_required { - McpConnectOutcome::NeedsOAuth { config, error } + if oauth_handler_factory.is_some() && oauth.is_some() && (authorization_required || challenge.is_some()) { + McpConnectOutcome::NeedsOAuth { config, challenge, error } } else { McpConnectOutcome::Failed { error } } diff --git a/crates/mcp-utils/src/client/manager.rs b/crates/mcp-utils/src/client/manager.rs index 8ed282452..d2f9d56c1 100644 --- a/crates/mcp-utils/src/client/manager.rs +++ b/crates/mcp-utils/src/client/manager.rs @@ -241,12 +241,13 @@ impl McpManager { let name = name.to_string(); let config = record.reauth_config.clone().expect("checked above"); + let challenge = record.oauth_challenge.clone(); let ctx = self.connect_config(); self.set_state(&name, ServerState::Authenticating); self.emit_server_statuses_changed().await; - Ok(async move { authenticate_http(name, config, ctx).await }) + Ok(async move { authenticate_http(name, config, challenge, ctx).await }) } pub async fn apply_connection_attempt(&mut self, attempt: McpConnectAttempt) { @@ -262,11 +263,12 @@ impl McpManager { Err(error) => self.apply_authentication_failure(name, error.to_string()).await, } } - McpConnectOutcome::NeedsOAuth { config, error } => { + McpConnectOutcome::NeedsOAuth { config, challenge, error } => { tracing::warn!("Server '{name}' needs OAuth: {error}"); if let Some(record) = self.servers.get_mut(&name) { record.state = ServerState::NeedsOAuth; record.reauth_config = Some(config); + record.oauth_challenge = challenge; } self.emit_server_statuses_changed().await; } @@ -495,6 +497,7 @@ impl McpManager { let record = self.servers.get_mut(name).expect("record checked above"); record.reauth_config = reauth_config.or_else(|| record.reauth_config.clone()); + record.oauth_challenge = None; record.state = ServerState::Connected { connection: conn, tools: tools.iter().map(Tool::from).collect() }; Ok(()) } @@ -572,6 +575,7 @@ impl Drop for McpManager { struct ServerRecord { state: ServerState, reauth_config: Option, + oauth_challenge: Option, exposure: ToolExposure, } @@ -597,7 +601,7 @@ impl From<&ServerState> for McpServerStatus { impl ServerRecord { fn new(state: ServerState, reauth_config: Option, exposure: ToolExposure) -> Self { - Self { state, reauth_config, exposure } + Self { state, reauth_config, oauth_challenge: None, exposure } } fn tools(&self) -> &[Tool] { @@ -681,7 +685,7 @@ mod tests { use crate::client::config::{McpHttpConfig, McpServer, McpTransport, ToolExposure}; use crate::client::connection::{McpConnectAttempt, McpConnectOutcome}; use crate::status::McpServerAuthCapability; - use aether_auth::{OAuthCallback, OAuthError, OAuthHandler}; + use aether_auth::{OAuthError, OAuthHandler}; use futures::future::BoxFuture; use rmcp::{ Json, RoleServer, ServerHandler, @@ -763,7 +767,7 @@ mod tests { "http://127.0.0.1:0/oauth2callback" } - fn authorize(&self, _auth_url: &str) -> BoxFuture<'_, Result> { + fn authorize(&self, _auth_url: &str) -> BoxFuture<'_, Result> { Box::pin(async { Err(OAuthError::UserCancelled) }) } } diff --git a/crates/mcp-utils/src/client/mod.rs b/crates/mcp-utils/src/client/mod.rs index 1154252c7..af0fcb71e 100644 --- a/crates/mcp-utils/src/client/mod.rs +++ b/crates/mcp-utils/src/client/mod.rs @@ -15,9 +15,10 @@ mod tool_proxy; pub use call_tool::{CallToolError, CallToolOptions, ToolCallEvent, call_tool}; pub use config::{ - InMemoryServerConfig, InMemoryType, McpConfig, McpHttpConfig, McpOAuthConfig, McpServer, McpServerCloneError, - McpServerConfig, McpTransport, ParseError, RemoteServerConfig, RemoteType, ServerFactory, StdioServerConfig, - StdioType, ToolExposure, ToolProxyRules, + AETHER_OAUTH_CALLBACK_PORT, AETHER_OAUTH_CLIENT_METADATA_URL, InMemoryServerConfig, InMemoryType, McpConfig, + McpHttpConfig, McpOAuthConfig, McpServer, McpServerCloneError, McpServerConfig, McpTransport, ParseError, + RemoteServerConfig, RemoteType, ResolvedOAuth, ServerFactory, StdioServerConfig, StdioType, ToolExposure, + ToolProxyRules, loopback_redirect_uri, }; pub use connection::{McpConnectAttempt, McpConnectOutcome, McpServerConnection}; pub use connection_attempt_manager::McpConnectionAttemptManager; diff --git a/crates/mcp-utils/src/client/oauth_handler.rs b/crates/mcp-utils/src/client/oauth_handler.rs index 279cdcb51..bab5ed48b 100644 --- a/crates/mcp-utils/src/client/oauth_handler.rs +++ b/crates/mcp-utils/src/client/oauth_handler.rs @@ -1,5 +1,6 @@ +use crate::client::config::loopback_redirect_uri; use crate::client::manager::{ElicitationRequest, McpClientEvent, OAuthHandlerContext}; -use aether_auth::{OAuthCallback, OAuthError, OAuthHandler, accept_oauth_callback}; +use aether_auth::{OAuthError, OAuthHandler, accept_oauth_callback}; use futures::future::BoxFuture; use rmcp::model::{ElicitRequestParams, ElicitationAction}; use std::num::NonZeroU16; @@ -8,7 +9,7 @@ use tokio::sync::{mpsc, oneshot}; const AETHER_OAUTH_ELICITATION_ID: &str = "aether-oauth"; -/// `OAuthHandler` that dispatches the OAuth authorization URL to the host +/// `OAuthHandler` that dispatches the OAuth authorization URL to the host. pub struct ElicitingOAuthHandler { listener: TcpListener, redirect_uri: String, @@ -18,17 +19,13 @@ pub struct ElicitingOAuthHandler { impl ElicitingOAuthHandler { pub fn new(ctx: OAuthHandlerContext) -> Result { - let (port, listener) = { - let port = ctx.callback_port.map_or(0, NonZeroU16::get); - let std_listener = std::net::TcpListener::bind(("127.0.0.1", port))?; - let port = std_listener.local_addr()?.port(); - std_listener.set_nonblocking(true)?; - (port, TcpListener::from_std(std_listener)?) - }; - + let port = ctx.callback_port.map_or(0, NonZeroU16::get); + let std_listener = std::net::TcpListener::bind(("127.0.0.1", port))?; + let port = std_listener.local_addr()?.port(); + std_listener.set_nonblocking(true)?; Ok(Self { - listener, - redirect_uri: format!("http://localhost:{port}/"), + listener: TcpListener::from_std(std_listener)?, + redirect_uri: loopback_redirect_uri(port), server_name: ctx.server_name, event_sender: ctx.tx, }) @@ -40,27 +37,25 @@ impl OAuthHandler for ElicitingOAuthHandler { &self.redirect_uri } - fn authorize(&self, auth_url: &str) -> BoxFuture<'_, Result> { + fn authorize(&self, auth_url: &str) -> BoxFuture<'_, Result> { let auth_url = auth_url.to_string(); Box::pin(async move { let (response_sender, response_rx) = oneshot::channel(); - let request = ElicitationRequest { - server_name: self.server_name.clone(), - request: ElicitRequestParams::UrlElicitationParams { - meta: None, - message: "Open this URL to authorize MCP server access.".to_string(), - url: auth_url, - elicitation_id: AETHER_OAUTH_ELICITATION_ID.to_string(), - }, - response_sender, - }; - self.event_sender - .send(McpClientEvent::Elicitation(request)) + .send(McpClientEvent::Elicitation(ElicitationRequest { + server_name: self.server_name.clone(), + request: ElicitRequestParams::UrlElicitationParams { + meta: None, + message: "Open this URL to authorize MCP server access.".to_string(), + url: auth_url, + elicitation_id: AETHER_OAUTH_ELICITATION_ID.to_string(), + }, + response_sender, + })) .await .map_err(|_| OAuthError::Rmcp("OAuth prompt channel closed".to_string()))?; - let callback = tokio::select! { + tokio::select! { callback = accept_oauth_callback(&self.listener) => callback, response = response_rx => match response { Ok(result) if matches!(result.action, ElicitationAction::Decline | ElicitationAction::Cancel) => { @@ -68,9 +63,7 @@ impl OAuthHandler for ElicitingOAuthHandler { } Ok(_) | Err(_) => accept_oauth_callback(&self.listener).await, }, - }?; - - Ok(callback) + } }) } } @@ -83,7 +76,7 @@ mod tests { use tokio::{io::AsyncWriteExt, task::yield_now}; #[tokio::test] - async fn accepting_browser_prompt_keeps_waiting_for_oauth_callback() { + async fn accepting_browser_prompt_keeps_waiting_for_callback() { let (tx, mut rx) = mpsc::channel(1); let handler = Arc::new( ElicitingOAuthHandler::new(OAuthHandlerContext { @@ -100,7 +93,6 @@ mod tests { .unwrap() .parse::() .unwrap(); - let authorize = { let handler = Arc::clone(&handler); tokio::spawn(async move { handler.authorize("https://example.com/oauth").await }) @@ -110,47 +102,41 @@ mod tests { }; request.response_sender.send(ElicitResult::new(ElicitationAction::Accept)).unwrap(); yield_now().await; - assert!(!authorize.is_finished(), "accepting the prompt must not complete OAuth"); + assert!(!authorize.is_finished()); let mut callback = tokio::net::TcpStream::connect(("127.0.0.1", port)).await.unwrap(); callback.write_all(b"GET /?code=test-code&state=test-state HTTP/1.1\r\nHost: localhost\r\n\r\n").await.unwrap(); - let result = authorize.await.unwrap().unwrap(); - assert_eq!(result.code, "test-code"); - assert_eq!(result.state, "test-state"); + assert!(authorize.await.unwrap().unwrap().contains("code=test-code&state=test-state")); } #[tokio::test] - async fn configured_callback_port_uses_registered_localhost_redirect() { + async fn configured_callback_port_uses_registered_redirect() { let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); let port = probe.local_addr().unwrap().port(); drop(probe); let (tx, _) = mpsc::channel(1); - let handler = ElicitingOAuthHandler::new(OAuthHandlerContext { server_name: "slack".to_string(), callback_port: NonZeroU16::new(port), tx, }) .unwrap(); - assert_eq!(handler.redirect_uri(), format!("http://localhost:{port}/")); } #[tokio::test] - async fn configured_callback_port_fails_when_already_in_use() { + async fn configured_callback_port_fails_when_in_use() { let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); let port = listener.local_addr().unwrap().port(); let (tx, _) = mpsc::channel(1); - let error = ElicitingOAuthHandler::new(OAuthHandlerContext { server_name: "slack".to_string(), callback_port: NonZeroU16::new(port), tx, }) .err() - .expect("occupied callback port should fail"); - + .unwrap(); assert_eq!(error.kind(), std::io::ErrorKind::AddrInUse); } } diff --git a/crates/mcp-utils/src/docs/mcp_server_config.md b/crates/mcp-utils/src/docs/mcp_server_config.md index b79a6ea86..b4de98fa6 100644 --- a/crates/mcp-utils/src/docs/mcp_server_config.md +++ b/crates/mcp-utils/src/docs/mcp_server_config.md @@ -25,10 +25,18 @@ A remote streamable-HTTP server using a pre-registered public OAuth client: } ``` -When `oauth` is omitted, Aether uses dynamic client registration and a random -loopback callback port. When present, Aether binds the configured callback port -and advertises `http://localhost:/`, which must match the redirect -URI registered for the OAuth client. +When a remote HTTP server returns an OAuth challenge, Aether uses its first-party +Client ID Metadata Document at +`https://aether-agent.io/oauth/client-metadata.json` and listens on +`127.0.0.1:3118`, advertising the exact redirect URI `http://localhost:3118/`. +If port 3118 is occupied, authentication fails until the port is available. + +The `oauth` object is optional. Set `clientMetadataUrl` for a custom CIMD, or +`clientId` for a pre-registered public client; a configured `clientId` takes +priority. `callbackPort` defaults to 3118 and must exactly match the client's +registered redirect URI. If the authorization server does not advertise CIMD, +Aether falls back to deprecated Dynamic Client Registration. An explicit +`Authorization` header bypasses OAuth entirely. A remote server using a bearer token: diff --git a/crates/mcp-utils/tests/oauth_config.rs b/crates/mcp-utils/tests/oauth_config.rs new file mode 100644 index 000000000..e6986eb81 --- /dev/null +++ b/crates/mcp-utils/tests/oauth_config.rs @@ -0,0 +1,64 @@ +use aether_auth::OAuthClientRegistration; +use mcp_utils::client::{McpConfig, McpHttpConfig, McpTransport}; +use std::collections::HashMap; +use utils::variables::Vars; + +#[tokio::test] +async fn http_oauth_defaults_to_aether_cimd_and_fixed_callback() { + let config = parse_http_config(r#"{ "type": "http", "url": "https://example.com/mcp" }"#, &Vars::new()).await; + let oauth = config.resolved_oauth().expect("OAuth should be offered without an auth header"); + + assert_eq!( + oauth.client_registration, + OAuthClientRegistration::ClientMetadata("https://aether-agent.io/oauth/client-metadata.json".to_string()) + ); + assert_eq!(oauth.callback_port.get(), 3118); + assert_eq!(oauth.redirect_uri(), "http://localhost:3118/"); +} + +#[tokio::test] +async fn custom_cimd_expands_variables_and_preregistered_client_takes_priority() { + let vars = Vars::new().with("CLIENT_ID", "registered").with("CIMD_URL", "https://client.example/oauth/client.json"); + let config = parse_http_config( + r#"{ + "type": "http", + "url": "https://example.com/mcp", + "oauth": { + "clientId": "$CLIENT_ID", + "clientMetadataUrl": "$CIMD_URL", + "callbackPort": 4000 + } + }"#, + &vars, + ) + .await; + let oauth = config.resolved_oauth().expect("OAuth should be offered without an auth header"); + + assert_eq!(oauth.client_registration, OAuthClientRegistration::PreRegistered("registered".to_string())); + assert_eq!(oauth.callback_port.get(), 4000); + assert_eq!(config.oauth.unwrap().client_metadata_url.as_deref(), Some("https://client.example/oauth/client.json")); +} + +#[tokio::test] +async fn bearer_header_bypasses_oauth_defaults() { + let config = parse_http_config( + r#"{ + "type": "http", + "url": "https://example.com/mcp", + "headers": { "Authorization": "Bearer token" } + }"#, + &Vars::new(), + ) + .await; + + assert!(config.resolved_oauth().is_none()); +} + +async fn parse_http_config(server: &str, vars: &Vars) -> McpHttpConfig { + let json = format!(r#"{{ "servers": {{ "remote": {server} }} }}"#); + let mut servers = McpConfig::from_json(&json).unwrap().into_servers(&HashMap::new(), vars).await.unwrap(); + assert_eq!(servers.len(), 1); + let server = servers.remove(0); + let McpTransport::Http(config) = server.transport else { panic!("expected HTTP transport") }; + config +} diff --git a/packages/website/src/content/docs/aether/settings/mcp-servers.mdx b/packages/website/src/content/docs/aether/settings/mcp-servers.mdx index 4db3136c8..29e45d544 100644 --- a/packages/website/src/content/docs/aether/settings/mcp-servers.mdx +++ b/packages/website/src/content/docs/aether/settings/mcp-servers.mdx @@ -125,7 +125,7 @@ If a server entry has a `command` and no `type`, it is treated as `"type": "stdi | `type` | `"http"` | Required. | | `url` | `string` | Streamable HTTP endpoint URL. | | `headers` | `object` | Optional headers. Aether currently forwards the `Authorization` header to the transport. | - | `oauth` | `object` | Optional pre-registered public OAuth client. See [Pre-registered OAuth clients](#pre-registered-oauth-clients). | + | `oauth` | `object` | Optional custom CIMD or pre-registered public OAuth client settings. See [OAuth clients](#oauth-clients). | | `proxy` | `boolean \| object` | `true` proxies every tool; an object partitions tools with `include` and `exclude` patterns. | @@ -151,15 +151,17 @@ If a server entry has a `command` and no `type`, it is treated as `"type": "stdi | `type` | `"sse"` | Required. | | `url` | `string` | SSE endpoint URL. | | `headers` | `object` | Optional headers. Aether currently forwards the `Authorization` header to the transport. | - | `oauth` | `object` | Optional pre-registered public OAuth client. See [Pre-registered OAuth clients](#pre-registered-oauth-clients). | + | `oauth` | `object` | Optional custom CIMD or pre-registered public OAuth client settings. See [OAuth clients](#oauth-clients). | | `proxy` | `boolean \| object` | `true` proxies every tool; an object partitions tools with `include` and `exclude` patterns. | -## Pre-registered OAuth clients +## OAuth clients -For `http` and `sse` servers that require OAuth, Aether defaults to dynamic client registration (RFC 7591) with a random loopback callback port. Some providers (Slack, for example) only accept a pre-registered public client. Set `oauth` to bind a fixed callback port and advertise `http://localhost:/`, which must match the redirect URI registered for the client: +For `http` and `sse` servers that return an OAuth challenge, Aether defaults to its first-party Client ID Metadata Document (CIMD) at `https://aether-agent.io/oauth/client-metadata.json`. The callback listener binds `http://localhost:3118/`. If port 3118 is occupied, authentication fails rather than selecting an unregistered redirect URI. + +Set `oauth.clientMetadataUrl` to use a custom CIMD, or `oauth.clientId` for a public client pre-registered with the provider. A pre-registered client ID takes priority when both are configured: ```json title=".aether/mcp.json" { @@ -178,8 +180,11 @@ For `http` and `sse` servers that require OAuth, Aether defaults to dynamic clie | Field | Type | Description | |-------|------|-------------| -| `oauth.clientId` | `string` | OAuth client ID registered with the provider. Supports `$VAR` / `${VAR}` expansion. | -| `oauth.callbackPort` | `number` | Loopback TCP port for the OAuth callback. Must match the redirect URI registered for the client. | +| `oauth.clientId` | `string` | Public OAuth client ID registered with the provider. Takes priority over CIMD. Supports `$VAR` / `${VAR}` expansion. | +| `oauth.clientMetadataUrl` | `string` | HTTPS URL of a custom Client ID Metadata Document. Supports variable expansion. | +| `oauth.callbackPort` | `number` | Loopback callback port; defaults to `3118`. The advertised URI is `http://localhost:/` and must exactly match the client metadata or registration. | + +When the authorization server does not advertise CIMD support, Aether falls back to deprecated Dynamic Client Registration for compatibility. OAuth discovery, PKCE, scopes, `offline_access`, resource indicators, refresh, and issuer validation are handled by `rmcp`. An explicit bearer `Authorization` header remains authoritative and bypasses OAuth. ## Built-in servers