From c22da5dbdf6dce959262e7b5ae934e481da98313 Mon Sep 17 00:00:00 2001 From: brigon Date: Sun, 9 Aug 2026 15:22:11 +1200 Subject: [PATCH 1/4] Add native MCP server to TrailBase Add OAuth login flow for native MCP Replace MCP sidecar with native deployment Add native MCP schema and config tools Bind OAuth tokens to the MCP resource Scope MCP authentication to the MCP route Clarify portable MCP client configuration Use localhost in MCP quick start Use depot flag in MCP deployment examples Explain native MCP deployment arguments Add local and public MCP client examples Explain MCP connection naming --- Cargo.lock | 99 ++++ Cargo.toml | 1 + crates/cli/src/args.rs | 4 + crates/cli/src/bin/trail.rs | 1 + crates/core/Cargo.toml | 1 + crates/core/src/auth/api/login.rs | 2 +- crates/core/src/auth/api/mod.rs | 6 +- crates/core/src/auth/jwt.rs | 11 + crates/core/src/lib.rs | 1 + crates/core/src/mcp.rs | 908 ++++++++++++++++++++++++++++++ crates/core/src/server/mod.rs | 17 +- 11 files changed, 1044 insertions(+), 7 deletions(-) create mode 100644 crates/core/src/mcp.rs diff --git a/Cargo.lock b/Cargo.lock index b4736a9c5..fb73587aa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1883,6 +1883,16 @@ dependencies = [ "darling_macro 0.21.3", ] +[[package]] +name = "darling" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88490bf1b990d87eaaa7ac8aa887f629a08e7359765b4911faf63c3763347d23" +dependencies = [ + "darling_core 0.24.0", + "darling_macro 0.24.0", +] + [[package]] name = "darling_core" version = "0.20.11" @@ -1910,6 +1920,19 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "darling_core" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "084e274f91c482280130e1e34e0b8d6e66776a060d7b6de7b84289ca778868c4" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 3.0.3", +] + [[package]] name = "darling_macro" version = "0.20.11" @@ -1932,6 +1955,17 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "darling_macro" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68f5792fa0d41cd2325ce0ffa64f0a340eaebd4971a3a0c5e1ffd2cc488a355e" +dependencies = [ + "darling_core 0.24.0", + "quote", + "syn 3.0.3", +] + [[package]] name = "dashmap" version = "5.5.3" @@ -5074,6 +5108,12 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pastey" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" + [[package]] name = "path-clean" version = "1.0.1" @@ -6379,6 +6419,50 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "rmcp" +version = "3.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8dddc5b1924b9a59fba420166160ca2c4663a4e01803e52eda33070f56d63c8" +dependencies = [ + "async-trait", + "base64 0.23.1", + "bytes", + "chrono", + "futures", + "http", + "http-body", + "http-body-util", + "pastey", + "pin-project-lite", + "rand 0.10.2", + "rmcp-macros", + "schemars", + "serde", + "serde_json", + "sse-stream", + "thiserror 2.0.20", + "tokio", + "tokio-stream", + "tokio-util", + "tower-service", + "tracing", + "uuid", +] + +[[package]] +name = "rmcp-macros" +version = "3.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6898e24cd16342b59bfa8a53c2c04b9cf62fc8a2cfea57b9c038b09984bfc521" +dependencies = [ + "darling 0.24.0", + "proc-macro2", + "quote", + "serde_json", + "syn 3.0.3", +] + [[package]] name = "rquickjs" version = "0.12.2" @@ -6792,6 +6876,7 @@ version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a" dependencies = [ + "chrono", "dyn-clone", "indexmap", "ref-cast", @@ -7325,6 +7410,19 @@ dependencies = [ "recursive", ] +[[package]] +name = "sse-stream" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c123f296ade4ec4b8b0f6162116e6629f5146922ca5ab40ca9d3c2e73ab4761e" +dependencies = [ + "bytes", + "futures-util", + "http-body", + "http-body-util", + "pin-project-lite", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -8208,6 +8306,7 @@ dependencies = [ "rcgen", "regex", "reqwest 0.13.4", + "rmcp", "rusqlite", "rust-embed", "schemars", diff --git a/Cargo.toml b/Cargo.toml index c21f46c2d..92c85cbf1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -70,6 +70,7 @@ minijinja = { version = "2.1.2", default-features = false } parking_lot = { version = "0.12.3", default-features = false, features = ["send_guard", "arc_lock"] } rand = "^0.10.0" reqwest = { version = "0.13.1", default-features = false, features = ["rustls", "json"] } +rmcp = { version = "3.1.2", features = ["transport-streamable-http-server"] } rusqlite = { version = "0.40.0", default-features = false, features = ["bundled", "cache", "column_decltype", "functions", "backup", "preupdate_hook"] } rust-embed = { version = "8.4.0", default-features = false, features = ["mime-guess"] } serde = { version = "^1.0.203", features = ["derive", "rc"] } diff --git a/crates/cli/src/args.rs b/crates/cli/src/args.rs index a66933ec4..e65680b3a 100644 --- a/crates/cli/src/args.rs +++ b/crates/cli/src/args.rs @@ -134,6 +134,10 @@ pub struct ServerArgs { #[arg(long)] pub demo: bool, + /// Enable the authenticated MCP endpoint at /mcp on the admin server. + #[arg(long, env, default_value_t = false)] + pub mcp: bool, + #[arg(long, default_value_t = false)] pub stderr_logging: bool, diff --git a/crates/cli/src/bin/trail.rs b/crates/cli/src/bin/trail.rs index 578726829..c51481d18 100644 --- a/crates/cli/src/bin/trail.rs +++ b/crates/cli/src/bin/trail.rs @@ -92,6 +92,7 @@ async fn async_main( tls_key: None, tls_cert: None, custom_router: None, + enable_mcp: cmd.mcp, }, ) .await?; diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index 49b6ccff3..b1e51ed6e 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -81,6 +81,7 @@ quick_cache = "0.7.0" rand = { workspace = true } regex = "1.11.0" reqwest = { workspace = true } +rmcp = { workspace = true } rusqlite = { workspace = true } rust-embed = { workspace = true } serde = { workspace = true } diff --git a/crates/core/src/auth/api/login.rs b/crates/core/src/auth/api/login.rs index d43e3c2ad..916892dd0 100644 --- a/crates/core/src/auth/api/login.rs +++ b/crates/core/src/auth/api/login.rs @@ -356,7 +356,7 @@ pub(crate) async fn build_auth_token_flow_response( /// /// An example using the two-step "authentication code flow" with PKCE can be found in /// `/examples/blog/flutter`. -async fn build_authorization_code_flow_and_pkce_response( +pub(crate) async fn build_authorization_code_flow_and_pkce_response( state: &AppState, db_user: &DbUser, redirect: String, diff --git a/crates/core/src/auth/api/mod.rs b/crates/core/src/auth/api/mod.rs index 67a5031c6..868109e72 100644 --- a/crates/core/src/auth/api/mod.rs +++ b/crates/core/src/auth/api/mod.rs @@ -3,15 +3,15 @@ pub(super) mod change_email; pub(super) mod change_password; pub(super) mod change_username; pub(super) mod delete; -pub(super) mod login; +pub(crate) mod login; pub(super) mod login_anonymous; pub(super) mod logout; pub(super) mod otp; pub(super) mod promote_anonymous; -pub(super) mod refresh; +pub(crate) mod refresh; pub(super) mod register; pub(super) mod reset_password; pub(super) mod status; -pub(super) mod token; +pub(crate) mod token; pub(super) mod totp; pub(super) mod verify_email; diff --git a/crates/core/src/auth/jwt.rs b/crates/core/src/auth/jwt.rs index 8fca78636..c7c610ff4 100644 --- a/crates/core/src/auth/jwt.rs +++ b/crates/core/src/auth/jwt.rs @@ -311,6 +311,17 @@ impl JwtHelper { .map(|data| data.claims); } + pub(crate) fn decode_with_audience( + &self, + token: &str, + audience: &str, + ) -> Result { + let mut validation = self.validation.clone(); + validation.set_audience(&[audience]); + return jsonwebtoken::decode::(token, &self.decoding_key, &validation) + .map(|data| data.claims); + } + pub fn encode(&self, claims: &T) -> Result { return jsonwebtoken::encode::(&self.header, claims, &self.encoding_key); } diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index 8419ae39f..cf3f1afa1 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -24,6 +24,7 @@ mod encryption; mod extract; mod init_error; mod listing; +mod mcp; mod migrations; mod scheduler; mod schema_metadata; diff --git a/crates/core/src/mcp.rs b/crates/core/src/mcp.rs new file mode 100644 index 000000000..850885228 --- /dev/null +++ b/crates/core/src/mcp.rs @@ -0,0 +1,908 @@ +use std::sync::Arc; + +use axum::Router; +use axum::body::Body; +use axum::extract::{Form, Json as AxumJson, Path, Query, Request, State}; +use axum::http::{HeaderValue, Method, StatusCode, header}; +use axum::middleware::{self, Next}; +use axum::response::{IntoResponse, Redirect, Response}; +use axum::routing::{get, post}; +use http_body_util::BodyExt; +use rmcp::handler::server::{router::tool::ToolRouter, wrapper::Parameters}; +use rmcp::model::{ErrorData as McpError, Implementation, ServerCapabilities, ServerInfo}; +use rmcp::transport::{ + StreamableHttpServerConfig, + streamable_http_server::{session::local::LocalSessionManager, tower::StreamableHttpService}, +}; +use rmcp::{Json, ServerHandler, schemars, tool, tool_handler, tool_router}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use tower::ServiceExt; + +use crate::admin; +use crate::app_state::AppState; +use crate::auth::util::is_admin; +use crate::auth::{AuthError, AuthTokenClaims, User}; + +const MCP_SCOPE: &str = "mcp"; +const MCP_PATH: &str = "/mcp"; + +#[derive(Debug, Deserialize, schemars::JsonSchema)] +struct AdminRequest { + /// HTTP method accepted by the TrailBase admin API. + method: String, + /// Admin API path relative to /api/_admin, including an optional query string. + path: String, + /// Optional JSON request body. + #[serde(default)] + body: Option, +} + +#[derive(Debug, Deserialize, schemars::JsonSchema)] +struct SqlRequest { + /// One or more SQLite statements. Schema-changing statements refresh TrailBase metadata. + query: String, + /// Optional configured attached database names. + #[serde(default)] + attached_databases: Option>, +} + +#[derive(Debug, Deserialize, schemars::JsonSchema)] +struct ConfigUpdateRequest { + /// Complete TrailBase config in protobuf text format, as returned by get_config. + config: String, +} + +#[derive(Clone)] +struct TrailBaseMcp { + state: AppState, + #[allow(dead_code)] + tool_router: ToolRouter, +} + +impl TrailBaseMcp { + fn new(state: AppState) -> Self { + Self { + state, + tool_router: Self::tool_router(), + } + } + + async fn dispatch_admin(&self, request: AdminRequest) -> Result, McpError> { + let method = Method::from_bytes(request.method.as_bytes()) + .map_err(|_| McpError::invalid_params("invalid HTTP method", None))?; + let path = normalize_admin_path(&request.path)?; + let body = request + .body + .map(|value| serde_json::to_vec(&value)) + .transpose() + .map_err(|err| McpError::invalid_params(err.to_string(), None))? + .unwrap_or_default(); + + let request = Request::builder() + .method(method) + .uri(path) + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(body)) + .map_err(|err| McpError::internal_error(err.to_string(), None))?; + + let (admin_router, _) = admin::router().split_for_parts(); + let response = admin_router + .with_state(self.state.clone()) + .oneshot(request) + .await + .map_err(|never| match never {})?; + let status = response.status(); + let bytes = response + .into_body() + .collect() + .await + .map_err(|err| McpError::internal_error(err.to_string(), None))? + .to_bytes(); + + let response_body = if bytes.is_empty() { + Value::Null + } else { + serde_json::from_slice(&bytes) + .unwrap_or_else(|_| Value::String(String::from_utf8_lossy(&bytes).into_owned())) + }; + + if !status.is_success() { + return Err(McpError::internal_error( + format!("TrailBase admin API returned {status}: {response_body}"), + Some(json!({ "status": status.as_u16(), "body": response_body })), + )); + } + + Ok(Json(json!({ + "status": status.as_u16(), + "body": response_body + }))) + } +} + +#[tool_router] +impl TrailBaseMcp { + #[tool( + description = "Call a TrailBase admin API in-process. Paths are relative to /api/_admin. This exposes the same table, index, row, config, schema, query, user, log, backup, job, and WASM operations as the admin dashboard." + )] + async fn call_admin_api( + &self, + Parameters(request): Parameters, + ) -> Result, McpError> { + self.dispatch_admin(request).await + } + + #[tool(description = "List TrailBase tables, views, columns, indexes, triggers, and metadata.")] + async fn list_tables(&self) -> Result, McpError> { + self + .dispatch_admin(AdminRequest { + method: "GET".to_string(), + path: "tables".to_string(), + body: None, + }) + .await + } + + #[tool( + description = "Execute SQL using TrailBase's admin query handler. Supports reads and writes; schema changes refresh cached metadata." + )] + async fn execute_sql( + &self, + Parameters(request): Parameters, + ) -> Result, McpError> { + self + .dispatch_admin(AdminRequest { + method: "POST".to_string(), + path: "query".to_string(), + body: Some(json!({ + "query": request.query, + "attached_databases": request.attached_databases, + })), + }) + .await + } + + #[tool( + description = "Get the complete TrailBase configuration as protobuf text. Secret values are redacted." + )] + fn get_config(&self) -> Result { + let (config, _) = crate::config::redact_secrets(&self.state.get_config()) + .map_err(|err| McpError::internal_error(err.to_string(), None))?; + config + .to_text() + .map_err(|err| McpError::internal_error(err.to_string(), None)) + } + + #[tool( + description = "Validate and replace the TrailBase configuration using protobuf text from get_config. Existing secret values are preserved." + )] + async fn update_config( + &self, + Parameters(request): Parameters, + ) -> Result { + if self.state.demo_mode() { + return Err(McpError::invalid_request( + "config updates are disabled in demo mode", + None, + )); + } + + let config = crate::config::proto::Config::from_text(&request.config) + .map_err(|err| McpError::invalid_params(err.to_string(), None))?; + let current = self.state.get_config(); + let hash = crate::config::proto::hash_config(¤t); + let (_, secrets) = crate::config::redact_secrets(¤t) + .map_err(|err| McpError::internal_error(err.to_string(), None))?; + let config = + crate::config::merge_vault_and_env(config, crate::config::proto::Vault { secrets }) + .map_err(|err| McpError::invalid_params(err.to_string(), None))?; + self + .state + .validate_and_update_config(config, Some(hash)) + .await + .map_err(|err| McpError::invalid_params(err.to_string(), None))?; + Ok("Config updated".to_string()) + } +} + +#[tool_handler] +impl ServerHandler for TrailBaseMcp { + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info( + Implementation::new("trailbase", env!("CARGO_PKG_VERSION")) + .with_title("TrailBase MCP") + .with_description("Native administrative MCP server for TrailBase") + .with_website_url("https://trailbase.io"), + ) + .with_instructions( + "TrailBase's native administrative MCP server. Use call_admin_api to perform the same operations as the admin dashboard. Destructive operations modify the active TrailBase depot.", + ) + } +} + +pub(crate) fn router(state: &AppState) -> Router { + let state_for_service = state.clone(); + let service: StreamableHttpService = + StreamableHttpService::new( + move || Ok(TrailBaseMcp::new(state_for_service.clone())), + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default() + .with_json_response(true) + // TrailBase supports operator-configured reverse proxies. Admin authentication below is + // the security boundary, so the MCP transport must accept the proxy's Host header. + .disable_allowed_hosts(), + ); + + let protected_mcp = Router::new() + .nest_service("/mcp", service) + // Keep MCP authentication scoped to the MCP route. A normal layer also wraps this + // router's fallback and would intercept auth UI routes after this router is merged. + .route_layer(middleware::from_fn_with_state( + state.clone(), + assert_mcp_access, + )); + + Router::new() + .merge(protected_mcp) + .route( + "/.well-known/oauth-protected-resource", + get(protected_resource_metadata), + ) + .route( + "/.well-known/oauth-protected-resource/mcp", + get(protected_resource_metadata), + ) + .route( + "/.well-known/oauth-authorization-server", + get(authorization_server_metadata), + ) + .route("/_/mcp/authorize", get(authorize)) + .route("/_/mcp/callback/{flow}", get(authorization_callback)) + .route("/_/mcp/register", post(register_client)) + .route("/_/mcp/token", post(oauth_token)) +} + +async fn assert_mcp_access( + State(state): State, + request: Request, + next: Next, +) -> Response { + let user_id = request + .headers() + .get(header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")) + .and_then(|token| mcp_user_id(&state, token)); + let authorized = match user_id { + Some(user_id) => is_admin(&state, &user_id).await, + None => false, + }; + if authorized { + return next.run(request).await; + } + + let metadata_url = external_url(&state, "/.well-known/oauth-protected-resource"); + let mut response = AuthError::Unauthorized.into_response(); + if let Ok(value) = HeaderValue::from_str(&format!( + "Bearer resource_metadata=\"{metadata_url}\", scope=\"{MCP_SCOPE}\"" + )) { + response + .headers_mut() + .insert(header::WWW_AUTHENTICATE, value); + } + response +} + +#[derive(Clone, Serialize, Deserialize)] +struct McpAccessTokenClaims { + #[serde(flatten)] + auth: AuthTokenClaims, + aud: String, + scope: String, +} + +#[derive(Clone, Deserialize)] +struct CompatibilityAccessTokenClaims { + #[serde(flatten)] + auth: AuthTokenClaims, + #[serde(default)] + aud: Option, +} + +fn mcp_user_id(state: &AppState, token: &str) -> Option { + let resource = external_url(state, MCP_PATH); + if let Ok(claims) = state + .jwt() + .decode_with_audience::(token, &resource) + { + if !claims.scope.split(' ').any(|scope| scope == MCP_SCOPE) { + return None; + } + return crate::util::b64_to_uuid(&claims.auth.sub).ok(); + } + + // Compatibility mode for callers that explicitly supply a normal TrailBase admin token. + let claims = state + .jwt() + .decode::(token) + .ok()?; + if claims.aud.is_some() { + return None; + } + crate::util::b64_to_uuid(&claims.auth.sub).ok() +} + +#[derive(Serialize)] +struct ProtectedResourceMetadata { + resource: String, + authorization_servers: Vec, + scopes_supported: Vec<&'static str>, + bearer_methods_supported: Vec<&'static str>, +} + +async fn protected_resource_metadata( + State(state): State, +) -> AxumJson { + let issuer = external_url(&state, ""); + AxumJson(ProtectedResourceMetadata { + resource: external_url(&state, MCP_PATH), + authorization_servers: vec![issuer], + scopes_supported: vec![MCP_SCOPE], + bearer_methods_supported: vec!["header"], + }) +} + +#[derive(Serialize)] +struct AuthorizationServerMetadata { + issuer: String, + authorization_endpoint: String, + token_endpoint: String, + registration_endpoint: String, + response_types_supported: Vec<&'static str>, + grant_types_supported: Vec<&'static str>, + code_challenge_methods_supported: Vec<&'static str>, + token_endpoint_auth_methods_supported: Vec<&'static str>, + scopes_supported: Vec<&'static str>, +} + +async fn authorization_server_metadata( + State(state): State, +) -> AxumJson { + AxumJson(AuthorizationServerMetadata { + issuer: external_url(&state, ""), + authorization_endpoint: external_url(&state, "/_/mcp/authorize"), + token_endpoint: external_url(&state, "/_/mcp/token"), + registration_endpoint: external_url(&state, "/_/mcp/register"), + response_types_supported: vec!["code"], + grant_types_supported: vec!["authorization_code", "refresh_token"], + code_challenge_methods_supported: vec!["S256"], + token_endpoint_auth_methods_supported: vec!["none"], + scopes_supported: vec![MCP_SCOPE], + }) +} + +#[derive(Debug, Deserialize, Serialize)] +struct ClientRegistration { + #[serde(default)] + redirect_uris: Vec, + #[serde(default)] + client_name: Option, + #[serde(default)] + scope: Option, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +struct ClientClaims { + exp: i64, + iat: i64, + redirect_uris: Vec, + client_name: Option, +} + +#[derive(Serialize)] +struct ClientRegistrationResponse { + client_id: String, + client_id_issued_at: i64, + redirect_uris: Vec, + client_name: Option, + token_endpoint_auth_method: &'static str, + grant_types: Vec<&'static str>, + response_types: Vec<&'static str>, + scope: &'static str, +} + +async fn register_client( + State(state): State, + AxumJson(request): AxumJson, +) -> Result, OAuthError> { + if request.redirect_uris.is_empty() + || request + .redirect_uris + .iter() + .any(|uri| !valid_client_redirect(uri)) + || request + .scope + .as_deref() + .is_some_and(|scope| !scope.split(' ').any(|value| value == MCP_SCOPE)) + { + return Err(OAuthError::invalid_client_metadata("invalid redirect_uris")); + } + + let now = chrono::Utc::now().timestamp(); + let claims = ClientClaims { + iat: now, + exp: now + chrono::Duration::days(30).num_seconds(), + redirect_uris: request.redirect_uris.clone(), + client_name: request.client_name.clone(), + }; + let client_id = state + .jwt() + .encode(&claims) + .map_err(|err| OAuthError::server(err.to_string()))?; + + Ok(AxumJson(ClientRegistrationResponse { + client_id, + client_id_issued_at: now, + redirect_uris: request.redirect_uris, + client_name: request.client_name, + token_endpoint_auth_method: "none", + grant_types: vec!["authorization_code", "refresh_token"], + response_types: vec!["code"], + scope: MCP_SCOPE, + })) +} + +#[derive(Deserialize)] +struct AuthorizeQuery { + client_id: String, + redirect_uri: String, + response_type: String, + code_challenge: String, + code_challenge_method: String, + state: Option, + scope: Option, + resource: Option, +} + +#[derive(Clone, Serialize, Deserialize)] +struct FlowClaims { + exp: i64, + redirect_uri: String, + client_id: String, + state: Option, +} + +async fn authorize( + State(state): State, + user: Option, + Query(query): Query, +) -> Result { + let client: ClientClaims = state + .jwt() + .decode(&query.client_id) + .map_err(|_| OAuthError::invalid_request("unknown or expired client_id"))?; + if query.response_type != "code" + || query.code_challenge_method != "S256" + || !client.redirect_uris.contains(&query.redirect_uri) + || query + .scope + .as_deref() + .is_some_and(|scope| !scope.split(' ').any(|s| s == MCP_SCOPE)) + || query + .resource + .as_deref() + .is_some_and(|resource| resource != external_url(&state, MCP_PATH)) + { + return Err(OAuthError::invalid_request("invalid authorization request")); + } + + let flow = state + .jwt() + .encode(&FlowClaims { + exp: chrono::Utc::now().timestamp() + chrono::Duration::minutes(10).num_seconds(), + redirect_uri: query.redirect_uri, + client_id: query.client_id, + state: query.state, + }) + .map_err(|err| OAuthError::server(err.to_string()))?; + let callback = format!("/_/mcp/callback/{flow}"); + let login_query = url::form_urlencoded::Serializer::new(String::new()) + .append_pair("redirect_uri", &callback) + .append_pair("response_type", "code") + .append_pair("pkce_code_challenge", &query.code_challenge) + .finish(); + + if let Some(user) = user { + if !is_admin(&state, &user.uuid).await { + return Err(OAuthError::access_denied( + "MCP access requires an administrator", + )); + } + let db_user = crate::auth::util::user_by_id(&state, &user.uuid) + .await + .map_err(OAuthError::from_auth)?; + return crate::auth::api::login::build_authorization_code_flow_and_pkce_response( + &state, + &db_user, + callback, + query.code_challenge, + ) + .await + .map_err(OAuthError::from_auth); + } + + Ok(Redirect::to(&format!("/_/auth/login?{login_query}")).into_response()) +} + +#[derive(Deserialize)] +struct CallbackQuery { + code: String, +} + +async fn authorization_callback( + State(state): State, + Path(flow): Path, + Query(query): Query, +) -> Result { + let flow: FlowClaims = state + .jwt() + .decode(&flow) + .map_err(|_| OAuthError::invalid_request("unknown or expired authorization flow"))?; + let mut redirect = url::Url::parse(&flow.redirect_uri) + .map_err(|_| OAuthError::invalid_request("invalid redirect_uri"))?; + redirect.query_pairs_mut().append_pair("code", &query.code); + if let Some(state) = flow.state { + redirect.query_pairs_mut().append_pair("state", &state); + } + Ok(Redirect::to(redirect.as_str())) +} + +#[derive(Deserialize)] +struct TokenRequest { + grant_type: String, + code: Option, + code_verifier: Option, + refresh_token: Option, + client_id: Option, + redirect_uri: Option, + resource: Option, +} + +#[derive(Serialize)] +struct TokenResponse { + access_token: String, + token_type: &'static str, + expires_in: i64, + refresh_token: Option, + scope: &'static str, +} + +async fn oauth_token( + State(state): State, + Form(request): Form, +) -> Result, OAuthError> { + if request + .resource + .as_deref() + .is_some_and(|resource| resource != external_url(&state, MCP_PATH)) + { + return Err(OAuthError::invalid_grant("invalid resource")); + } + + let client_id = request + .client_id + .as_deref() + .ok_or_else(|| OAuthError::invalid_grant("missing client_id"))?; + let client: ClientClaims = state + .jwt() + .decode(client_id) + .map_err(|_| OAuthError::invalid_grant("unknown or expired client_id"))?; + + let (access_token, refresh_token) = match request.grant_type.as_str() { + "authorization_code" => { + let redirect_uri = request + .redirect_uri + .as_deref() + .ok_or_else(|| OAuthError::invalid_grant("missing redirect_uri"))?; + if !client.redirect_uris.iter().any(|uri| uri == redirect_uri) { + return Err(OAuthError::invalid_grant( + "redirect_uri does not match client", + )); + } + let code = request + .code + .ok_or_else(|| OAuthError::invalid_grant("missing code"))?; + let verifier = request + .code_verifier + .ok_or_else(|| OAuthError::invalid_grant("missing code_verifier"))?; + let AxumJson(tokens) = crate::auth::api::token::auth_code_to_token_handler( + State(state.clone()), + AxumJson(crate::auth::api::token::AuthCodeToTokenRequest { + authorization_code: Some(code), + pkce_code_verifier: Some(verifier), + }), + ) + .await + .map_err(OAuthError::from_auth)?; + (tokens.auth_token, Some(tokens.refresh_token)) + } + "refresh_token" => { + let refresh_token = request + .refresh_token + .ok_or_else(|| OAuthError::invalid_grant("missing refresh_token"))?; + let AxumJson(tokens) = crate::auth::api::refresh::refresh_handler( + State(state.clone()), + AxumJson(crate::auth::api::refresh::RefreshRequest { refresh_token }), + ) + .await + .map_err(OAuthError::from_auth)?; + (tokens.auth_token, None) + } + _ => return Err(OAuthError::invalid_grant("unsupported grant_type")), + }; + let claims = AuthTokenClaims::from_auth_token(state.jwt(), &access_token) + .map_err(|_| OAuthError::server("failed to decode issued access token"))?; + let expires_in = (claims.exp - chrono::Utc::now().timestamp()).max(0); + let access_token = state + .jwt() + .encode(&McpAccessTokenClaims { + auth: claims, + aud: external_url(&state, MCP_PATH), + scope: MCP_SCOPE.to_string(), + }) + .map_err(|err| OAuthError::server(err.to_string()))?; + + Ok(AxumJson(TokenResponse { + access_token, + token_type: "Bearer", + expires_in, + refresh_token, + scope: MCP_SCOPE, + })) +} + +#[derive(Debug)] +struct OAuthError { + status: StatusCode, + code: &'static str, + description: String, +} + +impl OAuthError { + fn invalid_request(description: impl Into) -> Self { + Self { + status: StatusCode::BAD_REQUEST, + code: "invalid_request", + description: description.into(), + } + } + + fn invalid_client_metadata(description: impl Into) -> Self { + Self { + status: StatusCode::BAD_REQUEST, + code: "invalid_client_metadata", + description: description.into(), + } + } + + fn invalid_grant(description: impl Into) -> Self { + Self { + status: StatusCode::BAD_REQUEST, + code: "invalid_grant", + description: description.into(), + } + } + + fn access_denied(description: impl Into) -> Self { + Self { + status: StatusCode::FORBIDDEN, + code: "access_denied", + description: description.into(), + } + } + + fn server(description: impl Into) -> Self { + Self { + status: StatusCode::INTERNAL_SERVER_ERROR, + code: "server_error", + description: description.into(), + } + } + + fn from_auth(error: AuthError) -> Self { + Self { + status: StatusCode::BAD_REQUEST, + code: "invalid_grant", + description: error.to_string(), + } + } +} + +impl IntoResponse for OAuthError { + fn into_response(self) -> Response { + ( + self.status, + AxumJson(json!({ "error": self.code, "error_description": self.description })), + ) + .into_response() + } +} + +fn valid_client_redirect(uri: &str) -> bool { + let Ok(uri) = url::Url::parse(uri) else { + return false; + }; + uri.scheme() == "https" + || (uri.scheme() == "http" && matches!(uri.host_str(), Some("localhost" | "127.0.0.1" | "::1"))) +} + +fn external_url(state: &AppState, path: &str) -> String { + let mut base = state + .site_url() + .as_ref() + .clone() + .unwrap_or_else(|| url::Url::parse("http://localhost:4000").expect("constant URL")); + base.set_path(path); + base.set_query(None); + base.set_fragment(None); + base.to_string().trim_end_matches('/').to_string() +} + +fn normalize_admin_path(path: &str) -> Result { + let path = path.trim(); + if path.is_empty() || path.contains("://") || path.starts_with("//") { + return Err(McpError::invalid_params("invalid admin API path", None)); + } + + let path = path + .strip_prefix("/api/_admin") + .unwrap_or(path) + .trim_start_matches('/'); + Ok(format!("/{path}")) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::app_state::test_state; + + #[test] + fn normalizes_admin_paths() { + assert_eq!(normalize_admin_path("tables").unwrap(), "/tables"); + assert_eq!( + normalize_admin_path("/api/_admin/logs/list?limit=5").unwrap(), + "/logs/list?limit=5" + ); + assert!(normalize_admin_path("https://example.com").is_err()); + } + + #[tokio::test] + async fn registers_client_and_builds_pkce_login_redirect() { + let state = test_state(None).await.unwrap(); + let callback = "http://127.0.0.1:3334/oauth/callback".to_string(); + let AxumJson(registration) = register_client( + State(state.clone()), + AxumJson(ClientRegistration { + redirect_uris: vec![callback.clone()], + client_name: Some("test client".to_string()), + scope: Some(MCP_SCOPE.to_string()), + }), + ) + .await + .unwrap(); + + let redirect = authorize( + State(state), + None, + Query(AuthorizeQuery { + client_id: registration.client_id, + redirect_uri: callback, + response_type: "code".to_string(), + code_challenge: "ZmFrZS1jaGFsbGVuZ2U".to_string(), + code_challenge_method: "S256".to_string(), + state: Some("client-state".to_string()), + scope: Some(MCP_SCOPE.to_string()), + resource: None, + }), + ) + .await + .unwrap() + .into_response(); + + let location = redirect + .headers() + .get(header::LOCATION) + .unwrap() + .to_str() + .unwrap(); + assert!(location.starts_with("/_/auth/login?")); + assert!(location.contains("response_type=code")); + assert!(location.contains("pkce_code_challenge=")); + } + + #[tokio::test] + async fn native_tools_use_live_admin_state() { + let state = test_state(None).await.unwrap(); + let server = TrailBaseMcp::new(state); + let tool_names: Vec<_> = server + .tool_router + .list_all() + .into_iter() + .map(|tool| tool.name.to_string()) + .collect(); + assert_eq!( + tool_names, + [ + "call_admin_api", + "execute_sql", + "get_config", + "list_tables", + "update_config" + ] + ); + + let config = server.get_config().unwrap(); + assert!(config.contains("auth")); + assert_eq!( + server + .update_config(Parameters(ConfigUpdateRequest { config })) + .await + .unwrap(), + "Config updated" + ); + + server + .execute_sql(Parameters(SqlRequest { + query: "CREATE TABLE mcp_native_tool_test (id INTEGER PRIMARY KEY)".to_string(), + attached_databases: None, + })) + .await + .unwrap(); + let tables = server.list_tables().await.unwrap(); + assert!(tables.0.to_string().contains("mcp_native_tool_test")); + } + + #[tokio::test] + async fn mcp_tokens_are_bound_to_the_mcp_resource() { + let state = test_state(None).await.unwrap(); + let user_id = uuid::Uuid::new_v4(); + let now = chrono::Utc::now().timestamp(); + let claims = AuthTokenClaims { + sub: crate::util::uuid_to_b64(&user_id), + iat: now, + exp: now + 60, + r#type: 1, + admin: true, + mfa: false, + provider: 0, + email: Some("admin@localhost".to_string()), + username: Some("admin".to_string()), + csrf_token: "csrf".to_string(), + }; + + let scoped = state + .jwt() + .encode(&McpAccessTokenClaims { + auth: claims.clone(), + aud: external_url(&state, MCP_PATH), + scope: MCP_SCOPE.to_string(), + }) + .unwrap(); + assert_eq!(mcp_user_id(&state, &scoped), Some(user_id)); + + let wrong_audience = state + .jwt() + .encode(&McpAccessTokenClaims { + auth: claims.clone(), + aud: "https://other.example/mcp".to_string(), + scope: MCP_SCOPE.to_string(), + }) + .unwrap(); + assert_eq!(mcp_user_id(&state, &wrong_audience), None); + + let legacy = state.jwt().encode(&claims).unwrap(); + assert_eq!(mcp_user_id(&state, &legacy), Some(user_id)); + } +} diff --git a/crates/core/src/server/mod.rs b/crates/core/src/server/mod.rs index cabd98861..e3b3f66f1 100644 --- a/crates/core/src/server/mod.rs +++ b/crates/core/src/server/mod.rs @@ -76,6 +76,9 @@ pub struct ServerOptions { /// Custom axum router. pub custom_router: Option>, + + /// Expose the MCP endpoint on the admin server. + pub enable_mcp: bool, } pub struct Server { @@ -109,6 +112,7 @@ impl Server { tls_cert, tls_key, custom_router, + enable_mcp, } = opts; let version_info = trailbase_build::get_version_info!(); @@ -176,7 +180,7 @@ impl Server { None }; - let admin_router = Self::build_admin_router(&state); + let admin_router = Self::build_admin_router(&state, enable_mcp); let independent_admin_router = if let Some(admin_address) = admin_address && admin_address != address { @@ -370,8 +374,8 @@ impl Server { return Ok(()); } - pub(crate) fn build_admin_router(state: &AppState) -> OpenApiRouter { - return OpenApiRouter::new() + fn build_admin_router(state: &AppState, enable_mcp: bool) -> OpenApiRouter { + let mut router = OpenApiRouter::new() .nest( &format!("/{ADMIN_API_PATH}/"), admin::router().layer(middleware::from_fn_with_state( @@ -396,6 +400,13 @@ impl Server { }, ), ); + + if enable_mcp { + // TODO: Should be OpenApiTouter + router = router.merge(crate::mcp::router(state).into()); + } + + return router; } pub(crate) fn build_main_router( From bf8c4f2c180d1150b198e234131660a5ec563ccc Mon Sep 17 00:00:00 2001 From: Sebastian Jeltsch Date: Wed, 12 Aug 2026 22:24:59 +0200 Subject: [PATCH 2/4] First sweep of cleanups and questions. --- crates/core/src/mcp.rs | 242 +++++++++++++++++++++++------------------ 1 file changed, 134 insertions(+), 108 deletions(-) diff --git a/crates/core/src/mcp.rs b/crates/core/src/mcp.rs index 850885228..cf6683ce3 100644 --- a/crates/core/src/mcp.rs +++ b/crates/core/src/mcp.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use axum::Router; use axum::body::Body; use axum::extract::{Form, Json as AxumJson, Path, Query, Request, State}; -use axum::http::{HeaderValue, Method, StatusCode, header}; +use axum::http::{Method, StatusCode, header}; use axum::middleware::{self, Next}; use axum::response::{IntoResponse, Redirect, Response}; use axum::routing::{get, post}; @@ -69,42 +69,51 @@ impl TrailBaseMcp { } async fn dispatch_admin(&self, request: AdminRequest) -> Result, McpError> { - let method = Method::from_bytes(request.method.as_bytes()) - .map_err(|_| McpError::invalid_params("invalid HTTP method", None))?; - let path = normalize_admin_path(&request.path)?; - let body = request - .body - .map(|value| serde_json::to_vec(&value)) - .transpose() - .map_err(|err| McpError::invalid_params(err.to_string(), None))? - .unwrap_or_default(); - - let request = Request::builder() - .method(method) - .uri(path) - .header(header::CONTENT_TYPE, "application/json") - .body(Body::from(body)) - .map_err(|err| McpError::internal_error(err.to_string(), None))?; + let request = { + let method = Method::from_bytes(request.method.as_bytes()) + .map_err(|_| McpError::invalid_params("invalid HTTP method", None))?; + + let path = normalize_admin_path(&request.path)?; + let body = request + .body + .map(|value| serde_json::to_vec(&value)) + .transpose() + .map_err(|err| McpError::invalid_params(err.to_string(), None))? + .unwrap_or_default(); + + Request::builder() + .method(method) + .uri(path) + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(body)) + .map_err(|err| McpError::internal_error(err.to_string(), None))? + }; + // QUESTION: Rather than building a dedicated router, why not just call ourselves? Will make + // sure non-admin endpoints are available, logging works, ... let (admin_router, _) = admin::router().split_for_parts(); let response = admin_router .with_state(self.state.clone()) .oneshot(request) .await .map_err(|never| match never {})?; + let status = response.status(); - let bytes = response - .into_body() - .collect() - .await - .map_err(|err| McpError::internal_error(err.to_string(), None))? - .to_bytes(); - - let response_body = if bytes.is_empty() { - Value::Null - } else { - serde_json::from_slice(&bytes) - .unwrap_or_else(|_| Value::String(String::from_utf8_lossy(&bytes).into_owned())) + let response_body = { + let bytes = response + .into_body() + .collect() + .await + .map_err(|err| McpError::internal_error(err.to_string(), None))? + .to_bytes(); + + // TODO: Should we rely on content-type instead? + if bytes.is_empty() { + Value::Null + } else { + serde_json::from_slice(&bytes) + .unwrap_or_else(|_| Value::String(String::from_utf8_lossy(&bytes).into_owned())) + } }; if !status.is_success() { @@ -223,10 +232,12 @@ impl ServerHandler for TrailBaseMcp { } pub(crate) fn router(state: &AppState) -> Router { - let state_for_service = state.clone(); let service: StreamableHttpService = StreamableHttpService::new( - move || Ok(TrailBaseMcp::new(state_for_service.clone())), + { + let state = state.clone(); + move || Ok(TrailBaseMcp::new(state.clone())) + }, Arc::new(LocalSessionManager::default()), StreamableHttpServerConfig::default() .with_json_response(true) @@ -248,20 +259,24 @@ pub(crate) fn router(state: &AppState) -> Router { .merge(protected_mcp) .route( "/.well-known/oauth-protected-resource", - get(protected_resource_metadata), + get(protected_resource_metadata_handler), ) .route( "/.well-known/oauth-protected-resource/mcp", - get(protected_resource_metadata), + get(protected_resource_metadata_handler), ) .route( "/.well-known/oauth-authorization-server", - get(authorization_server_metadata), + get(authorization_server_metadata_handler), ) - .route("/_/mcp/authorize", get(authorize)) - .route("/_/mcp/callback/{flow}", get(authorization_callback)) - .route("/_/mcp/register", post(register_client)) - .route("/_/mcp/token", post(oauth_token)) + // TODO: Should the mcp be under /api/_mcp? + .route("/_/mcp/authorize", get(authorize_handler)) + .route( + "/_/mcp/callback/{flow}", + get(authorization_callback_handler), + ) + .route("/_/mcp/register", post(register_client_handler)) + .route("/_/mcp/token", post(oauth_token_handler)) } async fn assert_mcp_access( @@ -275,24 +290,24 @@ async fn assert_mcp_access( .and_then(|value| value.to_str().ok()) .and_then(|value| value.strip_prefix("Bearer ")) .and_then(|token| mcp_user_id(&state, token)); - let authorized = match user_id { - Some(user_id) => is_admin(&state, &user_id).await, - None => false, - }; - if authorized { + + if let Some(user_id) = user_id + && is_admin(&state, &user_id).await + { return next.run(request).await; } - let metadata_url = external_url(&state, "/.well-known/oauth-protected-resource"); - let mut response = AuthError::Unauthorized.into_response(); - if let Ok(value) = HeaderValue::from_str(&format!( - "Bearer resource_metadata=\"{metadata_url}\", scope=\"{MCP_SCOPE}\"" - )) { - response - .headers_mut() - .insert(header::WWW_AUTHENTICATE, value); - } - response + return Response::builder() + .status(StatusCode::UNAUTHORIZED) + .header( + header::WWW_AUTHENTICATE, + format!( + r#"Bearer resource_metadata="{metadata}", scope="{MCP_SCOPE}""#, + metadata = external_url(&state, "/.well-known/oauth-protected-resource"), + ), + ) + .body(axum::body::Body::empty()) + .unwrap_or_default(); } #[derive(Clone, Serialize, Deserialize)] @@ -331,6 +346,7 @@ fn mcp_user_id(state: &AppState, token: &str) -> Option { if claims.aud.is_some() { return None; } + crate::util::b64_to_uuid(&claims.auth.sub).ok() } @@ -342,10 +358,14 @@ struct ProtectedResourceMetadata { bearer_methods_supported: Vec<&'static str>, } -async fn protected_resource_metadata( +async fn protected_resource_metadata_handler( State(state): State, ) -> AxumJson { let issuer = external_url(&state, ""); + + // FIXME: Remove + println!("issuer: {issuer}"); + AxumJson(ProtectedResourceMetadata { resource: external_url(&state, MCP_PATH), authorization_servers: vec![issuer], @@ -367,10 +387,10 @@ struct AuthorizationServerMetadata { scopes_supported: Vec<&'static str>, } -async fn authorization_server_metadata( +async fn authorization_server_metadata_handler( State(state): State, ) -> AxumJson { - AxumJson(AuthorizationServerMetadata { + return AxumJson(AuthorizationServerMetadata { issuer: external_url(&state, ""), authorization_endpoint: external_url(&state, "/_/mcp/authorize"), token_endpoint: external_url(&state, "/_/mcp/token"), @@ -380,7 +400,7 @@ async fn authorization_server_metadata( code_challenge_methods_supported: vec!["S256"], token_endpoint_auth_methods_supported: vec!["none"], scopes_supported: vec![MCP_SCOPE], - }) + }); } #[derive(Debug, Deserialize, Serialize)] @@ -413,7 +433,7 @@ struct ClientRegistrationResponse { scope: &'static str, } -async fn register_client( +async fn register_client_handler( State(state): State, AxumJson(request): AxumJson, ) -> Result, OAuthError> { @@ -431,18 +451,17 @@ async fn register_client( } let now = chrono::Utc::now().timestamp(); - let claims = ClientClaims { - iat: now, - exp: now + chrono::Duration::days(30).num_seconds(), - redirect_uris: request.redirect_uris.clone(), - client_name: request.client_name.clone(), - }; let client_id = state .jwt() - .encode(&claims) + .encode(&ClientClaims { + iat: now, + exp: now + chrono::Duration::days(30).num_seconds(), + redirect_uris: request.redirect_uris.clone(), + client_name: request.client_name.clone(), + }) .map_err(|err| OAuthError::server(err.to_string()))?; - Ok(AxumJson(ClientRegistrationResponse { + return Ok(AxumJson(ClientRegistrationResponse { client_id, client_id_issued_at: now, redirect_uris: request.redirect_uris, @@ -451,7 +470,7 @@ async fn register_client( grant_types: vec!["authorization_code", "refresh_token"], response_types: vec!["code"], scope: MCP_SCOPE, - })) + })); } #[derive(Deserialize)] @@ -474,7 +493,7 @@ struct FlowClaims { state: Option, } -async fn authorize( +async fn authorize_handler( State(state): State, user: Option, Query(query): Query, @@ -483,6 +502,7 @@ async fn authorize( .jwt() .decode(&query.client_id) .map_err(|_| OAuthError::invalid_request("unknown or expired client_id"))?; + if query.response_type != "code" || query.code_challenge_method != "S256" || !client.redirect_uris.contains(&query.redirect_uri) @@ -498,42 +518,47 @@ async fn authorize( return Err(OAuthError::invalid_request("invalid authorization request")); } - let flow = state - .jwt() - .encode(&FlowClaims { - exp: chrono::Utc::now().timestamp() + chrono::Duration::minutes(10).num_seconds(), - redirect_uri: query.redirect_uri, - client_id: query.client_id, - state: query.state, - }) - .map_err(|err| OAuthError::server(err.to_string()))?; - let callback = format!("/_/mcp/callback/{flow}"); - let login_query = url::form_urlencoded::Serializer::new(String::new()) - .append_pair("redirect_uri", &callback) - .append_pair("response_type", "code") - .append_pair("pkce_code_challenge", &query.code_challenge) - .finish(); - - if let Some(user) = user { - if !is_admin(&state, &user.uuid).await { - return Err(OAuthError::access_denied( - "MCP access requires an administrator", - )); - } - let db_user = crate::auth::util::user_by_id(&state, &user.uuid) - .await - .map_err(OAuthError::from_auth)?; - return crate::auth::api::login::build_authorization_code_flow_and_pkce_response( - &state, - &db_user, - callback, - query.code_challenge, - ) - .await - .map_err(OAuthError::from_auth); - } + let callback = format!( + "/_/mcp/callback/{flow}", + flow = state + .jwt() + .encode(&FlowClaims { + exp: chrono::Utc::now().timestamp() + chrono::Duration::minutes(10).num_seconds(), + redirect_uri: query.redirect_uri, + client_id: query.client_id, + state: query.state, + }) + .map_err(|err| OAuthError::server(err.to_string()))? + ); + + return match user { + Some(user) if is_admin(&state, &user.uuid).await => { + let db_user = crate::auth::util::user_by_id(&state, &user.uuid) + .await + .map_err(OAuthError::from_auth)?; - Ok(Redirect::to(&format!("/_/auth/login?{login_query}")).into_response()) + crate::auth::api::login::build_authorization_code_flow_and_pkce_response( + &state, + &db_user, + callback, + query.code_challenge, + ) + .await + .map_err(OAuthError::from_auth) + } + Some(_user) => Err(OAuthError::access_denied("not an admin")), + None => Ok( + Redirect::to(&format!( + "/_/auth/login?{query}", + query = url::form_urlencoded::Serializer::new(String::new()) + .append_pair("redirect_uri", &callback) + .append_pair("response_type", "code") + .append_pair("pkce_code_challenge", &query.code_challenge) + .finish() + )) + .into_response(), + ), + }; } #[derive(Deserialize)] @@ -541,7 +566,7 @@ struct CallbackQuery { code: String, } -async fn authorization_callback( +async fn authorization_callback_handler( State(state): State, Path(flow): Path, Query(query): Query, @@ -579,7 +604,7 @@ struct TokenResponse { scope: &'static str, } -async fn oauth_token( +async fn oauth_token_handler( State(state): State, Form(request): Form, ) -> Result, OAuthError> { @@ -760,7 +785,8 @@ fn normalize_admin_path(path: &str) -> Result { .strip_prefix("/api/_admin") .unwrap_or(path) .trim_start_matches('/'); - Ok(format!("/{path}")) + + return Ok(format!("/{path}")); } #[cfg(test)] @@ -782,7 +808,7 @@ mod tests { async fn registers_client_and_builds_pkce_login_redirect() { let state = test_state(None).await.unwrap(); let callback = "http://127.0.0.1:3334/oauth/callback".to_string(); - let AxumJson(registration) = register_client( + let AxumJson(registration) = register_client_handler( State(state.clone()), AxumJson(ClientRegistration { redirect_uris: vec![callback.clone()], @@ -793,7 +819,7 @@ mod tests { .await .unwrap(); - let redirect = authorize( + let redirect = authorize_handler( State(state), None, Query(AuthorizeQuery { From 1e4f2e44b4b5cf375301b4df98351e6c1a83d7ca Mon Sep 17 00:00:00 2001 From: Sebastian Jeltsch Date: Fri, 14 Aug 2026 09:49:26 +0200 Subject: [PATCH 3/4] Hack: hard code localhost:4000 resource for local testing but pave the way for wiring up Host/RealIp instead. --- crates/core/src/extract/ip.rs | 64 +++++++++++++++++++++++++++++++++-- crates/core/src/mcp.rs | 47 +++++++++++++++++-------- 2 files changed, 95 insertions(+), 16 deletions(-) diff --git a/crates/core/src/extract/ip.rs b/crates/core/src/extract/ip.rs index 0354352c1..2cc576f60 100644 --- a/crates/core/src/extract/ip.rs +++ b/crates/core/src/extract/ip.rs @@ -1,5 +1,7 @@ -use axum::extract::ConnectInfo; -use axum::http::Request; +use axum::extract::{ConnectInfo, FromRequestParts}; +use axum::http::uri; +use axum::http::{Request, request::Parts}; +use std::convert::Infallible; use std::net::IpAddr; use tower_governor::GovernorError; use tower_governor::key_extractor::KeyExtractor; @@ -23,6 +25,27 @@ pub fn extract_ip(req: &Request) -> Option { }); } +// TODO: Should replace above function. +pub fn extract_ip2(parts: &Parts) -> Option { + let headers = &parts.headers; + + // NOTE: This code is mimicking axum_client_ip's pre v1 `InsecureClientIp::from`: + return client_ip::rightmost_x_forwarded_for(headers) + .or_else(|_| client_ip::x_real_ip(headers)) + .or_else(|_| client_ip::fly_client_ip(headers)) + .or_else(|_| client_ip::true_client_ip(headers)) + .or_else(|_| client_ip::cf_connecting_ip(headers)) + .or_else(|_| client_ip::cloudfront_viewer_address(headers)) + .ok() + .or_else(|| { + parts + .extensions + .get::>() + .map(|ConnectInfo(addr)| addr.ip()) + }); +} + +/// Key extractor for the Governor. #[derive(Debug, Clone)] pub struct RealIpKeyExtractor; @@ -42,6 +65,43 @@ impl KeyExtractor for RealIpKeyExtractor { // } } +// Extractor for handlers. +#[derive(Debug, Clone, Default)] +pub struct RealIp(pub Option); + +impl FromRequestParts for RealIp +where + S: Send + Sync, +{ + type Rejection = Infallible; + + async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result { + return Ok(Self(extract_ip2(parts))); + } +} + +// Extractor for handlers. +#[derive(Debug, Clone, Default)] +pub struct Host(pub Option); + +impl FromRequestParts for Host +where + S: Send + Sync, +{ + type Rejection = Infallible; + + async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result { + let header = parts + .headers + .get("X-Forwarded-Host") + .or_else(|| parts.headers.get("host")); + + return Ok(Self( + header.and_then(|h| uri::Authority::try_from(h.as_bytes()).ok()), + )); + } +} + #[allow(unused)] pub fn ipv6_privacy_mask(ip: IpAddr) -> IpAddr { return match ip { diff --git a/crates/core/src/mcp.rs b/crates/core/src/mcp.rs index cf6683ce3..4a42f9529 100644 --- a/crates/core/src/mcp.rs +++ b/crates/core/src/mcp.rs @@ -23,6 +23,7 @@ use crate::admin; use crate::app_state::AppState; use crate::auth::util::is_admin; use crate::auth::{AuthError, AuthTokenClaims, User}; +use crate::extract::ip::{Host, RealIp}; const MCP_SCOPE: &str = "mcp"; const MCP_PATH: &str = "/mcp"; @@ -261,6 +262,7 @@ pub(crate) fn router(state: &AppState) -> Router { "/.well-known/oauth-protected-resource", get(protected_resource_metadata_handler), ) + // QUESTION: Is this "/mcp" needed atop the root one above? .route( "/.well-known/oauth-protected-resource/mcp", get(protected_resource_metadata_handler), @@ -303,7 +305,7 @@ async fn assert_mcp_access( header::WWW_AUTHENTICATE, format!( r#"Bearer resource_metadata="{metadata}", scope="{MCP_SCOPE}""#, - metadata = external_url(&state, "/.well-known/oauth-protected-resource"), + metadata = external_url2(&state, "/.well-known/oauth-protected-resource"), ), ) .body(axum::body::Body::empty()) @@ -327,10 +329,9 @@ struct CompatibilityAccessTokenClaims { } fn mcp_user_id(state: &AppState, token: &str) -> Option { - let resource = external_url(state, MCP_PATH); if let Ok(claims) = state .jwt() - .decode_with_audience::(token, &resource) + .decode_with_audience::(token, &external_url2(state, MCP_PATH)) { if !claims.scope.split(' ').any(|scope| scope == MCP_SCOPE) { return None; @@ -360,14 +361,17 @@ struct ProtectedResourceMetadata { async fn protected_resource_metadata_handler( State(state): State, + RealIp(ip): RealIp, + Host(authority): Host, ) -> AxumJson { - let issuer = external_url(&state, ""); + let issuer = external_url2(&state, ""); + let resource = external_url2(&state, MCP_PATH); // FIXME: Remove - println!("issuer: {issuer}"); + log::debug!("ip: {ip:?}, authority: {authority:?}, issuer: {issuer}, resource: {resource}"); AxumJson(ProtectedResourceMetadata { - resource: external_url(&state, MCP_PATH), + resource, authorization_servers: vec![issuer], scopes_supported: vec![MCP_SCOPE], bearer_methods_supported: vec!["header"], @@ -391,10 +395,10 @@ async fn authorization_server_metadata_handler( State(state): State, ) -> AxumJson { return AxumJson(AuthorizationServerMetadata { - issuer: external_url(&state, ""), - authorization_endpoint: external_url(&state, "/_/mcp/authorize"), - token_endpoint: external_url(&state, "/_/mcp/token"), - registration_endpoint: external_url(&state, "/_/mcp/register"), + issuer: external_url2(&state, ""), + authorization_endpoint: external_url2(&state, "/_/mcp/authorize"), + token_endpoint: external_url2(&state, "/_/mcp/token"), + registration_endpoint: external_url2(&state, "/_/mcp/register"), response_types_supported: vec!["code"], grant_types_supported: vec!["authorization_code", "refresh_token"], code_challenge_methods_supported: vec!["S256"], @@ -513,7 +517,7 @@ async fn authorize_handler( || query .resource .as_deref() - .is_some_and(|resource| resource != external_url(&state, MCP_PATH)) + .is_some_and(|resource| resource != external_url2(&state, MCP_PATH)) { return Err(OAuthError::invalid_request("invalid authorization request")); } @@ -611,7 +615,7 @@ async fn oauth_token_handler( if request .resource .as_deref() - .is_some_and(|resource| resource != external_url(&state, MCP_PATH)) + .is_some_and(|resource| resource != external_url2(&state, MCP_PATH)) { return Err(OAuthError::invalid_grant("invalid resource")); } @@ -674,7 +678,7 @@ async fn oauth_token_handler( .jwt() .encode(&McpAccessTokenClaims { auth: claims, - aud: external_url(&state, MCP_PATH), + aud: external_url2(&state, MCP_PATH), scope: MCP_SCOPE.to_string(), }) .map_err(|err| OAuthError::server(err.to_string()))?; @@ -764,11 +768,26 @@ fn valid_client_redirect(uri: &str) -> bool { } fn external_url(state: &AppState, path: &str) -> String { + // FIXME: Hard-coded port 4000. + // FIXME: Doesn't work for local workflow were site_url may be set. let mut base = state .site_url() .as_ref() .clone() .unwrap_or_else(|| url::Url::parse("http://localhost:4000").expect("constant URL")); + // let mut base = url::Url::parse("http://localhost:4000").expect("constant URL"); + + base.set_path(path); + base.set_query(None); + base.set_fragment(None); + base.to_string().trim_end_matches('/').to_string() +} + +fn external_url2(state: &AppState, path: &str) -> String { + // FIXME: Hard-coded port 4000. + // FIXME: Doesn't work for local workflow were site_url may be set. + let mut base = url::Url::parse("http://localhost:4000").expect("constant URL"); + base.set_path(path); base.set_query(None); base.set_fragment(None); @@ -912,7 +931,7 @@ mod tests { .jwt() .encode(&McpAccessTokenClaims { auth: claims.clone(), - aud: external_url(&state, MCP_PATH), + aud: external_url2(&state, MCP_PATH), scope: MCP_SCOPE.to_string(), }) .unwrap(); From f632d77f0358b3b64a0e725db33c441ce212c348 Mon Sep 17 00:00:00 2001 From: Sebastian Jeltsch Date: Fri, 14 Aug 2026 11:51:46 +0200 Subject: [PATCH 4/4] Fix resource uri construction. --- crates/core/src/extract/ip.rs | 46 ++++-------- crates/core/src/logging.rs | 2 +- crates/core/src/mcp.rs | 128 +++++++++++++++++++++++++--------- 3 files changed, 107 insertions(+), 69 deletions(-) diff --git a/crates/core/src/extract/ip.rs b/crates/core/src/extract/ip.rs index 2cc576f60..3acd2cd07 100644 --- a/crates/core/src/extract/ip.rs +++ b/crates/core/src/extract/ip.rs @@ -1,34 +1,12 @@ use axum::extract::{ConnectInfo, FromRequestParts}; -use axum::http::uri; -use axum::http::{Request, request::Parts}; +use axum::http::uri::Authority; +use axum::http::{Extensions, HeaderMap, HeaderValue, Request, request::Parts}; use std::convert::Infallible; use std::net::IpAddr; use tower_governor::GovernorError; use tower_governor::key_extractor::KeyExtractor; -pub fn extract_ip(req: &Request) -> Option { - let headers = req.headers(); - - // NOTE: This code is mimicking axum_client_ip's pre v1 `InsecureClientIp::from`: - return client_ip::rightmost_x_forwarded_for(headers) - .or_else(|_| client_ip::x_real_ip(headers)) - .or_else(|_| client_ip::fly_client_ip(headers)) - .or_else(|_| client_ip::true_client_ip(headers)) - .or_else(|_| client_ip::cf_connecting_ip(headers)) - .or_else(|_| client_ip::cloudfront_viewer_address(headers)) - .ok() - .or_else(|| { - req - .extensions() - .get::>() - .map(|ConnectInfo(addr)| addr.ip()) - }); -} - -// TODO: Should replace above function. -pub fn extract_ip2(parts: &Parts) -> Option { - let headers = &parts.headers; - +pub fn extract_ip(headers: &HeaderMap, ext: &Extensions) -> Option { // NOTE: This code is mimicking axum_client_ip's pre v1 `InsecureClientIp::from`: return client_ip::rightmost_x_forwarded_for(headers) .or_else(|_| client_ip::x_real_ip(headers)) @@ -38,8 +16,7 @@ pub fn extract_ip2(parts: &Parts) -> Option { .or_else(|_| client_ip::cloudfront_viewer_address(headers)) .ok() .or_else(|| { - parts - .extensions + ext .get::>() .map(|ConnectInfo(addr)| addr.ip()) }); @@ -53,7 +30,8 @@ impl KeyExtractor for RealIpKeyExtractor { type Key = IpAddr; fn extract(&self, req: &Request) -> Result { - return extract_ip(req).ok_or_else(|| GovernorError::UnableToExtractKey); + return extract_ip(req.headers(), req.extensions()) + .ok_or_else(|| GovernorError::UnableToExtractKey); } // fn name(&self) -> &'static str { @@ -65,7 +43,7 @@ impl KeyExtractor for RealIpKeyExtractor { // } } -// Extractor for handlers. +// RealIp extractor for handlers. #[derive(Debug, Clone, Default)] pub struct RealIp(pub Option); @@ -76,13 +54,13 @@ where type Rejection = Infallible; async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result { - return Ok(Self(extract_ip2(parts))); + return Ok(Self(extract_ip(&parts.headers, &parts.extensions))); } } -// Extractor for handlers. +// Host extractor for handlers. #[derive(Debug, Clone, Default)] -pub struct Host(pub Option); +pub struct Host(pub Option); impl FromRequestParts for Host where @@ -97,13 +75,13 @@ where .or_else(|| parts.headers.get("host")); return Ok(Self( - header.and_then(|h| uri::Authority::try_from(h.as_bytes()).ok()), + header.and_then(|h| Authority::try_from(h.as_bytes()).ok()), )); } } #[allow(unused)] -pub fn ipv6_privacy_mask(ip: IpAddr) -> IpAddr { +fn ipv6_privacy_mask(ip: IpAddr) -> IpAddr { return match ip { IpAddr::V4(ip) => IpAddr::V4(ip), IpAddr::V6(ip) => IpAddr::V6(From::from( diff --git a/crates/core/src/logging.rs b/crates/core/src/logging.rs index dc148c88f..7e407e988 100644 --- a/crates/core/src/logging.rs +++ b/crates/core/src/logging.rs @@ -171,7 +171,7 @@ pub(super) fn sqlite_logger_make_span(request: &Request) -> Span { uri = %request.uri(), version = ?request.version(), host = get_header(headers, "host"), - client_ip = extract_ip(request).map(|ip| ip.to_string()), + client_ip = extract_ip(headers, request.extensions()).map(|ip| ip.to_string()), user_agent = get_header(headers, "user-agent"), referer = get_header(headers, "referer"), // Reserve placeholders that may be recorded later. diff --git a/crates/core/src/mcp.rs b/crates/core/src/mcp.rs index 4a42f9529..cfa36eb17 100644 --- a/crates/core/src/mcp.rs +++ b/crates/core/src/mcp.rs @@ -3,7 +3,8 @@ use std::sync::Arc; use axum::Router; use axum::body::Body; use axum::extract::{Form, Json as AxumJson, Path, Query, Request, State}; -use axum::http::{Method, StatusCode, header}; +use axum::http::uri::Scheme; +use axum::http::{Method, StatusCode, header, uri}; use axum::middleware::{self, Next}; use axum::response::{IntoResponse, Redirect, Response}; use axum::routing::{get, post}; @@ -23,7 +24,7 @@ use crate::admin; use crate::app_state::AppState; use crate::auth::util::is_admin; use crate::auth::{AuthError, AuthTokenClaims, User}; -use crate::extract::ip::{Host, RealIp}; +use crate::extract::ip::Host; const MCP_SCOPE: &str = "mcp"; const MCP_PATH: &str = "/mcp"; @@ -263,10 +264,10 @@ pub(crate) fn router(state: &AppState) -> Router { get(protected_resource_metadata_handler), ) // QUESTION: Is this "/mcp" needed atop the root one above? - .route( - "/.well-known/oauth-protected-resource/mcp", - get(protected_resource_metadata_handler), - ) + // .route( + // "/.well-known/oauth-protected-resource/mcp", + // get(protected_resource_metadata_handler), + // ) .route( "/.well-known/oauth-authorization-server", get(authorization_server_metadata_handler), @@ -283,15 +284,17 @@ pub(crate) fn router(state: &AppState) -> Router { async fn assert_mcp_access( State(state): State, + Host(authority): Host, request: Request, next: Next, ) -> Response { + let authority = authority.unwrap_or_else(|| uri::Authority::from_static("localhost:4000")); let user_id = request .headers() .get(header::AUTHORIZATION) .and_then(|value| value.to_str().ok()) .and_then(|value| value.strip_prefix("Bearer ")) - .and_then(|token| mcp_user_id(&state, token)); + .and_then(|token| mcp_user_id(&state, authority.clone(), token)); if let Some(user_id) = user_id && is_admin(&state, &user_id).await @@ -305,7 +308,7 @@ async fn assert_mcp_access( header::WWW_AUTHENTICATE, format!( r#"Bearer resource_metadata="{metadata}", scope="{MCP_SCOPE}""#, - metadata = external_url2(&state, "/.well-known/oauth-protected-resource"), + metadata = build_uri(&authority, "/.well-known/oauth-protected-resource").unwrap(), ), ) .body(axum::body::Body::empty()) @@ -328,10 +331,12 @@ struct CompatibilityAccessTokenClaims { aud: Option, } -fn mcp_user_id(state: &AppState, token: &str) -> Option { +fn mcp_user_id(state: &AppState, authority: uri::Authority, token: &str) -> Option { + let audience = build_uri(&authority, MCP_PATH)?.to_string(); + if let Ok(claims) = state .jwt() - .decode_with_audience::(token, &external_url2(state, MCP_PATH)) + .decode_with_audience::(token, &audience) { if !claims.scope.split(' ').any(|scope| scope == MCP_SCOPE) { return None; @@ -360,15 +365,14 @@ struct ProtectedResourceMetadata { } async fn protected_resource_metadata_handler( - State(state): State, - RealIp(ip): RealIp, Host(authority): Host, ) -> AxumJson { - let issuer = external_url2(&state, ""); - let resource = external_url2(&state, MCP_PATH); + let authority = authority.unwrap_or_else(|| uri::Authority::from_static("localhost:4000")); + let issuer = build_uri(&authority, "").unwrap().to_string(); + let resource = build_uri(&authority, MCP_PATH).unwrap().to_string(); // FIXME: Remove - log::debug!("ip: {ip:?}, authority: {authority:?}, issuer: {issuer}, resource: {resource}"); + log::debug!("authority: {authority:?}, issuer: {issuer}, resource: {resource}"); AxumJson(ProtectedResourceMetadata { resource, @@ -393,12 +397,19 @@ struct AuthorizationServerMetadata { async fn authorization_server_metadata_handler( State(state): State, + Host(authority): Host, ) -> AxumJson { + let authority = authority.unwrap_or_else(|| uri::Authority::from_static("localhost:4000")); + return AxumJson(AuthorizationServerMetadata { - issuer: external_url2(&state, ""), - authorization_endpoint: external_url2(&state, "/_/mcp/authorize"), - token_endpoint: external_url2(&state, "/_/mcp/token"), - registration_endpoint: external_url2(&state, "/_/mcp/register"), + issuer: build_uri(&authority, "").unwrap().to_string(), + authorization_endpoint: build_uri(&authority, "/_/mcp/authorize") + .unwrap() + .to_string(), + token_endpoint: build_uri(&authority, "/_/mcp/token").unwrap().to_string(), + registration_endpoint: build_uri(&authority, "/_/mcp/register") + .unwrap() + .to_string(), response_types_supported: vec!["code"], grant_types_supported: vec!["authorization_code", "refresh_token"], code_challenge_methods_supported: vec!["S256"], @@ -499,9 +510,12 @@ struct FlowClaims { async fn authorize_handler( State(state): State, + Host(authority): Host, user: Option, Query(query): Query, ) -> Result { + let authority = authority.unwrap_or_else(|| uri::Authority::from_static("localhost:4000")); + let expected_resource = build_uri(&authority, MCP_PATH).unwrap().to_string(); let client: ClientClaims = state .jwt() .decode(&query.client_id) @@ -517,7 +531,7 @@ async fn authorize_handler( || query .resource .as_deref() - .is_some_and(|resource| resource != external_url2(&state, MCP_PATH)) + .is_some_and(|resource| resource != expected_resource) { return Err(OAuthError::invalid_request("invalid authorization request")); } @@ -610,12 +624,15 @@ struct TokenResponse { async fn oauth_token_handler( State(state): State, + Host(authority): Host, Form(request): Form, ) -> Result, OAuthError> { + let authority = authority.unwrap_or_else(|| uri::Authority::from_static("localhost:4000")); + let expected_resource = build_uri(&authority, MCP_PATH).unwrap().to_string(); if request .resource .as_deref() - .is_some_and(|resource| resource != external_url2(&state, MCP_PATH)) + .is_some_and(|resource| resource != expected_resource) { return Err(OAuthError::invalid_grant("invalid resource")); } @@ -678,7 +695,7 @@ async fn oauth_token_handler( .jwt() .encode(&McpAccessTokenClaims { auth: claims, - aud: external_url2(&state, MCP_PATH), + aud: expected_resource, scope: MCP_SCOPE.to_string(), }) .map_err(|err| OAuthError::server(err.to_string()))?; @@ -783,15 +800,46 @@ fn external_url(state: &AppState, path: &str) -> String { base.to_string().trim_end_matches('/').to_string() } -fn external_url2(state: &AppState, path: &str) -> String { - // FIXME: Hard-coded port 4000. - // FIXME: Doesn't work for local workflow were site_url may be set. - let mut base = url::Url::parse("http://localhost:4000").expect("constant URL"); +// fn external_url2(state: &AppState, path: &str) -> String { +// // FIXME: Hard-coded port 4000. +// // FIXME: Doesn't work for local workflow were site_url may be set. +// let mut base = url::Url::parse("http://localhost:4000").expect("constant URL"); +// +// base.set_path(path); +// base.set_query(None); +// base.set_fragment(None); +// base.to_string().trim_end_matches('/').to_string() +// } + +fn build_uri(authority: &uri::Authority, path: &str) -> Option { + return axum::http::uri::Uri::builder() + .scheme(if is_local_host(authority) { + Scheme::HTTP + } else { + Scheme::HTTPS + }) + .authority(authority.clone()) + .path_and_query(path) + .build() + .ok(); +} - base.set_path(path); - base.set_query(None); - base.set_fragment(None); - base.to_string().trim_end_matches('/').to_string() +fn is_local_host(authority: &uri::Authority) -> bool { + let host = authority.host(); + + if host.eq_ignore_ascii_case("localhost") { + return true; + } + + use std::net::IpAddr; + if let Ok(ip) = host.parse::() { + return match ip { + IpAddr::V4(v4) => v4.is_loopback() || v4.is_private() || v4.is_link_local(), + IpAddr::V6(v6) => v6.is_loopback() || v6.is_unicast_link_local(), + }; + } + + return false; } fn normalize_admin_path(path: &str) -> Result { @@ -838,8 +886,10 @@ mod tests { .await .unwrap(); + let authority = uri::Authority::from_static("localhost:4000"); let redirect = authorize_handler( State(state), + Host(Some(authority)), None, Query(AuthorizeQuery { client_id: registration.client_id, @@ -927,15 +977,19 @@ mod tests { csrf_token: "csrf".to_string(), }; + let authority = uri::Authority::from_static("localhost:4000"); let scoped = state .jwt() .encode(&McpAccessTokenClaims { auth: claims.clone(), - aud: external_url2(&state, MCP_PATH), + aud: build_uri(&authority, MCP_PATH), scope: MCP_SCOPE.to_string(), }) .unwrap(); - assert_eq!(mcp_user_id(&state, &scoped), Some(user_id)); + assert_eq!( + mcp_user_id(&state, authority.clone(), &scoped), + Some(user_id) + ); let wrong_audience = state .jwt() @@ -945,9 +999,15 @@ mod tests { scope: MCP_SCOPE.to_string(), }) .unwrap(); - assert_eq!(mcp_user_id(&state, &wrong_audience), None); + assert_eq!( + mcp_user_id(&state, authority.clone(), &wrong_audience), + None + ); let legacy = state.jwt().encode(&claims).unwrap(); - assert_eq!(mcp_user_id(&state, &legacy), Some(user_id)); + assert_eq!( + mcp_user_id(&state, authority.clone(), &legacy), + Some(user_id) + ); } }