diff --git a/Cargo.lock b/Cargo.lock index 8b941554..eab90297 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -309,6 +309,7 @@ dependencies = [ "aisix-core", "aisix-gateway", "aisix-guardrails", + "aisix-mcp", "aisix-obs", "aisix-provider-anthropic", "aisix-provider-openai", diff --git a/crates/aisix-mcp/src/gateway.rs b/crates/aisix-mcp/src/gateway.rs index 1f9acbfa..22ec78ea 100644 --- a/crates/aisix-mcp/src/gateway.rs +++ b/crates/aisix-mcp/src/gateway.rs @@ -221,6 +221,15 @@ pub fn streamable_http_service( let mut config = StreamableHttpServerConfig::default(); config.stateful_mode = false; config.json_response = true; + // Disable rmcp's `Host`-header allowlist. Its default + // (`localhost`/`127.0.0.1`/`::1`) is a DNS-rebinding guard for + // browser-driven local servers — it 403s every request whose `Host` is + // the deployment's real DNS name. This endpoint is not browser-driven: it + // is reached server-to-server by agents and is gated by the AISIX API key, + // which is the real access control. An empty allowlist accepts any `Host`; + // the request is still authenticated upstream of this service. (Operators + // who want Host pinning can layer it at their ingress.) + config.allowed_hosts = Vec::new(); StreamableHttpService::new( move || Ok(gateway.clone()), Arc::new(LocalSessionManager::default()), diff --git a/crates/aisix-proxy/Cargo.toml b/crates/aisix-proxy/Cargo.toml index b443d49d..06d1a424 100644 --- a/crates/aisix-proxy/Cargo.toml +++ b/crates/aisix-proxy/Cargo.toml @@ -11,6 +11,9 @@ description = "aisix: /v1/* proxy router (OpenAI-compatible) + middleware + hook [dependencies] aisix-core = { path = "../aisix-core" } aisix-gateway = { path = "../aisix-gateway" } +# Mounts the downstream-facing `/mcp` MCP gateway endpoint (mcp.rs): the +# proxy builds an `McpGateway` from the snapshot's `mcp_servers` per request. +aisix-mcp = { path = "../aisix-mcp" } # /v1/messages handler reaches into the Anthropic wire helpers # (parse_inbound_request / chat_response_into_anthropic_json / # AnthropicSseEncoder) so it can translate Anthropic-protocol diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 26a76e0f..63777545 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -42,6 +42,7 @@ mod error_translate; pub mod health; mod http_client; mod images; +mod mcp; mod messages; mod models; mod passthrough; @@ -105,6 +106,10 @@ pub fn build_router(state: ProxyState) -> Router { "/passthrough/:provider/*rest", any(passthrough::passthrough), ) + // Downstream-facing MCP gateway. Authentication (AISIX API key) is + // enforced inside the handler via the `AuthenticatedKey` extractor. + .route("/mcp", any(mcp::mcp_endpoint)) + .route("/mcp/", any(mcp::mcp_endpoint)) // Wire the configured cap into axum's request-body extractor // chain (`Json` defers to `Bytes` which honors this layer). // Without this, axum 0.7's `DefaultBodyLimit` falls back to @@ -178,6 +183,7 @@ fn normalize_endpoint_label(path: &str) -> &'static str { "/v1/audio/transcriptions" => "/v1/audio/transcriptions", "/v1/audio/translations" => "/v1/audio/translations", "/v1/audio/speech" => "/v1/audio/speech", + "/mcp" | "/mcp/" => "/mcp", _ if path.starts_with("/passthrough/") => "/passthrough/:provider/*rest", _ => "other", } @@ -186,6 +192,8 @@ fn normalize_endpoint_label(path: &str) -> &'static str { fn inbound_protocol_for_endpoint(endpoint: &str) -> &'static str { if endpoint == "/v1/messages" || endpoint == "/v1/messages/count_tokens" { "anthropic" + } else if endpoint == "/mcp" { + "mcp" } else { "openai" } diff --git a/crates/aisix-proxy/src/mcp.rs b/crates/aisix-proxy/src/mcp.rs new file mode 100644 index 00000000..d449ac1b --- /dev/null +++ b/crates/aisix-proxy/src/mcp.rs @@ -0,0 +1,202 @@ +//! `/mcp` — the downstream-facing MCP gateway endpoint. +//! +//! AISIX presents as a single MCP server to a downstream agent: it aggregates +//! the tools of the registered `mcp_servers` and routes tool calls back to +//! them. The caller authenticates with an AISIX API key — the +//! [`AuthenticatedKey`] extractor rejects a missing or invalid key with `401` +//! before the request reaches the gateway. The gateway is rebuilt from the +//! current configuration snapshot on each request, so it always reflects the +//! live `mcp_servers` set. +//! +//! Per-tool access control, guardrail / quota reuse, and usage logging over MCP +//! traffic are layered on in subsequent steps; this step establishes the +//! authenticated, snapshot-sourced endpoint. + +use axum::body::Body; +use axum::extract::{Request, State}; +use axum::response::Response; +use tower::ServiceExt; + +use crate::auth::AuthenticatedKey; +use crate::state::ProxyState; + +/// Serve a `/mcp` request. The [`AuthenticatedKey`] extractor enforces a valid +/// AISIX API key (responding `401` otherwise); the request is then handled by an +/// MCP gateway built from the current snapshot's `mcp_servers`. +pub async fn mcp_endpoint( + _auth: AuthenticatedKey, + State(state): State, + request: Request, +) -> Response { + let snapshot = state.snapshot.load(); + let gateway = aisix_mcp::McpGateway::from_snapshot(&snapshot); + let service = aisix_mcp::streamable_http_service(gateway); + // `StreamableHttpService` is a tower service that dispatches on method and + // never fails (`Error = Infallible`); map its boxed body back to axum's. + match service.oneshot(request).await { + Ok(response) => response.map(Body::new), + Err(infallible) => match infallible {}, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::build_router; + use aisix_core::{AisixSnapshot, ApiKey, ProxyConfig, ResourceEntry, SnapshotHandle}; + use axum::body::Body; + use axum::http::{Request as HttpRequest, StatusCode}; + use std::sync::Arc; + + fn cfg() -> ProxyConfig { + ProxyConfig { + addr: "127.0.0.1:0".into(), + request_body_limit_bytes: 1_048_576, + real_ip: Default::default(), + tls: None, + } + } + + const TOKEN: &str = "sk-mcp-endpoint-test"; + + /// A snapshot carrying one valid API key (and no MCP servers — the MCP + /// `initialize` handshake is answered by the gateway itself, no upstream + /// needed). + fn snapshot_with_key() -> AisixSnapshot { + let key_hash = ApiKey::hash_bearer(TOKEN); + let apikey: ApiKey = serde_json::from_value(serde_json::json!({ + "key_hash": key_hash, + "allowed_models": ["*"], + })) + .expect("valid apikey"); + let snapshot = AisixSnapshot::new(); + snapshot + .apikeys + .insert(ResourceEntry::new("ak-1", apikey, 1)); + snapshot + } + + fn router_with(snapshot: AisixSnapshot) -> axum::Router { + let handle = SnapshotHandle::new(snapshot); + let hub = Arc::new(aisix_gateway::Hub::new()); + build_router(ProxyState::new(handle, hub, &cfg()).without_cache()) + } + + /// A minimal MCP `initialize` request body + the headers the Streamable + /// HTTP transport requires (Accept must list both content types). + fn initialize_request(auth: Option<&str>) -> HttpRequest { + let body = serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-11-25", + "capabilities": {}, + "clientInfo": { "name": "endpoint-test", "version": "0.1" } + } + }); + // A non-loopback Host on purpose: proves the gateway accepts the + // deployment's real DNS name (rmcp's default Host allowlist is disabled + // for this key-authenticated endpoint). + let mut builder = HttpRequest::post("/mcp") + .header("host", "mcp.aisix.example.com") + .header("content-type", "application/json") + .header("accept", "application/json, text/event-stream"); + if let Some(token) = auth { + builder = builder.header("authorization", format!("Bearer {token}")); + } + builder.body(Body::from(body.to_string())).unwrap() + } + + #[tokio::test] + async fn rejects_request_without_api_key() { + let router = router_with(snapshot_with_key()); + let resp = router + .oneshot(initialize_request(None)) + .await + .expect("router responds"); + assert_eq!( + resp.status(), + StatusCode::UNAUTHORIZED, + "missing API key must be rejected at the /mcp edge" + ); + } + + #[tokio::test] + async fn rejects_request_with_invalid_api_key() { + let router = router_with(snapshot_with_key()); + let resp = router + .oneshot(initialize_request(Some("sk-wrong"))) + .await + .expect("router responds"); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + } + + #[tokio::test] + async fn auth_gates_non_post_methods() { + // The route is `any(...)`, so every method must be auth-gated — a GET + // with no key must 401 (not fall through to rmcp's 405). + let router = router_with(snapshot_with_key()); + let req = HttpRequest::get("/mcp") + .header("host", "mcp.aisix.example.com") + .body(Body::empty()) + .unwrap(); + let resp = router.oneshot(req).await.expect("router responds"); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + } + + #[tokio::test] + async fn trailing_slash_route_is_auth_gated() { + let router = router_with(snapshot_with_key()); + let req = HttpRequest::post("/mcp/") + .header("host", "mcp.aisix.example.com") + .header("content-type", "application/json") + .header("accept", "application/json, text/event-stream") + .body(Body::from("{}")) + .unwrap(); + let resp = router.oneshot(req).await.expect("router responds"); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + } + + #[tokio::test] + async fn oversized_unauthenticated_body_is_limited_before_handler() { + // A declared Content-Length over the cap is rejected (413) by the + // body-limit layer, which wraps the route — before auth or the handler, + // so an oversized unauthenticated body can't pin resources. + let router = router_with(snapshot_with_key()); + let big = "a".repeat(1_048_577); // cfg() cap is 1 MiB + let req = HttpRequest::post("/mcp") + .header("host", "mcp.aisix.example.com") + .header("content-type", "application/json") + .header("content-length", big.len().to_string()) + .body(Body::from(big)) + .unwrap(); + let resp = router.oneshot(req).await.expect("router responds"); + let status = resp.status(); + assert_eq!(status, StatusCode::PAYLOAD_TOO_LARGE, "got {status}"); + } + + #[tokio::test] + async fn authenticated_request_reaches_the_mcp_gateway() { + let router = router_with(snapshot_with_key()); + let resp = router + .oneshot(initialize_request(Some(TOKEN))) + .await + .expect("router responds"); + // Auth passed and the request was served by the MCP gateway (not a 401). + let status = resp.status(); + let body = axum::body::to_bytes(resp.into_body(), 64 * 1024) + .await + .expect("read body"); + let text = String::from_utf8_lossy(&body); + assert_eq!( + status, + StatusCode::OK, + "a valid key should reach the gateway and complete the MCP initialize handshake; body: {text}" + ); + assert!( + text.contains("serverInfo") || text.contains("protocolVersion"), + "initialize result should carry the server info, got: {text}" + ); + } +}