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/extract/ip.rs b/crates/core/src/extract/ip.rs index 0354352c1..3acd2cd07 100644 --- a/crates/core/src/extract/ip.rs +++ b/crates/core/src/extract/ip.rs @@ -1,12 +1,12 @@ -use axum::extract::ConnectInfo; -use axum::http::Request; +use axum::extract::{ConnectInfo, FromRequestParts}; +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(); - +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)) @@ -16,13 +16,13 @@ pub fn extract_ip(req: &Request) -> Option { .or_else(|_| client_ip::cloudfront_viewer_address(headers)) .ok() .or_else(|| { - req - .extensions() + ext .get::>() .map(|ConnectInfo(addr)| addr.ip()) }); } +/// Key extractor for the Governor. #[derive(Debug, Clone)] pub struct RealIpKeyExtractor; @@ -30,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 { @@ -42,8 +43,45 @@ impl KeyExtractor for RealIpKeyExtractor { // } } +// RealIp 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_ip(&parts.headers, &parts.extensions))); + } +} + +// Host 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| 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/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/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 new file mode 100644 index 000000000..cfa36eb17 --- /dev/null +++ b/crates/core/src/mcp.rs @@ -0,0 +1,1013 @@ +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::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}; +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}; +use crate::extract::ip::Host; + +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 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 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() { + 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 service: StreamableHttpService = + StreamableHttpService::new( + { + let state = state.clone(); + move || Ok(TrailBaseMcp::new(state.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_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-authorization-server", + get(authorization_server_metadata_handler), + ) + // 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( + 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, authority.clone(), token)); + + if let Some(user_id) = user_id + && is_admin(&state, &user_id).await + { + return next.run(request).await; + } + + return Response::builder() + .status(StatusCode::UNAUTHORIZED) + .header( + header::WWW_AUTHENTICATE, + format!( + r#"Bearer resource_metadata="{metadata}", scope="{MCP_SCOPE}""#, + metadata = build_uri(&authority, "/.well-known/oauth-protected-resource").unwrap(), + ), + ) + .body(axum::body::Body::empty()) + .unwrap_or_default(); +} + +#[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, 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, &audience) + { + 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_handler( + Host(authority): Host, +) -> AxumJson { + 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!("authority: {authority:?}, issuer: {issuer}, resource: {resource}"); + + AxumJson(ProtectedResourceMetadata { + resource, + 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_handler( + State(state): State, + Host(authority): Host, +) -> AxumJson { + let authority = authority.unwrap_or_else(|| uri::Authority::from_static("localhost:4000")); + + return AxumJson(AuthorizationServerMetadata { + 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"], + 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_handler( + 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 client_id = state + .jwt() + .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()))?; + + return 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_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) + .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 != expected_resource) + { + return Err(OAuthError::invalid_request("invalid authorization request")); + } + + 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)?; + + 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)] +struct CallbackQuery { + code: String, +} + +async fn authorization_callback_handler( + 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_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 != expected_resource) + { + 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: expected_resource, + 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 { + // 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); +// 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(); +} + +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 { + 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('/'); + + return 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_handler( + 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 authority = uri::Authority::from_static("localhost:4000"); + let redirect = authorize_handler( + State(state), + Host(Some(authority)), + 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 authority = uri::Authority::from_static("localhost:4000"); + let scoped = state + .jwt() + .encode(&McpAccessTokenClaims { + auth: claims.clone(), + aud: build_uri(&authority, MCP_PATH), + scope: MCP_SCOPE.to_string(), + }) + .unwrap(); + assert_eq!( + mcp_user_id(&state, authority.clone(), &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, authority.clone(), &wrong_audience), + None + ); + + let legacy = state.jwt().encode(&claims).unwrap(); + assert_eq!( + mcp_user_id(&state, authority.clone(), &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(