From 8259ec3017432c0b0182d7f8f1c084414a23b451 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Tue, 4 Aug 2026 16:57:14 +0800 Subject: [PATCH 1/5] feat(mcp): single-server endpoint /mcp/{server} with original tool names The path names a registered mcp_servers entry; the endpoint presents that server alone: initialize reports its name, tools/list returns the upstream's original (un-namespaced) tool names, and tools/call accepts both the bare and the namespaced spelling. A name prefixed with a different registered server fails closed; an unregistered prefix stays a bare name. Unknown/disabled servers 404 after auth, with no fallback to the aggregate. Governance is the same pipeline as /mcp: per-tool ACL still evaluates the namespaced form (a grant means one thing on every endpoint, and the scoped surface cannot widen a key's grant), per-server rate limits and usage attribution key on the path's server rather than the tool-name prefix, and guardrails run both directions. /mcp itself is unchanged. Part of AISIX-Cloud#1219: gives per-server-URL clients a native scoped endpoint; the entry-level URL rewrite that maps legacy URL shapes onto it ships separately. --- crates/aisix-mcp/src/gateway.rs | 155 +++++++-- crates/aisix-mcp/tests/gateway_scoped.rs | 310 ++++++++++++++++++ crates/aisix-proxy/src/lib.rs | 6 +- crates/aisix-proxy/src/mcp.rs | 275 +++++++++++++++- .../src/cases/mcp-scoped-endpoint-e2e.test.ts | 302 +++++++++++++++++ 5 files changed, 1016 insertions(+), 32 deletions(-) create mode 100644 crates/aisix-mcp/tests/gateway_scoped.rs create mode 100644 tests/e2e/src/cases/mcp-scoped-endpoint-e2e.test.ts diff --git a/crates/aisix-mcp/src/gateway.rs b/crates/aisix-mcp/src/gateway.rs index ff1f5b3a4..080a49da4 100644 --- a/crates/aisix-mcp/src/gateway.rs +++ b/crates/aisix-mcp/src/gateway.rs @@ -14,6 +14,11 @@ //! - `tools/call` strips the namespace prefix and routes to the owning //! upstream. //! +//! A gateway may instead be **scoped** to a single upstream +//! ([`McpGateway::from_snapshot_scoped`], mounted at `/mcp/{server}`): it then +//! serves that server's tools under their original, un-namespaced names while +//! ACL decisions keep evaluating the namespaced form. +//! //! The aggregator holds no per-request or per-session state, so governance //! never depends on a transport session — which keeps it aligned with the //! stateless direction of the MCP 2026-07-28 revision. @@ -229,6 +234,24 @@ impl ToolAcl { pub struct McpGateway { upstreams: Arc<[NamedUpstream]>, tool_acl: ToolAcl, + /// When set, this gateway serves exactly one upstream under its **original** + /// tool names: `tools/list` strips the `__` namespace prefix and + /// `tools/call` accepts both the bare and the prefixed form. ACL decisions + /// still evaluate the namespaced form, so per-tool grants keep one meaning + /// across the aggregated and the scoped endpoint. + scoped: Option>, +} + +/// The single-server scope of a gateway built by +/// [`McpGateway::from_snapshot_scoped`]. +struct ScopedServer { + /// The scoped upstream's registered name — the namespace every ACL check + /// re-applies and the name `initialize` reports. + name: String, + /// The **other** registered, enabled server names. A `tools/call` whose + /// name is prefixed with one of these is a cross-server mistake and fails + /// closed rather than being silently served as a bare name. + foreign: std::collections::HashSet, } impl McpGateway { @@ -265,6 +288,7 @@ impl McpGateway { Self { upstreams: deduped.into(), tool_acl: ToolAcl::allow_all(), + scoped: None, } } @@ -303,6 +327,36 @@ impl McpGateway { McpGateway::new(upstreams) } + /// Build a gateway scoped to the single **enabled** `mcp_servers` entry + /// named `server`, serving its tools under their original names (see + /// [`McpGateway::scoped`]). Returns `None` when the server is not + /// registered or is disabled — a disabled server is treated as absent, + /// same as the aggregated endpoint skipping it. + pub fn from_snapshot_scoped(snapshot: &AisixSnapshot, server: &str) -> Option { + let entry = snapshot.mcp_servers.get_by_name(server)?; + if !entry.value.enabled { + return None; + } + let name = entry.value.name.clone(); + let bridge: Arc = match entry.value.server_type { + McpServerType::Mcp => { + let upstream = upstream_from_mcp_server(&entry.value); + Arc::new(EphemeralBridge::new(upstream)) + } + McpServerType::Openapi => Arc::new(OpenApiBridge::new(entry)), + }; + let foreign = snapshot + .mcp_servers + .entries() + .into_iter() + .filter(|e| e.value.enabled && e.value.name != name) + .map(|e| e.value.name.clone()) + .collect(); + let mut gateway = McpGateway::new([(name.clone(), bridge)]); + gateway.scoped = Some(Arc::new(ScopedServer { name, foreign })); + Some(gateway) + } + fn find(&self, server: &str) -> Option<&Arc> { self.upstreams .iter() @@ -345,6 +399,18 @@ impl ServerHandler for McpGateway { } // Per-tool ACL: expose only the tools this caller's key permits. tools.retain(|tool| self.tool_acl.permits(tool.name.as_ref())); + // A scoped gateway serves its single upstream's tools under their + // original names — the namespace prefix exists to disambiguate the + // aggregate, and a single-server endpoint has nothing to disambiguate. + // ACL filtering above ran on the namespaced form. + if let Some(scoped) = &self.scoped { + let prefix = format!("{}{TOOL_NAMESPACE_SEPARATOR}", scoped.name); + for tool in &mut tools { + if let Some(bare) = tool.name.strip_prefix(&prefix) { + tool.name = Cow::Owned(bare.to_string()); + } + } + } Ok(ListToolsResult::with_all_items(tools)) } @@ -353,26 +419,63 @@ impl ServerHandler for McpGateway { request: CallToolRequestParams, _context: RequestContext, ) -> Result { - let (server, tool) = request - .name - .split_once(TOOL_NAMESPACE_SEPARATOR) - .ok_or_else(|| { - ErrorData::invalid_params( - format!( - "tool name '{}' is missing a 'server__tool' prefix", - request.name - ), - None, + // Resolve `(server, tool, namespaced)` from the caller's tool name. + // Scoped: the name is the upstream's original one, but a caller that + // sends the namespaced form anyway (an aggregated-endpoint client + // pointed at the scoped URL) keeps working — `__x` reduces to + // `x`. (An upstream tool literally named `__x` is therefore + // only callable as `____x` — prefix-stripping takes + // precedence, pinned by test.) A name prefixed with a *different* + // registered server's name fails closed: the scoped endpoint never + // cross-routes, and silently serving it as a bare name would mask the + // caller's mistake. An unregistered prefix stays a bare name, since + // tool names may legitimately contain the separator. + let request_name = request.name.as_ref(); + let (server, tool, namespaced): (&str, &str, Cow<'_, str>) = match &self.scoped { + Some(scoped) => { + let scope = scoped.name.as_str(); + let bare = match request_name.split_once(TOOL_NAMESPACE_SEPARATOR) { + Some((prefix, rest)) if prefix == scope => rest, + Some((prefix, _)) if scoped.foreign.contains(prefix) => { + // Same neutral wording as the ACL reject: don't + // confirm what the other server serves. + return Err(ErrorData::invalid_params( + format!("tool '{request_name}' is not available"), + None, + )); + } + _ => request_name, + }; + ( + scope, + bare, + Cow::Owned(format!("{scope}{TOOL_NAMESPACE_SEPARATOR}{bare}")), ) - })?; + } + None => { + let (server, tool) = request_name + .split_once(TOOL_NAMESPACE_SEPARATOR) + .ok_or_else(|| { + ErrorData::invalid_params( + format!( + "tool name '{request_name}' is missing a 'server__tool' prefix" + ), + None, + ) + })?; + (server, tool, Cow::Borrowed(request_name)) + } + }; // Per-tool ACL: reject a call the caller's key doesn't permit. A // disallowed tool is also absent from `tools/list`, so this is // defense-in-depth; the message stays neutral and does not reveal - // whether the tool exists upstream. - if !self.tool_acl.permits(request.name.as_ref()) { + // whether the tool exists upstream. The check runs on the namespaced + // form so grants mean the same thing on every endpoint; the message + // echoes the caller's own spelling. + if !self.tool_acl.permits(&namespaced) { return Err(ErrorData::invalid_params( - format!("tool '{}' is not available", request.name), + format!("tool '{request_name}' is not available"), None, )); } @@ -407,11 +510,25 @@ impl ServerHandler for McpGateway { fn get_info(&self) -> ServerInfo { let mut info = ServerInfo::default(); info.capabilities = ServerCapabilities::builder().enable_tools().build(); - info.instructions = Some( - "AISIX MCP gateway: aggregates tools from registered upstream MCP \ - servers, namespaced as `server__tool`." - .to_string(), - ); + match &self.scoped { + // The scoped endpoint presents as the upstream server itself, so + // `initialize` reports that server's registered name. + Some(scoped) => { + info.server_info.name = scoped.name.clone(); + info.instructions = Some(format!( + "AISIX MCP gateway: serves the tools of MCP server `{}` \ + under their original names.", + scoped.name + )); + } + None => { + info.instructions = Some( + "AISIX MCP gateway: aggregates tools from registered upstream MCP \ + servers, namespaced as `server__tool`." + .to_string(), + ); + } + } info } } diff --git a/crates/aisix-mcp/tests/gateway_scoped.rs b/crates/aisix-mcp/tests/gateway_scoped.rs new file mode 100644 index 000000000..264d765fe --- /dev/null +++ b/crates/aisix-mcp/tests/gateway_scoped.rs @@ -0,0 +1,310 @@ +//! End-to-end test of the **scoped** gateway (`/mcp/{server}`): AISIX as an +//! MCP server fronting exactly one registered upstream, serving its tools +//! under their original (un-namespaced) names. +//! +//! Topology, all real Streamable HTTP over ephemeral ports (no mock +//! transport): +//! +//! downstream rmcp client ──► McpGateway (scoped "alpha") ──► upstream "alpha" (echo) +//! +//! Pins: `initialize` reports the scoped server's name; `tools/list` strips +//! the `alpha__` prefix; `tools/call` accepts both the bare and the +//! namespaced spelling; a foreign prefix stays pinned to the scoped upstream; +//! ACL patterns keep their namespaced meaning; unknown/disabled servers +//! resolve to no gateway at all. + +use std::net::SocketAddr; +use std::sync::Arc; + +use aisix_core::{AisixSnapshot, McpServer, ResourceEntry}; +use aisix_mcp::{streamable_http_service, McpGateway, ToolAcl}; +use rmcp::model::{ + CallToolRequestParams, CallToolResult, Content, ErrorData, ListToolsResult, + PaginatedRequestParams, ServerCapabilities, ServerInfo, Tool, +}; +use rmcp::service::RequestContext; +use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; +use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService}; +use rmcp::transport::StreamableHttpClientTransport; +use rmcp::{RoleServer, ServerHandler, ServiceExt}; + +/// A real upstream MCP server exposing one echo tool under `tool_name`, +/// prefixing its reply with `label` so routing is observable. +#[derive(Clone)] +struct LabeledEcho { + label: &'static str, + tool_name: &'static str, +} + +impl ServerHandler for LabeledEcho { + async fn list_tools( + &self, + _request: Option, + _context: RequestContext, + ) -> Result { + let schema = serde_json::json!({ + "type": "object", + "properties": { "text": { "type": "string" } }, + "required": ["text"], + }); + let tool = Tool::new( + self.tool_name, + "Echo back the provided text", + schema.as_object().expect("schema is an object").clone(), + ); + Ok(ListToolsResult::with_all_items(vec![tool])) + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + _context: RequestContext, + ) -> Result { + if request.name != self.tool_name { + return Err(ErrorData::invalid_params( + format!("unknown tool: {}", request.name), + None, + )); + } + let text = request + .arguments + .as_ref() + .and_then(|m| m.get("text")) + .and_then(|v| v.as_str()) + .unwrap_or_default(); + Ok(CallToolResult::success(vec![Content::text(format!( + "{}:{text}", + self.label + ))])) + } + + fn get_info(&self) -> ServerInfo { + let mut info = ServerInfo::default(); + info.capabilities = ServerCapabilities::builder().enable_tools().build(); + info + } +} + +/// Start a labeled upstream echo server; return its bound address. +async fn spawn_upstream(label: &'static str, tool_name: &'static str) -> SocketAddr { + let service = StreamableHttpService::new( + move || Ok(LabeledEcho { label, tool_name }), + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default(), + ); + serve(axum::Router::new().nest_service("/mcp", service)).await +} + +/// Serve the gateway itself; return its bound address. +async fn spawn_gateway(gateway: McpGateway) -> SocketAddr { + serve(axum::Router::new().nest_service("/mcp", streamable_http_service(gateway))).await +} + +async fn serve(app: axum::Router) -> SocketAddr { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind ephemeral port"); + let addr = listener.local_addr().expect("local addr"); + tokio::spawn(async move { + axum::serve(listener, app).await.expect("serve"); + }); + addr +} + +/// Build a snapshot resource entry for an upstream at `addr`. +fn mcp_entry(id: &str, name: &str, addr: &SocketAddr, enabled: bool) -> ResourceEntry { + let server: McpServer = serde_json::from_value(serde_json::json!({ + "display_name": name, + "url": format!("http://{addr}/mcp"), + "enabled": enabled + })) + .unwrap(); + ResourceEntry::new(id, server, 1) +} + +/// A snapshot with one enabled server `alpha` at `addr` and one disabled +/// server `dark`. +fn snapshot_with_alpha(addr: &SocketAddr) -> AisixSnapshot { + let snapshot = AisixSnapshot::new(); + snapshot + .mcp_servers + .insert(mcp_entry("e1", "alpha", addr, true)); + snapshot + .mcp_servers + .insert(mcp_entry("e2", "dark", addr, false)); + snapshot +} + +async fn connect(gw_addr: SocketAddr) -> rmcp::service::RunningService { + ().serve(StreamableHttpClientTransport::from_uri(format!( + "http://{gw_addr}/mcp" + ))) + .await + .expect("downstream client connects to gateway") +} + +fn call(name: &'static str, text: &str) -> CallToolRequestParams { + let args = serde_json::json!({ "text": text }); + CallToolRequestParams::new(name).with_arguments(args.as_object().unwrap().clone()) +} + +/// Decode the first text content block of a tool result. +fn first_text(result: &CallToolResult) -> String { + let value = serde_json::to_value(&result.content).expect("encode content"); + value[0]["text"].as_str().unwrap_or_default().to_string() +} + +#[tokio::test] +async fn scoped_serves_original_names_and_accepts_both_call_forms() { + let upstream = spawn_upstream("alpha", "echo").await; + let snapshot = snapshot_with_alpha(&upstream); + let gateway = + McpGateway::from_snapshot_scoped(&snapshot, "alpha").expect("alpha is registered"); + let gw_addr = spawn_gateway(gateway).await; + let client = connect(gw_addr).await; + + // `initialize` presents the scoped server, not the aggregate. + let info = client.peer_info().expect("initialize completed"); + assert_eq!(info.server_info.name, "alpha"); + + // tools/list carries the upstream's original names — no `alpha__` prefix. + let tools = client.list_all_tools().await.expect("list tools"); + let names: Vec<&str> = tools.iter().map(|t| t.name.as_ref()).collect(); + assert_eq!(names, vec!["echo"], "original name only: {names:?}"); + + // The bare (original) spelling calls through… + let bare = client + .call_tool(call("echo", "hi")) + .await + .expect("bare call"); + assert_eq!(first_text(&bare), "alpha:hi"); + + // …and the namespaced spelling still works, so an aggregated-endpoint + // client pointed at the scoped URL keeps functioning. + let prefixed = client + .call_tool(call("alpha__echo", "hi")) + .await + .expect("namespaced call"); + assert_eq!(first_text(&prefixed), "alpha:hi"); +} + +#[tokio::test] +async fn scoped_missing_or_disabled_server_resolves_to_none() { + let upstream = spawn_upstream("alpha", "echo").await; + let snapshot = snapshot_with_alpha(&upstream); + + assert!( + McpGateway::from_snapshot_scoped(&snapshot, "ghost").is_none(), + "unregistered server must not resolve" + ); + assert!( + McpGateway::from_snapshot_scoped(&snapshot, "dark").is_none(), + "disabled server must resolve like a missing one" + ); +} + +#[tokio::test] +async fn scoped_acl_keeps_namespaced_meaning() { + let upstream = spawn_upstream("alpha", "echo").await; + let snapshot = snapshot_with_alpha(&upstream); + + // A grant written against the aggregated form covers the scoped endpoint. + let gateway = McpGateway::from_snapshot_scoped(&snapshot, "alpha") + .expect("alpha is registered") + .with_tool_acl(ToolAcl::from_allowed(Some(&["alpha__echo".to_string()]))); + let client = connect(spawn_gateway(gateway).await).await; + let tools = client.list_all_tools().await.expect("list tools"); + assert_eq!(tools.len(), 1, "granted tool listed (bare)"); + assert_eq!(tools[0].name.as_ref(), "echo"); + let ok = client.call_tool(call("echo", "hi")).await.expect("allowed"); + assert_eq!(first_text(&ok), "alpha:hi"); + + // A grant for a different server admits nothing here — the bare name is + // re-namespaced before the check, so the scoped surface cannot widen a + // key's grant. + let gateway = McpGateway::from_snapshot_scoped(&snapshot, "alpha") + .expect("alpha is registered") + .with_tool_acl(ToolAcl::from_allowed(Some(&["beta__*".to_string()]))); + let client = connect(spawn_gateway(gateway).await).await; + let tools = client.list_all_tools().await.expect("list tools"); + assert!(tools.is_empty(), "foreign grant must expose nothing"); + assert!( + client.call_tool(call("echo", "hi")).await.is_err(), + "foreign grant must not admit a bare-name call" + ); +} + +#[tokio::test] +async fn scoped_registered_foreign_prefix_fails_closed() { + // The upstream serves a tool literally named `beta__echo`, and `beta` IS + // a registered, enabled server. On alpha's scoped gateway that spelling + // must fail closed — a cross-server mistake, never silently served as a + // bare name (which this upstream would happily answer). + let upstream = spawn_upstream("alpha", "beta__echo").await; + let snapshot = snapshot_with_alpha(&upstream); + snapshot + .mcp_servers + .insert(mcp_entry("e3", "beta", &upstream, true)); + let gateway = + McpGateway::from_snapshot_scoped(&snapshot, "alpha").expect("alpha is registered"); + let client = connect(spawn_gateway(gateway).await).await; + + assert!( + client.call_tool(call("beta__echo", "x")).await.is_err(), + "a registered foreign prefix must fail closed, never route or serve" + ); + + // The shadowed literal tool stays reachable through its namespaced + // spelling: `alpha__beta__echo` strips alpha's own prefix first. + let escaped = client + .call_tool(call("alpha__beta__echo", "hi")) + .await + .expect("the namespaced spelling reaches the literal tool"); + assert_eq!(first_text(&escaped), "alpha:hi"); +} + +#[tokio::test] +async fn scoped_unregistered_prefix_is_a_bare_name() { + // `ghost` is not a registered server, so `ghost__echo` is just a tool + // name that happens to contain the separator — it must reach the scoped + // upstream verbatim (which serves exactly that name here). + let upstream = spawn_upstream("alpha", "ghost__echo").await; + let snapshot = snapshot_with_alpha(&upstream); + let gateway = + McpGateway::from_snapshot_scoped(&snapshot, "alpha").expect("alpha is registered"); + let client = connect(spawn_gateway(gateway).await).await; + + let served = client + .call_tool(call("ghost__echo", "hi")) + .await + .expect("an unregistered prefix stays a bare tool name"); + assert_eq!(first_text(&served), "alpha:hi"); +} + +#[tokio::test] +async fn scoped_upstream_tool_spelled_like_the_prefix_keeps_precedence() { + // Pathological upstream: a tool literally named `alpha__echo` on server + // `alpha`. Prefix-stripping takes precedence (documented in `call_tool`): + // the tool lists as `alpha__echo` and is callable as `alpha__alpha__echo`; + // the spelling `alpha__echo` reduces to `echo`, which does not exist. + let upstream = spawn_upstream("alpha", "alpha__echo").await; + let snapshot = snapshot_with_alpha(&upstream); + let gateway = + McpGateway::from_snapshot_scoped(&snapshot, "alpha").expect("alpha is registered"); + let client = connect(spawn_gateway(gateway).await).await; + + let tools = client.list_all_tools().await.expect("list tools"); + let names: Vec<&str> = tools.iter().map(|t| t.name.as_ref()).collect(); + assert_eq!(names, vec!["alpha__echo"], "one strip only: {names:?}"); + + let double = client + .call_tool(call("alpha__alpha__echo", "hi")) + .await + .expect("double-prefixed spelling reaches the literal tool"); + assert_eq!(first_text(&double), "alpha:hi"); + + assert!( + client.call_tool(call("alpha__echo", "x")).await.is_err(), + "the single-prefixed spelling reduces to `echo`, which must error" + ); +} diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index d9b328985..1e42521d1 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -160,8 +160,11 @@ pub fn build_router(state: ProxyState) -> Router { ) // Downstream-facing MCP gateway. Authentication (AISIX API key) is // enforced inside the handler via the `AuthenticatedKey` extractor. + // `/mcp/{server}` is the single-server variant (original tool names); + // the static `/mcp/` route wins over the param route for that path. .route("/mcp", any(mcp::mcp_endpoint)) .route("/mcp/", any(mcp::mcp_endpoint)) + .route("/mcp/:server", any(mcp::mcp_scoped_endpoint)) // Downstream-facing A2A gateway. One route per registered agent; the // agent's card (with the service URL rewritten to the gateway) is served // at the RFC 8615 well-known path under it. Authentication (AISIX API @@ -275,6 +278,7 @@ fn normalize_endpoint_label(path: &str) -> &'static str { _ if path.starts_with("/v1/files/") => "/v1/files/:id", _ if path.starts_with("/v1/batches/") => "/v1/batches/:id", _ if path.starts_with("/v1/fine_tuning/jobs/") => "/v1/fine_tuning/jobs/:id", + _ if path.starts_with("/mcp/") => "/mcp/{server}", _ if path.starts_with("/a2a/") => "/a2a", _ if path.starts_with("/passthrough/") => "/passthrough/:provider/*rest", _ => "other", @@ -284,7 +288,7 @@ 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" { + } else if endpoint == "/mcp" || endpoint == "/mcp/{server}" { "mcp" } else if endpoint == "/a2a" { "a2a" diff --git a/crates/aisix-proxy/src/mcp.rs b/crates/aisix-proxy/src/mcp.rs index ee2452d14..3290d444e 100644 --- a/crates/aisix-proxy/src/mcp.rs +++ b/crates/aisix-proxy/src/mcp.rs @@ -1,8 +1,9 @@ -//! `/mcp` — the downstream-facing MCP gateway endpoint. +//! `/mcp` and `/mcp/{server}` — the downstream-facing MCP gateway endpoints. //! //! 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 +//! them. `/mcp/{server}` scopes the same gateway to one registered server, +//! serving its tools under their original (un-namespaced) names. 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 @@ -63,6 +64,32 @@ pub async fn mcp_endpoint( auth: AuthenticatedKey, State(state): State, request: Request, +) -> Response { + serve(auth, state, request, None).await +} + +/// Serve a `/mcp/{server}` request: the single-server variant of +/// [`mcp_endpoint`]. The path names a registered MCP server; the gateway is +/// scoped to it and serves its tools under their original names (`tools/call` +/// also accepts the namespaced form). Everything else — auth, per-tool ACL, +/// quota, guardrails, usage — is the same pipeline as the aggregated endpoint; +/// only the server selection and the tool-name surface differ. An unknown or +/// disabled server is `404` (after auth, so an unauthenticated caller learns +/// nothing about which servers exist). +pub async fn mcp_scoped_endpoint( + auth: AuthenticatedKey, + axum::extract::Path(server): axum::extract::Path, + State(state): State, + request: Request, +) -> Response { + serve(auth, state, request, Some(server)).await +} + +async fn serve( + auth: AuthenticatedKey, + state: ProxyState, + request: Request, + scope: Option, ) -> Response { // #698: /mcp emits the same access log + request metrics as every other // handler — pre-fix the endpoint was invisible in both. One wrapper @@ -77,13 +104,19 @@ pub async fn mcp_endpoint( let api_key_id = auth.entry.id.clone(); let method = request.method().clone(); - let response = dispatch(auth, &state, request, &request_id).await; + let response = dispatch(auth, scope.as_deref(), &state, request, &request_id).await; let elapsed = started.elapsed(); let status = response.status().as_u16(); AccessLog { method: method.as_str(), - path: "/mcp", + // Bounded route template, mirroring `/a2a` (the per-request server is + // on the usage event, not the access log). + path: if scope.is_some() { + "/mcp/{server}" + } else { + "/mcp" + }, status, latency: elapsed, provider: Some("mcp"), @@ -115,10 +148,33 @@ pub async fn mcp_endpoint( async fn dispatch( auth: AuthenticatedKey, + scope: Option<&str>, state: &ProxyState, request: Request, request_id: &str, ) -> Response { + // One snapshot for the whole request: the scoped-server resolution below + // and the gateway construction further down must see the same resource + // set. + let snapshot = state.snapshot.load(); + + // Scoped endpoint: resolve the path's server before doing any work. A + // disabled server is treated as absent — not served, same as the + // aggregated endpoint skipping it (and same as `/a2a/:agent`). + if let Some(server) = scope { + let known = snapshot + .mcp_servers + .get_by_name(server) + .is_some_and(|entry| entry.value.enabled); + if !known { + return ( + StatusCode::NOT_FOUND, + format!("unknown MCP server: {server}"), + ) + .into_response(); + } + } + // Buffer the body so the JSON-RPC method can be inspected, then rebuilt for // the gateway. The global body-limit layer has already capped the size. let (parts, body) = request.into_parts(); @@ -143,15 +199,31 @@ async fn dispatch( let peek = serde_json::from_slice::(&bytes).ok(); let is_tool_call = peek.as_ref().and_then(|p| p.method.as_deref()) == Some("tools/call"); - // Split the namespaced tool name into (server, tool) up front, owned, so it - // survives the body being consumed when the request is rebuilt. + // Resolve the called (server, tool) up front, owned, so it survives the + // body being consumed when the request is rebuilt. Aggregated: split the + // namespaced name. Scoped: the server comes from the path and the name is + // the upstream's original one — but strip the scope prefix if the caller + // sent the namespaced form anyway, mirroring the gateway's own acceptance + // of both, so quota and usage attribute the same tool either way. let (mcp_server, mcp_tool) = if is_tool_call { - peek.as_ref() + let name = peek + .as_ref() .and_then(|p| p.params.as_ref()) .and_then(|p| p.name.as_deref()) - .and_then(|name| name.split_once(aisix_mcp::TOOL_NAMESPACE_SEPARATOR)) - .map(|(server, tool)| (server.to_string(), tool.to_string())) - .unwrap_or_default() + .unwrap_or_default(); + match scope { + Some(server) => { + let bare = name + .strip_prefix(server) + .and_then(|rest| rest.strip_prefix(aisix_mcp::TOOL_NAMESPACE_SEPARATOR)) + .unwrap_or(name); + (server.to_string(), bare.to_string()) + } + None => name + .split_once(aisix_mcp::TOOL_NAMESPACE_SEPARATOR) + .map(|(server, tool)| (server.to_string(), tool.to_string())) + .unwrap_or_default(), + } } else { (String::new(), String::new()) }; @@ -243,12 +315,25 @@ async fn dispatch( } } - let snapshot = state.snapshot.load(); // Scope the gateway to the tools this caller's key permits — resolved // from the key together with the environment/team MCP access policies — // so MCP tool access is governed by the same key object as LLM access. let acl = aisix_mcp::ToolAcl::resolve(&snapshot, auth.key()); - let gateway = aisix_mcp::McpGateway::from_snapshot(&snapshot).with_tool_acl(acl); + let gateway = match scope { + // Same snapshot as the resolution above, so the entry is still there. + Some(server) => match aisix_mcp::McpGateway::from_snapshot_scoped(&snapshot, server) { + Some(gateway) => gateway, + None => { + return ( + StatusCode::NOT_FOUND, + format!("unknown MCP server: {server}"), + ) + .into_response() + } + }, + None => aisix_mcp::McpGateway::from_snapshot(&snapshot), + } + .with_tool_acl(acl); let service = aisix_mcp::streamable_http_service(gateway); let request = Request::from_parts(parts, Body::from(bytes)); // `StreamableHttpService` is a tower service that dispatches on method and @@ -1318,4 +1403,170 @@ mod tests { tokio::time::sleep(std::time::Duration::from_millis(25)).await; } } + + // ── /mcp/{server} — the scoped, single-server endpoint ── + + /// Seed an MCP server row. The URL is unroutable on purpose: these tests + /// pin routing, resolution and attribution, not upstream success (that is + /// covered by the aisix-mcp integration tests and the e2e suite). + fn insert_mcp_server(snapshot: &AisixSnapshot, id: &str, name: &str, enabled: bool) { + let server: aisix_core::McpServer = serde_json::from_value(serde_json::json!({ + "display_name": name, + "url": "http://127.0.0.1:9/mcp", + "enabled": enabled, + })) + .expect("valid mcp server"); + snapshot + .mcp_servers + .insert(ResourceEntry::new(id, server, 1)); + } + + /// A JSON-RPC request to `/mcp/{server}`, optionally authenticated. + fn scoped_request( + server: &str, + auth: Option<&str>, + method: &str, + params: serde_json::Value, + ) -> HttpRequest { + let body = serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "method": method, + "params": params + }); + let mut builder = HttpRequest::post(format!("/mcp/{server}")) + .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() + } + + fn scoped_tools_call(server: &str, name: &str) -> HttpRequest { + scoped_request( + server, + Some(TOKEN), + "tools/call", + serde_json::json!({ "name": name, "arguments": {} }), + ) + } + + /// A snapshot with one enabled server `alpha` and a key that may call + /// every tool. + fn scoped_snapshot() -> AisixSnapshot { + let apikey: ApiKey = serde_json::from_value(serde_json::json!({ + "key_hash": ApiKey::hash_bearer(TOKEN), + "allowed_models": ["*"], + "allowed_tools": ["*"], + })) + .expect("valid apikey"); + let snapshot = AisixSnapshot::new(); + snapshot + .apikeys + .insert(ResourceEntry::new("ak-1", apikey, 1)); + insert_mcp_server(&snapshot, "mcp-1", "alpha", true); + snapshot + } + + #[tokio::test] + async fn scoped_endpoint_auth_precedes_server_resolution() { + // No credentials → 401, even for a server that does not exist: an + // unauthenticated caller must not learn which servers are registered. + let router = router_with(snapshot_with_key()); + let response = router + .oneshot(scoped_request( + "ghost", + None, + "initialize", + serde_json::json!({}), + )) + .await + .expect("router responds"); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + } + + #[tokio::test] + async fn scoped_endpoint_unknown_or_disabled_server_is_404() { + // Unregistered server. + let router = router_with(scoped_snapshot()); + let response = router + .oneshot(scoped_tools_call("ghost", "tool")) + .await + .expect("router responds"); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + + // Registered but disabled server: treated as absent, no fallback to + // the aggregated surface. + let snapshot = scoped_snapshot(); + insert_mcp_server(&snapshot, "mcp-2", "dark", false); + let router = router_with(snapshot); + let response = router + .oneshot(scoped_tools_call("dark", "tool")) + .await + .expect("router responds"); + assert_eq!(response.status(), StatusCode::NOT_FOUND); + } + + #[tokio::test] + async fn scoped_tool_call_attributes_server_from_path() { + use aisix_obs::{UsageEvent, UsageSink}; + + let (tx, mut rx) = tokio::sync::mpsc::channel::(8); + let handle = SnapshotHandle::new(scoped_snapshot()); + let hub = Arc::new(aisix_gateway::Hub::new()); + let state = ProxyState::new(handle, hub, &cfg()) + .without_cache() + .with_usage_sink(UsageSink::new(tx)); + let router = build_router(state); + + // A bare (original) tool name: the server comes from the path, not + // from a namespace prefix inside the name. + let _ = router + .clone() + .oneshot(scoped_tools_call("alpha", "tool")) + .await + .expect("router responds"); + let event = rx.try_recv().expect("usage event for the bare-name call"); + assert_eq!(event.inbound_protocol, "mcp"); + assert_eq!(event.mcp_server_name, "alpha"); + assert_eq!(event.mcp_tool_name, "tool"); + + // The namespaced spelling attributes identically — same tool, same + // server — so quota and usage cannot be split by client spelling. + let _ = router + .oneshot(scoped_tools_call("alpha", "alpha__tool")) + .await + .expect("router responds"); + let event = rx.try_recv().expect("usage event for the prefixed call"); + assert_eq!(event.mcp_server_name, "alpha"); + assert_eq!(event.mcp_tool_name, "tool"); + } + + #[tokio::test] + async fn scoped_per_server_rate_limit_keys_on_path_server() { + // The key may call `alpha` once a minute. On the scoped endpoint the + // limit must key on the path's server even for bare tool names. + let snapshot = snapshot_with_mcp_server_limits(&[( + "ak-1", + TOKEN, + serde_json::json!({ "alpha": { "rpm": 1 } }), + )]); + insert_mcp_server(&snapshot, "mcp-1", "alpha", true); + let router = router_with(snapshot); + + let first = router + .clone() + .oneshot(scoped_tools_call("alpha", "tool")) + .await + .expect("router responds"); + assert_eq!(first.status(), StatusCode::OK); + + let second = router + .oneshot(scoped_tools_call("alpha", "tool")) + .await + .expect("router responds"); + assert_eq!(second.status(), StatusCode::TOO_MANY_REQUESTS); + } } diff --git a/tests/e2e/src/cases/mcp-scoped-endpoint-e2e.test.ts b/tests/e2e/src/cases/mcp-scoped-endpoint-e2e.test.ts new file mode 100644 index 000000000..30743c6d5 --- /dev/null +++ b/tests/e2e/src/cases/mcp-scoped-endpoint-e2e.test.ts @@ -0,0 +1,302 @@ +import { createHash, randomUUID } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + spawnApp, + startMcpUpstream, + waitConfigPropagation, + type McpUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: the scoped single-server MCP endpoint `/mcp/{server}` against a real +// gateway + etcd + two real MCP upstreams (official TypeScript SDK servers). +// +// Pinned contract: +// - `initialize` reports the scoped server's registered name; +// - `tools/list` returns the upstream's ORIGINAL tool names (no `server__` +// prefix), filtered by the caller's ACL evaluated in namespaced form; +// - `tools/call` accepts both the bare and the namespaced spelling; the +// path — not the tool name — picks the server, so the same bare name +// reaches different servers on different URLs; +// - a foreign `other__` prefix stays pinned to the path's server; +// - unknown and disabled servers are 404 (no fallback to the aggregate); +// - the aggregated `/mcp` surface is unchanged (still namespaced). + +const KEY_WILD = "sk-scoped-wild"; +const KEY_ALPHA_ECHO = "sk-scoped-alpha-echo"; +const KEY_BETA_ONLY = "sk-scoped-beta-only"; + +const sha256 = (s: string) => createHash("sha256").update(s).digest("hex"); + +interface RpcReply { + status: number; + json?: { + result?: { + serverInfo?: { name?: string }; + tools?: Array<{ name: string }>; + content?: Array<{ type: string; text?: string }>; + isError?: boolean; + }; + error?: { code: number; message: string }; + }; +} + +describe("mcp scoped endpoint e2e: /mcp/{server}", () => { + let app: SpawnedApp | undefined; + let alpha: McpUpstream | undefined; + let beta: McpUpstream | undefined; + let etcdReachable = false; + let seed: SeedClient; + + const post = async ( + path: string, + token: string, + body: unknown, + ): Promise => { + const res = await fetch(`${app!.proxyUrl}${path}`, { + method: "POST", + headers: { + authorization: `Bearer ${token}`, + "content-type": "application/json", + accept: "application/json, text/event-stream", + }, + body: JSON.stringify(body), + }); + const text = await res.text(); + let json: RpcReply["json"]; + try { + json = text ? JSON.parse(text) : undefined; + } catch { + json = undefined; + } + return { status: res.status, json }; + }; + + /** Spec-faithful per-operation handshake (the endpoint is stateless). */ + const initialize = async ( + path: string, + token: string, + ): Promise<{ status: number; serverName?: string }> => { + const init = await post(path, token, { + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "mcp-scoped-e2e", version: "0.1" }, + }, + }); + if (init.status !== 200) return { status: init.status }; + await post(path, token, { + jsonrpc: "2.0", + method: "notifications/initialized", + }); + return { + status: init.status, + serverName: init.json?.result?.serverInfo?.name, + }; + }; + + /** Sorted tool names visible to `token` at `path`, or an HTTP status. */ + const listToolNames = async ( + path: string, + token: string, + ): Promise<{ status: number; names?: string[] }> => { + const { status } = await initialize(path, token); + if (status !== 200) return { status }; + const r = await post(path, token, { + jsonrpc: "2.0", + id: 2, + method: "tools/list", + params: {}, + }); + const tools = r.json?.result?.tools; + if (r.status !== 200 || !tools) return { status: r.status }; + return { status: r.status, names: tools.map((t) => t.name).sort() }; + }; + + const callTool = async ( + path: string, + token: string, + name: string, + text: string, + ): Promise<{ ok: boolean; text?: string; error?: string }> => { + await initialize(path, token); + const r = await post(path, token, { + jsonrpc: "2.0", + id: 3, + method: "tools/call", + params: { name, arguments: { text } }, + }); + if (r.json?.error) return { ok: false, error: r.json.error.message }; + const result = r.json?.result; + if (!result || result.isError) { + return { ok: false, error: JSON.stringify(r.json ?? r.status) }; + } + return { ok: true, text: result.content?.[0]?.text }; + }; + + /** True once `token` lists exactly `names` at `path` — propagation probe. */ + const listMatches = async ( + path: string, + token: string, + names: string[], + ): Promise => { + const listed = await listToolNames(path, token); + return ( + listed.status === 200 && + JSON.stringify(listed.names) === JSON.stringify(names) + ); + }; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + alpha = await startMcpUpstream("alpha"); + beta = await startMcpUpstream("beta"); + app = await spawnApp(); + seed = new SeedClient(etcd, app.etcdPrefix); + + await seed.update("mcp_servers", randomUUID(), { + display_name: "alpha", + url: alpha.url, + enabled: true, + }); + await seed.update("mcp_servers", randomUUID(), { + display_name: "beta", + url: beta.url, + enabled: true, + }); + await seed.update("mcp_servers", randomUUID(), { + display_name: "dark", + url: alpha.url, + enabled: false, + }); + + const keyDoc = ( + plaintext: string, + allowedTools: string[], + ): Record => ({ + key_hash: sha256(plaintext), + allowed_models: [], + allowed_tools: allowedTools, + }); + await seed.createApiKey(keyDoc(KEY_WILD, ["*"])); + await seed.createApiKey(keyDoc(KEY_ALPHA_ECHO, ["alpha__echo"])); + await seed.createApiKey(keyDoc(KEY_BETA_ONLY, ["beta__*"])); + + // Probe every key to its steady state so no later assertion races a row + // that has not landed yet. + await waitConfigPropagation(async () => { + if (!(await listMatches("/mcp/alpha", KEY_WILD, ["echo", "reverse"]))) { + return false; + } + if (!(await listMatches("/mcp/alpha", KEY_ALPHA_ECHO, ["echo"]))) { + return false; + } + return listMatches("/mcp/beta", KEY_BETA_ONLY, ["echo", "reverse"]); + }); + }, 60_000); + + afterAll(async () => { + await app?.exit(); + await alpha?.close(); + await beta?.close(); + }); + + test("initialize presents the scoped server, not the aggregate", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + const init = await initialize("/mcp/alpha", KEY_WILD); + expect(init.status).toBe(200); + expect(init.serverName).toBe("alpha"); + }); + + test("tools/list returns original names; the aggregate stays namespaced", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + const scoped = await listToolNames("/mcp/alpha", KEY_WILD); + expect(scoped.status).toBe(200); + expect(scoped.names).toEqual(["echo", "reverse"]); + + // Regression: the aggregated endpoint is untouched by the scoped surface. + const aggregated = await listToolNames("/mcp", KEY_WILD); + expect(aggregated.status).toBe(200); + expect(aggregated.names).toEqual([ + "alpha__echo", + "alpha__reverse", + "beta__echo", + "beta__reverse", + ]); + }); + + test("the path picks the server: the same bare name reaches different upstreams", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + const viaAlpha = await callTool("/mcp/alpha", KEY_WILD, "echo", "hi"); + expect(viaAlpha).toEqual({ ok: true, text: "alpha:hi" }); + + const viaBeta = await callTool("/mcp/beta", KEY_WILD, "echo", "hi"); + expect(viaBeta).toEqual({ ok: true, text: "beta:hi" }); + + // The namespaced spelling keeps working on the scoped URL. + const namespaced = await callTool( + "/mcp/alpha", + KEY_WILD, + "alpha__echo", + "hi", + ); + expect(namespaced).toEqual({ ok: true, text: "alpha:hi" }); + }); + + test("a registered foreign prefix fails closed; an unregistered one is a bare name", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + // `beta` is a registered server, so `beta__echo` on /mcp/alpha is a + // cross-server mistake: rejected, never routed to beta and never served + // bare by alpha (the harness upstream would answer any name). + const crossed = await callTool("/mcp/alpha", KEY_WILD, "beta__echo", "hi"); + expect(crossed.ok).toBe(false); + expect(crossed.error).toContain("not available"); + + // `ghost` is not registered, so `ghost__echo` is just a tool name that + // contains the separator — it reaches alpha verbatim (the harness + // upstream echoes for any non-`reverse` name, labelling the answer). + const bare = await callTool("/mcp/alpha", KEY_WILD, "ghost__echo", "hi"); + expect(bare).toEqual({ ok: true, text: "alpha:hi" }); + }); + + test("ACL keeps its namespaced meaning on the scoped endpoint", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + // A grant written as `alpha__echo` admits the bare `echo` here… + const listed = await listToolNames("/mcp/alpha", KEY_ALPHA_ECHO); + expect(listed.names).toEqual(["echo"]); + const ok = await callTool("/mcp/alpha", KEY_ALPHA_ECHO, "echo", "hi"); + expect(ok).toEqual({ ok: true, text: "alpha:hi" }); + const denied = await callTool("/mcp/alpha", KEY_ALPHA_ECHO, "reverse", "hi"); + expect(denied.ok).toBe(false); + expect(denied.error).toContain("not available"); + + // …and a key scoped to beta reaches nothing on alpha's URL. + const foreign = await listToolNames("/mcp/alpha", KEY_BETA_ONLY); + expect(foreign.names).toEqual([]); + const rejected = await callTool("/mcp/alpha", KEY_BETA_ONLY, "echo", "hi"); + expect(rejected.ok).toBe(false); + }); + + test("unknown and disabled servers are 404 with no fallback", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + const unknown = await initialize("/mcp/ghost", KEY_WILD); + expect(unknown.status).toBe(404); + + const disabled = await initialize("/mcp/dark", KEY_WILD); + expect(disabled.status).toBe(404); + }); +}); From fd95de08940d2d50a769b8d775175e62fe56e37e Mon Sep 17 00:00:00 2001 From: Jarvis Date: Tue, 4 Aug 2026 17:21:54 +0800 Subject: [PATCH 2/5] fix(mcp): make scoped tools/list round-trip and share one name parser MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Independent review findings on the scoped endpoint: - tools/list stripped the namespace prefix unconditionally, so an upstream tool whose literal name starts with a registered server's prefix was advertised under a spelling tools/call would re-strip (mis-dispatch) or fail closed on. Strip only when the bare name round-trips; colliding names stay namespaced — that spelling is the callable one. - The proxy's attribution peek parsed the name with its own logic, which diverged from the gateway for server names ending in '_'. Both sides now share one primitive (strip_server_prefix), which is whole-string-prefix based, so such names namespace cleanly too. - Pin the remaining review gaps in tests: the namespaced spelling and the aggregated endpoint share the per-server rate-limit bucket, and /mcp/{server} vs /mcp/ metric labels can't regress by arm reorder. --- crates/aisix-mcp/src/gateway.rs | 54 +++++++++++++++------ crates/aisix-mcp/src/lib.rs | 4 +- crates/aisix-mcp/tests/gateway_scoped.rs | 60 ++++++++++++++++++++---- crates/aisix-proxy/src/lib.rs | 8 ++++ crates/aisix-proxy/src/mcp.rs | 48 ++++++++++++++++--- 5 files changed, 143 insertions(+), 31 deletions(-) diff --git a/crates/aisix-mcp/src/gateway.rs b/crates/aisix-mcp/src/gateway.rs index 080a49da4..c5effd8d1 100644 --- a/crates/aisix-mcp/src/gateway.rs +++ b/crates/aisix-mcp/src/gateway.rs @@ -59,6 +59,17 @@ struct NamedUpstream { bridge: Arc, } +/// Strip `server`'s namespace prefix from `name`: `Some(bare)` when `name` +/// is `__`, `None` otherwise. The one primitive both the +/// scoped gateway and the proxy's attribution peek use to interpret a tool +/// name on `/mcp/{server}`, so the two can never drift. Prefix matching is +/// by whole-string prefix (not first-separator split), so a server name +/// that itself ends in `_` still namespaces cleanly. +pub fn strip_server_prefix<'a>(server: &str, name: &'a str) -> Option<&'a str> { + name.strip_prefix(server) + .and_then(|rest| rest.strip_prefix(TOOL_NAMESPACE_SEPARATOR)) +} + /// Which tools a gateway instance may expose and call, in the namespaced /// `__` form. Built per request from the caller's API key and /// the environment's / the key's team's MCP access policies, so MCP tool @@ -402,12 +413,22 @@ impl ServerHandler for McpGateway { // A scoped gateway serves its single upstream's tools under their // original names — the namespace prefix exists to disambiguate the // aggregate, and a single-server endpoint has nothing to disambiguate. - // ACL filtering above ran on the namespaced form. + // ACL filtering above ran on the namespaced form. Strip only when the + // bare name round-trips through `call_tool`'s parsing: a literal + // upstream name that itself starts with a registered server's prefix + // would be re-stripped (or fail closed) if advertised bare, so those + // stay namespaced — that spelling is the one `call_tool` accepts. if let Some(scoped) = &self.scoped { - let prefix = format!("{}{TOOL_NAMESPACE_SEPARATOR}", scoped.name); for tool in &mut tools { - if let Some(bare) = tool.name.strip_prefix(&prefix) { - tool.name = Cow::Owned(bare.to_string()); + if let Some(bare) = strip_server_prefix(&scoped.name, tool.name.as_ref()) { + let round_trips = strip_server_prefix(&scoped.name, bare).is_none() + && !scoped + .foreign + .iter() + .any(|f| strip_server_prefix(f, bare).is_some()); + if round_trips { + tool.name = Cow::Owned(bare.to_string()); + } } } } @@ -424,19 +445,24 @@ impl ServerHandler for McpGateway { // sends the namespaced form anyway (an aggregated-endpoint client // pointed at the scoped URL) keeps working — `__x` reduces to // `x`. (An upstream tool literally named `__x` is therefore - // only callable as `____x` — prefix-stripping takes - // precedence, pinned by test.) A name prefixed with a *different* - // registered server's name fails closed: the scoped endpoint never - // cross-routes, and silently serving it as a bare name would mask the - // caller's mistake. An unregistered prefix stays a bare name, since - // tool names may legitimately contain the separator. + // only callable as `____x` — one strip, own prefix + // first; `tools/list` advertises exactly that spelling.) A name + // prefixed with a *different* registered server's name fails closed: + // the scoped endpoint never cross-routes, and silently serving it as + // a bare name would mask the caller's mistake. An unregistered + // prefix stays a bare name, since tool names may legitimately + // contain the separator. let request_name = request.name.as_ref(); let (server, tool, namespaced): (&str, &str, Cow<'_, str>) = match &self.scoped { Some(scoped) => { let scope = scoped.name.as_str(); - let bare = match request_name.split_once(TOOL_NAMESPACE_SEPARATOR) { - Some((prefix, rest)) if prefix == scope => rest, - Some((prefix, _)) if scoped.foreign.contains(prefix) => { + let bare = match strip_server_prefix(scope, request_name) { + Some(rest) => rest, + None if scoped + .foreign + .iter() + .any(|f| strip_server_prefix(f, request_name).is_some()) => + { // Same neutral wording as the ACL reject: don't // confirm what the other server serves. return Err(ErrorData::invalid_params( @@ -444,7 +470,7 @@ impl ServerHandler for McpGateway { None, )); } - _ => request_name, + None => request_name, }; ( scope, diff --git a/crates/aisix-mcp/src/lib.rs b/crates/aisix-mcp/src/lib.rs index 09f15952c..736a561dc 100644 --- a/crates/aisix-mcp/src/lib.rs +++ b/crates/aisix-mcp/src/lib.rs @@ -24,5 +24,7 @@ pub use bridge::{ McpUpstream, OAuthClientConfig, RmcpBridge, }; pub use error::McpError; -pub use gateway::{streamable_http_service, McpGateway, ToolAcl, TOOL_NAMESPACE_SEPARATOR}; +pub use gateway::{ + streamable_http_service, strip_server_prefix, McpGateway, ToolAcl, TOOL_NAMESPACE_SEPARATOR, +}; pub use openapi::{validate_spec, OpenApiBridge}; diff --git a/crates/aisix-mcp/tests/gateway_scoped.rs b/crates/aisix-mcp/tests/gateway_scoped.rs index 264d765fe..f95978837 100644 --- a/crates/aisix-mcp/tests/gateway_scoped.rs +++ b/crates/aisix-mcp/tests/gateway_scoped.rs @@ -249,13 +249,19 @@ async fn scoped_registered_foreign_prefix_fails_closed() { McpGateway::from_snapshot_scoped(&snapshot, "alpha").expect("alpha is registered"); let client = connect(spawn_gateway(gateway).await).await; + // tools/list keeps the colliding literal name namespaced: advertising + // the bare `beta__echo` would advertise a spelling `tools/call` rejects. + let tools = client.list_all_tools().await.expect("list tools"); + let names: Vec<&str> = tools.iter().map(|t| t.name.as_ref()).collect(); + assert_eq!(names, vec!["alpha__beta__echo"], "list must round-trip"); + assert!( client.call_tool(call("beta__echo", "x")).await.is_err(), "a registered foreign prefix must fail closed, never route or serve" ); - // The shadowed literal tool stays reachable through its namespaced - // spelling: `alpha__beta__echo` strips alpha's own prefix first. + // The advertised spelling reaches the literal tool: `alpha__beta__echo` + // strips alpha's own prefix first. let escaped = client .call_tool(call("alpha__beta__echo", "hi")) .await @@ -263,6 +269,37 @@ async fn scoped_registered_foreign_prefix_fails_closed() { assert_eq!(first_text(&escaped), "alpha:hi"); } +#[tokio::test] +async fn scoped_server_name_ending_in_underscore_namespaces_cleanly() { + // `data_` is a legal server name (only `__` inside a name is rejected). + // Prefix parsing is whole-string based, not first-separator based, so + // the namespaced spelling `data___query` (= `data_` + `__` + `query`) + // resolves to `query` even while a server named `data` also exists. + let upstream = spawn_upstream("data_", "query").await; + let snapshot = AisixSnapshot::new(); + snapshot + .mcp_servers + .insert(mcp_entry("e1", "data_", &upstream, true)); + snapshot + .mcp_servers + .insert(mcp_entry("e2", "data", &upstream, true)); + let gateway = + McpGateway::from_snapshot_scoped(&snapshot, "data_").expect("data_ is registered"); + let client = connect(spawn_gateway(gateway).await).await; + + let tools = client.list_all_tools().await.expect("list tools"); + let names: Vec<&str> = tools.iter().map(|t| t.name.as_ref()).collect(); + assert_eq!(names, vec!["query"]); + + let bare = client.call_tool(call("query", "hi")).await.expect("bare"); + assert_eq!(first_text(&bare), "data_:hi"); + let namespaced = client + .call_tool(call("data___query", "hi")) + .await + .expect("namespaced spelling of a trailing-underscore server"); + assert_eq!(first_text(&namespaced), "data_:hi"); +} + #[tokio::test] async fn scoped_unregistered_prefix_is_a_bare_name() { // `ghost` is not a registered server, so `ghost__echo` is just a tool @@ -284,9 +321,10 @@ async fn scoped_unregistered_prefix_is_a_bare_name() { #[tokio::test] async fn scoped_upstream_tool_spelled_like_the_prefix_keeps_precedence() { // Pathological upstream: a tool literally named `alpha__echo` on server - // `alpha`. Prefix-stripping takes precedence (documented in `call_tool`): - // the tool lists as `alpha__echo` and is callable as `alpha__alpha__echo`; - // the spelling `alpha__echo` reduces to `echo`, which does not exist. + // `alpha`. Prefix-stripping takes precedence (documented in `call_tool`), + // so the bare spelling would not round-trip — tools/list therefore keeps + // the namespaced `alpha__alpha__echo`, which IS the callable spelling; + // `alpha__echo` reduces to `echo`, which does not exist. let upstream = spawn_upstream("alpha", "alpha__echo").await; let snapshot = snapshot_with_alpha(&upstream); let gateway = @@ -295,13 +333,17 @@ async fn scoped_upstream_tool_spelled_like_the_prefix_keeps_precedence() { let tools = client.list_all_tools().await.expect("list tools"); let names: Vec<&str> = tools.iter().map(|t| t.name.as_ref()).collect(); - assert_eq!(names, vec!["alpha__echo"], "one strip only: {names:?}"); + assert_eq!( + names, + vec!["alpha__alpha__echo"], + "a non-round-tripping name stays namespaced: {names:?}" + ); - let double = client + let advertised = client .call_tool(call("alpha__alpha__echo", "hi")) .await - .expect("double-prefixed spelling reaches the literal tool"); - assert_eq!(first_text(&double), "alpha:hi"); + .expect("the advertised spelling reaches the literal tool"); + assert_eq!(first_text(&advertised), "alpha:hi"); assert!( client.call_tool(call("alpha__echo", "x")).await.is_err(), diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 1e42521d1..c683c9e9e 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -622,6 +622,14 @@ mod tests { // Arbitrary unauthenticated paths bucket to a single label. assert_eq!(normalize_endpoint_label("/random/x"), "other"); assert_eq!(normalize_endpoint_label("/random/y"), "other"); + // The scoped MCP endpoint collapses to one label regardless of the + // server segment; the aggregated endpoint (with and without the + // trailing slash) keeps its own — the exact `"/mcp/"` arm must stay + // ahead of the `/mcp/` prefix arm. + assert_eq!(normalize_endpoint_label("/mcp/alpha"), "/mcp/{server}"); + assert_eq!(normalize_endpoint_label("/mcp/unique-xyz"), "/mcp/{server}"); + assert_eq!(normalize_endpoint_label("/mcp"), "/mcp"); + assert_eq!(normalize_endpoint_label("/mcp/"), "/mcp"); } use aisix_core::resource::ResourceEntry; diff --git a/crates/aisix-proxy/src/mcp.rs b/crates/aisix-proxy/src/mcp.rs index 3290d444e..460117363 100644 --- a/crates/aisix-proxy/src/mcp.rs +++ b/crates/aisix-proxy/src/mcp.rs @@ -202,9 +202,10 @@ async fn dispatch( // Resolve the called (server, tool) up front, owned, so it survives the // body being consumed when the request is rebuilt. Aggregated: split the // namespaced name. Scoped: the server comes from the path and the name is - // the upstream's original one — but strip the scope prefix if the caller - // sent the namespaced form anyway, mirroring the gateway's own acceptance - // of both, so quota and usage attribute the same tool either way. + // the upstream's original one — stripped through the same primitive the + // gateway parses with (`strip_server_prefix`), so quota and usage + // attribute the same tool for both spellings and can never drift from + // what actually dispatches. let (mcp_server, mcp_tool) = if is_tool_call { let name = peek .as_ref() @@ -213,10 +214,7 @@ async fn dispatch( .unwrap_or_default(); match scope { Some(server) => { - let bare = name - .strip_prefix(server) - .and_then(|rest| rest.strip_prefix(aisix_mcp::TOOL_NAMESPACE_SEPARATOR)) - .unwrap_or(name); + let bare = aisix_mcp::strip_server_prefix(server, name).unwrap_or(name); (server.to_string(), bare.to_string()) } None => name @@ -1564,9 +1562,45 @@ mod tests { assert_eq!(first.status(), StatusCode::OK); let second = router + .clone() .oneshot(scoped_tools_call("alpha", "tool")) .await .expect("router responds"); assert_eq!(second.status(), StatusCode::TOO_MANY_REQUESTS); + + // The namespaced spelling shares the SAME bucket — a client cannot + // double its per-server allowance by switching spellings. + let prefixed = router + .oneshot(scoped_tools_call("alpha", "alpha__tool")) + .await + .expect("router responds"); + assert_eq!(prefixed.status(), StatusCode::TOO_MANY_REQUESTS); + } + + #[tokio::test] + async fn scoped_and_aggregated_endpoints_share_the_per_server_bucket() { + // rpm=1 for `alpha`: one call through the aggregated endpoint must + // exhaust the allowance for the scoped endpoint too — switching + // endpoints cannot double the limit. + let snapshot = snapshot_with_mcp_server_limits(&[( + "ak-1", + TOKEN, + serde_json::json!({ "alpha": { "rpm": 1 } }), + )]); + insert_mcp_server(&snapshot, "mcp-1", "alpha", true); + let router = router_with(snapshot); + + let aggregated = router + .clone() + .oneshot(tools_call_on(TOKEN, "alpha")) + .await + .expect("router responds"); + assert_eq!(aggregated.status(), StatusCode::OK); + + let scoped = router + .oneshot(scoped_tools_call("alpha", "tool")) + .await + .expect("router responds"); + assert_eq!(scoped.status(), StatusCode::TOO_MANY_REQUESTS); } } From 37bf39b246c902ff5bca321a7b9bd82377313d35 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Tue, 4 Aug 2026 17:22:38 +0800 Subject: [PATCH 3/5] feat(proxy): entry-level URL rewriting via proxy.url_rewrites MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An ordered list of {match, rewrite} regex rules applied to every proxy-listener request before route matching (admin/metrics listeners unaffected): the first matching rule rewrites the path once — no cascading — and the request then flows through the normal endpoint (auth, ACL, quota, metrics labelling) as if the client had sent the rewritten path. Replacement substitutes the matched portion with $1/${name} capture-group expansion; the query string is preserved; a miss leaves the request untouched. Invalid regexes fail startup. Because axum's Router::layer middleware runs after route matching, the rewrite gets its pre-routing seat by wrapping the whole router as the fallback of an outer router; the wrapper is only built when rules are configured, so the default path pays nothing. Lets operators map legacy URL shapes onto AISIX endpoints without client changes — e.g. per-server MCP paths like /mcp-servers/{svc}/mcp onto the /mcp/{server} endpoint, completing the migration scenario of api7/AISIX-Cloud#1219 together with the scoped-endpoint PR. --- Cargo.lock | 2 + config.example.yaml | 12 + config.managed.yaml | 8 + crates/aisix-admin/src/playground_handler.rs | 1 + crates/aisix-core/Cargo.toml | 1 + crates/aisix-core/src/config.rs | 88 +++++++ crates/aisix-core/src/lib.rs | 2 +- crates/aisix-proxy/Cargo.toml | 1 + crates/aisix-proxy/src/a2a.rs | 1 + crates/aisix-proxy/src/audio.rs | 1 + crates/aisix-proxy/src/completions.rs | 1 + crates/aisix-proxy/src/count_tokens.rs | 1 + crates/aisix-proxy/src/embeddings.rs | 1 + crates/aisix-proxy/src/images.rs | 1 + crates/aisix-proxy/src/jobs.rs | 1 + crates/aisix-proxy/src/lib.rs | 26 ++- crates/aisix-proxy/src/mcp.rs | 1 + crates/aisix-proxy/src/messages.rs | 1 + crates/aisix-proxy/src/models.rs | 1 + crates/aisix-proxy/src/passthrough.rs | 1 + crates/aisix-proxy/src/realtime.rs | 1 + crates/aisix-proxy/src/rerank.rs | 1 + crates/aisix-proxy/src/responses.rs | 1 + crates/aisix-proxy/src/rewrite.rs | 227 +++++++++++++++++++ crates/aisix-proxy/src/state.rs | 7 + crates/aisix-proxy/src/videos.rs | 1 + tests/e2e/src/cases/url-rewrite-e2e.test.ts | 219 ++++++++++++++++++ tests/e2e/src/harness/app.ts | 8 + 28 files changed, 614 insertions(+), 3 deletions(-) create mode 100644 crates/aisix-proxy/src/rewrite.rs create mode 100644 tests/e2e/src/cases/url-rewrite-e2e.test.ts diff --git a/Cargo.lock b/Cargo.lock index b89e4d40a..b6fec76c6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -105,6 +105,7 @@ dependencies = [ "ipnet", "jsonschema", "once_cell", + "regex", "schemars 0.8.22", "serde", "serde_json", @@ -373,6 +374,7 @@ dependencies = [ "jsonwebtoken", "lofty", "rand 0.8.5", + "regex", "reqwest 0.12.28", "serde", "serde_json", diff --git a/config.example.yaml b/config.example.yaml index c66ca767b..5f5b1c6b7 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -51,6 +51,18 @@ proxy: # trusted_proxies: ["10.0.0.0/8", "127.0.0.1/32"] # recursive: true # header: x-forwarded-for + # Entry-level URL rewriting: map legacy URL shapes onto AISIX endpoints + # before routing. The first rule whose `match` regex matches the request + # path rewrites it (once, no cascading); the request then flows through + # the normal endpoint — auth, ACL, quota — as if the client had sent the + # rewritten path. `rewrite` replaces the matched portion of the path + # ($1/${name} expand capture groups); the query string is preserved. An + # invalid regex fails startup. Example: serve per-server MCP URLs like + # /mcp-servers/github/mcp on the /mcp/{server} endpoint. + # url_rewrites: + # - name: per-server-mcp-compat + # match: "^/mcp-servers/([^/]+)/mcp$" + # rewrite: "/mcp/$1" admin: addr: "127.0.0.1:3001" diff --git a/config.managed.yaml b/config.managed.yaml index c3cf4debf..5e35e8770 100644 --- a/config.managed.yaml +++ b/config.managed.yaml @@ -44,6 +44,14 @@ proxy: # 0 = no request-body cap (the default); set a value to bound # per-request memory. # request_body_limit_bytes: 0 + # Entry-level URL rewriting (first matching rule wins; `rewrite` + # replaces the matched portion of the path, $1 = capture group; the + # query string is preserved). Lets clients keep legacy URL shapes, + # e.g. per-server MCP paths served on the /mcp/{server} endpoint. + # url_rewrites: + # - name: per-server-mcp-compat + # match: "^/mcp-servers/([^/]+)/mcp$" + # rewrite: "/mcp/$1" admin: # Bind to an unbindable port so even if managed mode somehow diff --git a/crates/aisix-admin/src/playground_handler.rs b/crates/aisix-admin/src/playground_handler.rs index cc0e70431..3477f6a88 100644 --- a/crates/aisix-admin/src/playground_handler.rs +++ b/crates/aisix-admin/src/playground_handler.rs @@ -79,6 +79,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-core/Cargo.toml b/crates/aisix-core/Cargo.toml index 717095b92..d513a2c1a 100644 --- a/crates/aisix-core/Cargo.toml +++ b/crates/aisix-core/Cargo.toml @@ -17,6 +17,7 @@ thiserror.workspace = true anyhow.workspace = true config.workspace = true jsonschema.workspace = true +regex.workspace = true schemars = { workspace = true, features = ["chrono"] } # ApiKey::expires_at parses RFC 3339 timestamps for key-expiry # enforcement in the proxy auth path (#933). diff --git a/crates/aisix-core/src/config.rs b/crates/aisix-core/src/config.rs index a4a973154..27a28fa8e 100644 --- a/crates/aisix-core/src/config.rs +++ b/crates/aisix-core/src/config.rs @@ -393,6 +393,16 @@ pub struct ProxyConfig { /// an L7 LB / ingress that sets `x-forwarded-for`. #[serde(default)] pub real_ip: RealIpConfig, + /// Entry-level URL rewrite rules, applied to every proxy-listener + /// request **before** routing (the admin and metrics listeners are + /// unaffected). The first rule whose `match` regex matches the request + /// path rewrites it — once, no cascading — and the request then flows + /// through the normal endpoint (auth, ACL, quota, …) as if the client + /// had sent the rewritten path. Lets operators map legacy URL shapes + /// onto AISIX endpoints, e.g. per-server MCP paths onto + /// `/mcp/{server}`. Empty (the default) = no rewriting. + #[serde(default)] + pub url_rewrites: Vec, } impl ProxyConfig { @@ -401,6 +411,26 @@ impl ProxyConfig { } } +/// One entry-level URL rewrite rule (see [`ProxyConfig::url_rewrites`]). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UrlRewriteRule { + /// Optional name, used in logs when the rule fires. + #[serde(default)] + pub name: Option, + /// Regex matched against the request path (never the query string). + /// Anchor with `^`/`$` to match the whole path; an unanchored pattern + /// matches anywhere in it. + #[serde(rename = "match")] + pub pattern: String, + /// Replacement for the matched portion of the path. Capture groups are + /// available as `$1`… / `${name}`; use `${1}x` (braced) when a literal + /// character follows a group reference. The query string is preserved + /// as sent. + #[serde(rename = "rewrite")] + pub replacement: String, +} + /// nginx `set_real_ip_from` + `real_ip_recursive` equivalent. Resolves /// the downstream client IP for usage logs (#492) from a forwarded /// header, trusting only addresses inside `trusted_proxies`. @@ -1094,6 +1124,20 @@ impl Config { } fn validate(&self) -> Result<(), BootstrapError> { + // Fail fast on a rewrite rule that can never compile: it would + // otherwise surface as every legacy-path request 404ing, which is + // much harder to trace back to a typo in one regex. + for (i, rule) in self.proxy.url_rewrites.iter().enumerate() { + if let Err(e) = regex::Regex::new(&rule.pattern) { + let ctx = rule + .name + .clone() + .unwrap_or_else(|| format!("proxy.url_rewrites[{i}]")); + return Err(BootstrapError::Config(format!( + "{ctx}: invalid match regex: {e}" + ))); + } + } if let Some(path) = self.resources_file.as_deref() { // File source selected: exactly one resource source may be // active. A configured etcd endpoint list alongside the file @@ -1306,6 +1350,50 @@ admin: assert!(nets.iter().any(|n| n.to_string() == "10.0.0.0/8")); } + #[test] + fn loads_url_rewrites_and_rejects_an_invalid_regex() { + let f = write_yaml( + r#" +etcd: + endpoints: ["http://127.0.0.1:2379"] + prefix: "/aisix" +proxy: + addr: "0.0.0.0:3000" + url_rewrites: + - name: per-server-mcp-compat + match: "^/mcp-servers/([^/]+)/mcp$" + rewrite: "/mcp/$1" +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +"#, + ); + let cfg = Config::load_from_path(Some(f.path())).unwrap(); + assert_eq!(cfg.proxy.url_rewrites.len(), 1); + assert_eq!(cfg.proxy.url_rewrites[0].replacement, "/mcp/$1"); + + let f = write_yaml( + r#" +etcd: + endpoints: ["http://127.0.0.1:2379"] + prefix: "/aisix" +proxy: + addr: "0.0.0.0:3000" + url_rewrites: + - match: "^/mcp-servers/([^/+/mcp$" + rewrite: "/mcp/$1" +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +"#, + ); + let err = Config::load_from_path(Some(f.path())).unwrap_err(); + assert!( + format!("{err}").contains("url_rewrites"), + "error should name the bad rule: {err}" + ); + } + #[test] fn rejects_malformed_trusted_proxy_cidr() { let f = write_yaml( diff --git a/crates/aisix-core/src/lib.rs b/crates/aisix-core/src/lib.rs index a379c6b87..bc3fda26d 100644 --- a/crates/aisix-core/src/lib.rs +++ b/crates/aisix-core/src/lib.rs @@ -30,7 +30,7 @@ pub mod wildcard; pub use config::{ AdminConfig, CacheBackend, CacheConfig, ClientTypeRule, Config, EtcdConfig, EtcdTlsConfig, HistogramBucketsConfig, ManagedConfig, ObservabilityConfig, ProxyConfig, RateLimitBackend, - RateLimitConfig, RealIpConfig, RedisConnConfig, RedisMode, TlsConfig, + RateLimitConfig, RealIpConfig, RedisConnConfig, RedisMode, TlsConfig, UrlRewriteRule, }; pub use config_status::{ hash_bytes, hash_entries, AppliedSnapshot, ConfigMetricsView, ConfigState, ConfigStatus, diff --git a/crates/aisix-proxy/Cargo.toml b/crates/aisix-proxy/Cargo.toml index 946a22a33..8211ca817 100644 --- a/crates/aisix-proxy/Cargo.toml +++ b/crates/aisix-proxy/Cargo.toml @@ -42,6 +42,7 @@ tower.workspace = true tower-http.workspace = true http-body-util.workspace = true bytes.workspace = true +regex.workspace = true serde.workspace = true serde_json.workspace = true # Token-estimation fallback for usage telemetry (token_estimate.rs). diff --git a/crates/aisix-proxy/src/a2a.rs b/crates/aisix-proxy/src/a2a.rs index e722b36ab..3f95aaa3a 100644 --- a/crates/aisix-proxy/src/a2a.rs +++ b/crates/aisix-proxy/src/a2a.rs @@ -375,6 +375,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index bcc6fb343..d18aa7015 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -1439,6 +1439,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 10_485_760, // 10 MB for audio real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index 5ba0fa18c..2b788835b 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -744,6 +744,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index ea0f8bf8d..28bfc1101 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/crates/aisix-proxy/src/count_tokens.rs @@ -544,6 +544,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index 738e4c828..612a6f695 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -669,6 +669,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/images.rs b/crates/aisix-proxy/src/images.rs index 5019831e4..5ec13b6b0 100644 --- a/crates/aisix-proxy/src/images.rs +++ b/crates/aisix-proxy/src/images.rs @@ -541,6 +541,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/jobs.rs b/crates/aisix-proxy/src/jobs.rs index 7519b3999..4e215ecd0 100644 --- a/crates/aisix-proxy/src/jobs.rs +++ b/crates/aisix-proxy/src/jobs.rs @@ -1802,6 +1802,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 1e42521d1..8041fc9b9 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -60,6 +60,7 @@ mod request_id; mod rerank; mod responses; mod responses_bridge; +mod rewrite; mod routing; mod semantic; pub mod sse_keepalive; @@ -100,7 +101,7 @@ static SERVER_HEADER_VALUE: std::sync::LazyLock = std::sync::LazyLo /// OpenAI-compatible proxy surface. pub fn build_router(state: ProxyState) -> Router { let body_limit = state.request_body_limit_bytes; - Router::new() + let router = Router::new() .route("/livez", get(livez)) .route("/readyz", get(readyz)) .route("/v1/models", get(models::list_models)) @@ -220,13 +221,32 @@ pub fn build_router(state: ProxyState) -> Router { header::SERVER, SERVER_HEADER_VALUE.clone(), )) + // Rewrite the request path per `proxy.url_rewrites` before route // Outermost: mint the request id into the request extensions // before any handler/extractor runs, and stamp it onto every // response (including the short-circuited 4xx from the layers // above) so the whole proxy family carries `x-aisix-request-id` // and it equals the telemetry request_id. See request_id.rs. .layer(middleware::from_fn(request_id::ensure_request_id)) - .with_state(state) + .with_state(state.clone()); + + if state.url_rewrites.is_empty() { + return router; + } + // Pre-routing URL rewriting (`proxy.url_rewrites`). `Router::layer` + // middleware runs AFTER route matching, so a URI rewritten there could + // never change which route matches. Wrapping the whole router as the + // fallback of an outer router — whose "routing" trivially resolves to + // that fallback — gives the rewrite layer a genuine pre-routing seat: + // it mutates the URI, then the inner router matches on the rewritten + // path. Built only when rules are configured, so the default path pays + // nothing. See rewrite.rs. + Router::new() + .fallback_service(router) + .layer(middleware::from_fn_with_state( + state, + rewrite::rewrite_request_uri, + )) } async fn record_in_flight_request( @@ -643,6 +663,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } @@ -1616,6 +1637,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: limit, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, }; ProxyState::new(handle, hub, &cfg).without_cache() diff --git a/crates/aisix-proxy/src/mcp.rs b/crates/aisix-proxy/src/mcp.rs index 3290d444e..3a516004d 100644 --- a/crates/aisix-proxy/src/mcp.rs +++ b/crates/aisix-proxy/src/mcp.rs @@ -548,6 +548,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index b6caeee54..f9b73e58b 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -3578,6 +3578,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/models.rs b/crates/aisix-proxy/src/models.rs index 4583add07..d1868aad4 100644 --- a/crates/aisix-proxy/src/models.rs +++ b/crates/aisix-proxy/src/models.rs @@ -119,6 +119,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/passthrough.rs b/crates/aisix-proxy/src/passthrough.rs index 805449fed..1e8221374 100644 --- a/crates/aisix-proxy/src/passthrough.rs +++ b/crates/aisix-proxy/src/passthrough.rs @@ -894,6 +894,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/realtime.rs b/crates/aisix-proxy/src/realtime.rs index 3a975feb5..1c20953a6 100644 --- a/crates/aisix-proxy/src/realtime.rs +++ b/crates/aisix-proxy/src/realtime.rs @@ -799,6 +799,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 4be1aef96..3c4128e1e 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -735,6 +735,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index be9543e6e..d0379151a 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -3065,6 +3065,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/crates/aisix-proxy/src/rewrite.rs b/crates/aisix-proxy/src/rewrite.rs new file mode 100644 index 000000000..462eccee8 --- /dev/null +++ b/crates/aisix-proxy/src/rewrite.rs @@ -0,0 +1,227 @@ +//! Entry-level URL rewriting (`proxy.url_rewrites`). +//! +//! Applied to every proxy-listener request **before** route matching (the +//! admin and metrics listeners never see this layer): the first rule whose +//! `match` regex matches the request path rewrites it — once, no cascading — +//! and the request then flows through the normal endpoint (auth, ACL, quota, +//! metrics labelling) as if the client had sent the rewritten path. A miss +//! leaves the request untouched. +//! +//! Replacement substitutes the **matched portion** of the path, with +//! `$1`/`${name}` expanding capture groups; the query string is preserved as +//! sent. Rules were syntax-checked at config load (`Config::validate`), so +//! compiling them here cannot fail on operator input. + +use std::borrow::Cow; +use std::sync::Arc; + +use axum::extract::{Request, State}; +use axum::http; +use axum::middleware::Next; +use axum::response::Response; +use regex::Regex; + +use aisix_core::UrlRewriteRule; + +use crate::state::ProxyState; + +/// One boot-compiled rewrite rule. +pub struct CompiledRewrite { + name: Option, + pattern: Regex, + replacement: String, +} + +impl CompiledRewrite { + /// The rewritten path, or `None` when the rule does not match. + fn apply(&self, path: &str) -> Option { + match self.pattern.replace(path, self.replacement.as_str()) { + // `replace` hands the haystack back untouched on a miss. + Cow::Borrowed(_) => None, + Cow::Owned(rewritten) => Some(rewritten), + } + } + + fn label(&self) -> &str { + self.name.as_deref().unwrap_or("") + } +} + +/// Compile the configured rules in declaration order. +/// +/// Invariant: every pattern was syntax-checked by `Config::validate` at +/// load, so a panic here means a caller constructed a `ProxyConfig` with an +/// unvalidated pattern. +pub fn compile(rules: &[UrlRewriteRule]) -> Arc<[CompiledRewrite]> { + rules + .iter() + .map(|rule| CompiledRewrite { + name: rule.name.clone(), + pattern: Regex::new(&rule.pattern) + .expect("proxy.url_rewrites pattern is validated at config load"), + replacement: rule.replacement.clone(), + }) + .collect() +} + +/// Middleware: apply the first matching rewrite rule to the request path. +pub async fn rewrite_request_uri( + State(state): State, + mut request: Request, + next: Next, +) -> Response { + if state.url_rewrites.is_empty() { + return next.run(request).await; + } + let fired = state.url_rewrites.iter().find_map(|rule| { + let path = request.uri().path(); + rule.apply(path).map(|to| (rule, path.to_owned(), to)) + }); + if let Some((rule, from, to)) = fired { + match with_path(request.uri(), &to) { + Ok(uri) => { + tracing::debug!(rule = rule.label(), %from, %to, "url rewrite applied"); + *request.uri_mut() = uri; + } + Err(error) => { + // A template can assemble an invalid path out of matched + // input (e.g. an empty string). Serving the original path is + // the conservative outcome: the request 404s the same way it + // would have without the layer, and the warn names the rule. + tracing::warn!( + rule = rule.label(), + %from, + rewritten = %to, + %error, + "url rewrite produced an invalid path; leaving the request unrewritten" + ); + } + } + } + next.run(request).await +} + +/// `uri` with its path replaced by `new_path`, query preserved. +fn with_path(uri: &http::Uri, new_path: &str) -> Result { + let path_and_query = match uri.query() { + Some(query) => format!("{new_path}?{query}"), + None => new_path.to_owned(), + }; + let mut parts = uri.clone().into_parts(); + parts.path_and_query = Some(path_and_query.parse()?); + Ok(http::Uri::from_parts(parts)?) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rule(pattern: &str, replacement: &str) -> UrlRewriteRule { + UrlRewriteRule { + name: None, + pattern: pattern.to_string(), + replacement: replacement.to_string(), + } + } + + #[test] + fn apply_substitutes_capture_groups() { + let compiled = compile(&[rule("^/mcp-servers/([^/]+)/mcp$", "/mcp/$1")]); + assert_eq!( + compiled[0].apply("/mcp-servers/github/mcp").as_deref(), + Some("/mcp/github") + ); + assert_eq!(compiled[0].apply("/mcp-servers/github/sse"), None); + assert_eq!(compiled[0].apply("/v1/chat/completions"), None); + } + + #[test] + fn apply_replaces_only_the_matched_portion() { + // Unanchored pattern: the unmatched prefix survives, mirroring the + // replace-matched-portion semantics of mainstream gateways. + let compiled = compile(&[rule("/legacy$", "/current")]); + assert_eq!( + compiled[0].apply("/api/legacy").as_deref(), + Some("/api/current") + ); + } + + #[test] + fn apply_supports_named_groups_and_braced_references() { + let compiled = compile(&[rule( + "^/gw/(?P[^/]+)/v(\\d+)$", + "/mcp/${server}-v${2}", + )]); + assert_eq!( + compiled[0].apply("/gw/github/v2").as_deref(), + Some("/mcp/github-v2") + ); + } + + #[test] + fn with_path_preserves_the_query_string() { + let uri: http::Uri = "http://gw.example/mcp-servers/github/mcp?probe=1" + .parse() + .unwrap(); + let rewritten = with_path(&uri, "/mcp/github").unwrap(); + assert_eq!(rewritten.path(), "/mcp/github"); + assert_eq!(rewritten.query(), Some("probe=1")); + assert_eq!(rewritten.host(), Some("gw.example")); + } + + #[test] + fn with_path_rejects_an_invalid_path() { + let uri: http::Uri = "/x".parse().unwrap(); + // A template can assemble characters that are invalid in a request + // path; the middleware then keeps the original URI (warn + no-op). + assert!(with_path(&uri, "/a b").is_err()); + } + + fn router_with_rules(rules: Vec) -> axum::Router { + use aisix_core::snapshot::SnapshotHandle; + let cfg = aisix_core::ProxyConfig { + addr: "127.0.0.1:0".into(), + request_body_limit_bytes: 0, + tls: None, + real_ip: Default::default(), + url_rewrites: rules, + }; + let state = ProxyState::new( + SnapshotHandle::new(aisix_core::AisixSnapshot::new()), + Arc::new(aisix_gateway::Hub::new()), + &cfg, + ) + .without_cache(); + crate::build_router(state) + } + + async fn get(router: axum::Router, path: &str) -> axum::http::StatusCode { + use tower::ServiceExt; + let request = axum::http::Request::get(path) + .body(axum::body::Body::empty()) + .unwrap(); + router.oneshot(request).await.unwrap().status() + } + + #[tokio::test] + async fn router_serves_a_rewritten_legacy_path_and_first_rule_wins() { + // Two rules match the legacy path; the first rewrites onto a real + // endpoint, the second onto a 404. Declaration order must win. + let router = router_with_rules(vec![ + rule("^/legacy/health$", "/livez"), + rule("^/legacy/health$", "/nonexistent"), + ]); + assert_eq!(get(router.clone(), "/legacy/health").await, 200); + // A miss flows through unrewritten. + assert_eq!(get(router.clone(), "/legacy/other").await, 404); + // The canonical path keeps working alongside the legacy one. + assert_eq!(get(router, "/livez").await, 200); + } + + #[tokio::test] + async fn router_without_rules_is_untouched() { + let router = router_with_rules(Vec::new()); + assert_eq!(get(router.clone(), "/livez").await, 200); + assert_eq!(get(router, "/legacy/health").await, 404); + } +} diff --git a/crates/aisix-proxy/src/state.rs b/crates/aisix-proxy/src/state.rs index e278573ab..1f7b8f10f 100644 --- a/crates/aisix-proxy/src/state.rs +++ b/crates/aisix-proxy/src/state.rs @@ -145,6 +145,9 @@ pub struct ProxyStateInner { /// client IP on each request (#492). Default = trust nothing → the /// logged source IP is the immediate TCP peer. pub real_ip: Arc, + /// Boot-compiled `proxy.url_rewrites` rules, applied in order to every + /// request before routing (first match wins). Empty = layer no-ops. + pub url_rewrites: Arc<[crate::rewrite::CompiledRewrite]>, /// Optional config-freshness probe for `GET /readyz`: returns the time /// since the etcd watch last applied config (`None` = never applied). /// Wired from the watch supervisor in aisix-server; `None` here means @@ -208,6 +211,7 @@ impl ProxyState { otlp_fan_out: OtlpHttpFanOut::new(), request_body_limit_bytes: cfg.request_body_limit_bytes, real_ip: Arc::new(ResolvedRealIp::from_config(&cfg.real_ip)), + url_rewrites: crate::rewrite::compile(&cfg.url_rewrites), billed_batches: Arc::new(dashmap::DashSet::new()), client_classifier: Arc::new(ClientTypeClassifier::builtin()), default_retries: aisix_core::config::DEFAULT_UPSTREAM_RETRIES, @@ -244,6 +248,7 @@ impl ProxyState { otlp_fan_out: OtlpHttpFanOut::new(), request_body_limit_bytes: cfg.request_body_limit_bytes, real_ip: Arc::new(ResolvedRealIp::from_config(&cfg.real_ip)), + url_rewrites: crate::rewrite::compile(&cfg.url_rewrites), billed_batches: Arc::new(dashmap::DashSet::new()), client_classifier: Arc::new(ClientTypeClassifier::builtin()), default_retries: aisix_core::config::DEFAULT_UPSTREAM_RETRIES, @@ -290,6 +295,7 @@ impl ProxyState { otlp_fan_out: OtlpHttpFanOut::new(), request_body_limit_bytes: cfg.request_body_limit_bytes, real_ip: Arc::new(ResolvedRealIp::from_config(&cfg.real_ip)), + url_rewrites: crate::rewrite::compile(&cfg.url_rewrites), billed_batches: Arc::new(dashmap::DashSet::new()), client_classifier: Arc::new(ClientTypeClassifier::builtin()), default_retries: aisix_core::config::DEFAULT_UPSTREAM_RETRIES, @@ -388,6 +394,7 @@ mod tests { request_body_limit_bytes: 1_048_576, tls: None, real_ip: Default::default(), + url_rewrites: Vec::new(), }, ) } diff --git a/crates/aisix-proxy/src/videos.rs b/crates/aisix-proxy/src/videos.rs index ecb0968b2..a16339de8 100644 --- a/crates/aisix-proxy/src/videos.rs +++ b/crates/aisix-proxy/src/videos.rs @@ -2327,6 +2327,7 @@ mod tests { addr: "127.0.0.1:0".into(), request_body_limit_bytes: 1_048_576, real_ip: Default::default(), + url_rewrites: Vec::new(), tls: None, } } diff --git a/tests/e2e/src/cases/url-rewrite-e2e.test.ts b/tests/e2e/src/cases/url-rewrite-e2e.test.ts new file mode 100644 index 000000000..874790fae --- /dev/null +++ b/tests/e2e/src/cases/url-rewrite-e2e.test.ts @@ -0,0 +1,219 @@ +import { createHash, randomUUID } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + spawnApp, + startMcpUpstream, + waitConfigPropagation, + type McpUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: entry-level URL rewriting (`proxy.url_rewrites`) against a real +// gateway + etcd + a real MCP upstream. +// +// Pinned contract: +// - the first rule whose `match` regex matches the path rewrites it; the +// request then flows through the normal endpoint (auth, ACL, quota) as +// if the client had sent the rewritten path; +// - the flagship scenario: a client keeping its legacy per-server MCP URL +// (`/mcp-servers/{service}/mcp`) and ORIGINAL tool names works end to +// end through the rewritten `/mcp/{server}` endpoint; +// - rewriting is generic, not MCP-specific (any path → any endpoint); +// - a miss leaves the request untouched (canonical paths keep working, +// unmatched legacy shapes 404); +// - the query string survives the rewrite. + +const KEY = "sk-url-rewrite-e2e"; +const sha256 = (s: string) => createHash("sha256").update(s).digest("hex"); + +interface RpcReply { + status: number; + json?: { + result?: { + serverInfo?: { name?: string }; + tools?: Array<{ name: string }>; + content?: Array<{ type: string; text?: string }>; + isError?: boolean; + }; + error?: { code: number; message: string }; + }; +} + +describe("url rewrite e2e: proxy.url_rewrites", () => { + let app: SpawnedApp | undefined; + let alpha: McpUpstream | undefined; + let etcdReachable = false; + + const post = async ( + path: string, + body: unknown, + ): Promise => { + const res = await fetch(`${app!.proxyUrl}${path}`, { + method: "POST", + headers: { + authorization: `Bearer ${KEY}`, + "content-type": "application/json", + accept: "application/json, text/event-stream", + }, + body: JSON.stringify(body), + }); + const text = await res.text(); + let json: RpcReply["json"]; + try { + json = text ? JSON.parse(text) : undefined; + } catch { + json = undefined; + } + return { status: res.status, json }; + }; + + const initialize = async ( + path: string, + ): Promise<{ status: number; serverName?: string }> => { + const init = await post(path, { + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-11-25", + capabilities: {}, + clientInfo: { name: "url-rewrite-e2e", version: "0.1" }, + }, + }); + if (init.status !== 200) return { status: init.status }; + await post(path, { jsonrpc: "2.0", method: "notifications/initialized" }); + return { + status: init.status, + serverName: init.json?.result?.serverInfo?.name, + }; + }; + + const listToolNames = async ( + path: string, + ): Promise<{ status: number; names?: string[] }> => { + const { status } = await initialize(path); + if (status !== 200) return { status }; + const r = await post(path, { + jsonrpc: "2.0", + id: 2, + method: "tools/list", + params: {}, + }); + const tools = r.json?.result?.tools; + if (r.status !== 200 || !tools) return { status: r.status }; + return { status: r.status, names: tools.map((t) => t.name).sort() }; + }; + + const callTool = async ( + path: string, + name: string, + text: string, + ): Promise<{ ok: boolean; text?: string; error?: string }> => { + await initialize(path); + const r = await post(path, { + jsonrpc: "2.0", + id: 3, + method: "tools/call", + params: { name, arguments: { text } }, + }); + if (r.json?.error) return { ok: false, error: r.json.error.message }; + const result = r.json?.result; + if (!result || result.isError) { + return { ok: false, error: JSON.stringify(r.json ?? r.status) }; + } + return { ok: true, text: result.content?.[0]?.text }; + }; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + alpha = await startMcpUpstream("alpha"); + app = await spawnApp({ + urlRewrites: [ + { + name: "per-server-mcp-compat", + match: "^/mcp-servers/([^/]+)/mcp$", + rewrite: "/mcp/$1", + }, + // Generic, non-MCP mapping: proves the layer rewrites any path + // onto any endpoint. + { match: "^/compat/health$", rewrite: "/livez" }, + ], + }); + const seed = new SeedClient(new EtcdClient(), app.etcdPrefix); + + await seed.update("mcp_servers", randomUUID(), { + display_name: "alpha", + url: alpha.url, + enabled: true, + }); + await seed.createApiKey({ + key_hash: sha256(KEY), + allowed_models: [], + allowed_tools: ["*"], + }); + + await waitConfigPropagation(async () => { + const listed = await listToolNames("/mcp-servers/alpha/mcp"); + return ( + listed.status === 200 && + JSON.stringify(listed.names) === JSON.stringify(["echo", "reverse"]) + ); + }); + }, 60_000); + + afterAll(async () => { + await app?.exit(); + await alpha?.close(); + }); + + test("legacy per-server MCP URL + original tool names work end to end", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + // The migration scenario in one flow: the client keeps its legacy URL + // and its original tool name; the rewrite maps the URL onto + // /mcp/{server}, which serves original names. + const init = await initialize("/mcp-servers/alpha/mcp"); + expect(init.status).toBe(200); + expect(init.serverName).toBe("alpha"); + + const listed = await listToolNames("/mcp-servers/alpha/mcp"); + expect(listed.names).toEqual(["echo", "reverse"]); + + const called = await callTool("/mcp-servers/alpha/mcp", "echo", "hi"); + expect(called).toEqual({ ok: true, text: "alpha:hi" }); + }); + + test("rewriting is generic: any path maps onto any endpoint", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + const res = await fetch(`${app.proxyUrl}/compat/health`); + expect(res.status).toBe(200); + + // The query string survives the rewrite (the endpoint ignores it, but + // a broken rewrite would 404 or drop it before routing). + const withQuery = await fetch(`${app.proxyUrl}/compat/health?probe=1`); + expect(withQuery.status).toBe(200); + }); + + test("a miss leaves the request untouched", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + // Canonical paths keep working alongside the legacy shapes… + const canonical = await initialize("/mcp/alpha"); + expect(canonical.status).toBe(200); + const livez = await fetch(`${app.proxyUrl}/livez`); + expect(livez.status).toBe(200); + + // …and a legacy shape the rule does not match is NOT silently routed + // anywhere (`backend_path` other than /mcp stays unserved). + const unmatchedTail = await post("/mcp-servers/alpha/sse", {}); + expect(unmatchedTail.status).toBe(404); + const unknownShape = await fetch(`${app.proxyUrl}/compat/other`); + expect(unknownShape.status).toBe(404); + }); +}); diff --git a/tests/e2e/src/harness/app.ts b/tests/e2e/src/harness/app.ts index d5bc9ea88..e41510f26 100644 --- a/tests/e2e/src/harness/app.ts +++ b/tests/e2e/src/harness/app.ts @@ -42,6 +42,13 @@ export interface AppOverrides { recursive?: boolean; header?: string; }; + /** + * `proxy.url_rewrites` block. Merged into the base proxy config (like + * `realIp`) so the listener addr is preserved. Entry-level path + * rewriting: first matching rule wins, `rewrite` replaces the matched + * portion of the path. + */ + urlRewrites?: Array<{ name?: string; match: string; rewrite: string }>; /** * `proxy.request_body_limit_bytes`. A dedicated override (like * `realIp`) because `extra` replaces whole top-level blocks and the @@ -211,6 +218,7 @@ async function spawnAppOnce(overrides: AppOverrides = {}): Promise { addr: `127.0.0.1:${proxyPort}`, request_body_limit_bytes: overrides.requestBodyLimitBytes ?? 10485760, ...(overrides.realIp ? { real_ip: overrides.realIp } : {}), + ...(overrides.urlRewrites ? { url_rewrites: overrides.urlRewrites } : {}), }, admin: adminEnabled ? { addr: `127.0.0.1:${adminPort}`, admin_keys: [adminKey] } From f947ef477241f902d53ec0384d2b4e61fe077533 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Tue, 4 Aug 2026 17:38:26 +0800 Subject: [PATCH 4/5] fix(mcp): reserve disabled servers' prefixes on the scoped endpoint too Review follow-up: the foreign set excluded disabled entries, so toggling another server's enabled flag changed which names a scope would serve bare vs fail closed on. Reserve every other registered name regardless of enabled state; the round-trip listing keeps the colliding literal names namespaced, so they stay callable. Also assert the ACL rejection wording in the scoped e2e denial case. --- crates/aisix-mcp/src/gateway.rs | 10 ++++--- crates/aisix-mcp/tests/gateway_scoped.rs | 28 +++++++++++++++++++ .../src/cases/mcp-scoped-endpoint-e2e.test.ts | 1 + 3 files changed, 35 insertions(+), 4 deletions(-) diff --git a/crates/aisix-mcp/src/gateway.rs b/crates/aisix-mcp/src/gateway.rs index c5effd8d1..d9034f9e5 100644 --- a/crates/aisix-mcp/src/gateway.rs +++ b/crates/aisix-mcp/src/gateway.rs @@ -259,9 +259,11 @@ struct ScopedServer { /// The scoped upstream's registered name — the namespace every ACL check /// re-applies and the name `initialize` reports. name: String, - /// The **other** registered, enabled server names. A `tools/call` whose - /// name is prefixed with one of these is a cross-server mistake and fails - /// closed rather than being silently served as a bare name. + /// Every **other** registered server name, enabled or not. A `tools/call` + /// whose name is prefixed with one of these is a cross-server mistake and + /// fails closed rather than being silently served as a bare name. + /// Disabled servers stay reserved so this scope's callable name surface + /// does not shift when another server's `enabled` flag is toggled. foreign: std::collections::HashSet, } @@ -360,7 +362,7 @@ impl McpGateway { .mcp_servers .entries() .into_iter() - .filter(|e| e.value.enabled && e.value.name != name) + .filter(|e| e.value.name != name) .map(|e| e.value.name.clone()) .collect(); let mut gateway = McpGateway::new([(name.clone(), bridge)]); diff --git a/crates/aisix-mcp/tests/gateway_scoped.rs b/crates/aisix-mcp/tests/gateway_scoped.rs index f95978837..1e4328786 100644 --- a/crates/aisix-mcp/tests/gateway_scoped.rs +++ b/crates/aisix-mcp/tests/gateway_scoped.rs @@ -269,6 +269,34 @@ async fn scoped_registered_foreign_prefix_fails_closed() { assert_eq!(first_text(&escaped), "alpha:hi"); } +#[tokio::test] +async fn scoped_disabled_foreign_prefix_still_fails_closed() { + // `dark` is registered but disabled. Its name stays reserved: a scope's + // callable name surface must not shift when another server's `enabled` + // flag is toggled, so `dark__echo` fails closed on alpha even though the + // upstream would serve that literal name. + let upstream = spawn_upstream("alpha", "dark__echo").await; + let snapshot = snapshot_with_alpha(&upstream); + let gateway = + McpGateway::from_snapshot_scoped(&snapshot, "alpha").expect("alpha is registered"); + let client = connect(spawn_gateway(gateway).await).await; + + assert!( + client.call_tool(call("dark__echo", "x")).await.is_err(), + "a disabled registered server's prefix must still fail closed" + ); + // The literal tool stays reachable through the advertised namespaced + // spelling, same as the enabled-foreign case. + let tools = client.list_all_tools().await.expect("list tools"); + let names: Vec<&str> = tools.iter().map(|t| t.name.as_ref()).collect(); + assert_eq!(names, vec!["alpha__dark__echo"]); + let escaped = client + .call_tool(call("alpha__dark__echo", "hi")) + .await + .expect("the namespaced spelling reaches the literal tool"); + assert_eq!(first_text(&escaped), "alpha:hi"); +} + #[tokio::test] async fn scoped_server_name_ending_in_underscore_namespaces_cleanly() { // `data_` is a legal server name (only `__` inside a name is rejected). diff --git a/tests/e2e/src/cases/mcp-scoped-endpoint-e2e.test.ts b/tests/e2e/src/cases/mcp-scoped-endpoint-e2e.test.ts index 30743c6d5..17fc8c5c6 100644 --- a/tests/e2e/src/cases/mcp-scoped-endpoint-e2e.test.ts +++ b/tests/e2e/src/cases/mcp-scoped-endpoint-e2e.test.ts @@ -288,6 +288,7 @@ describe("mcp scoped endpoint e2e: /mcp/{server}", () => { expect(foreign.names).toEqual([]); const rejected = await callTool("/mcp/alpha", KEY_BETA_ONLY, "echo", "hi"); expect(rejected.ok).toBe(false); + expect(rejected.error).toContain("not available"); }); test("unknown and disabled servers are 404 with no fallback", async (ctx) => { From 1483221c07dae6580516e7adaf92a0d9608798a8 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Tue, 4 Aug 2026 17:58:28 +0800 Subject: [PATCH 5/5] fix(config): validate rewrite templates and accept the env JSON form MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Independent review findings on the rewrite layer: - proxy.url_rewrites is a struct list, which AISIX_* env vars — the only config channel in chart-driven deployments — cannot express. The field now also accepts one JSON array in a string, so AISIX_PROXY__URL_REWRITES='[{...}]' works. - Only the match regex was validated; the template could reference an unknown capture group (expands to empty — every legacy request lands on the wrong endpoint, silently), or carry '?'/'#'/whitespace (absorbs the caller's query / truncates the path as a fragment). Config::validate now rejects unknown group references, forbidden template characters, and patterns that match the empty string (which would fire on every request). - Request ids now cover the rewrite layer: ensure_request_id moved outside the wrapper, so fired/failed rewrite logs carry the request span. - Pin the remaining review gaps: warn-fallback serves the original path (router-level), raw percent-encoded matching (no decode/normalize), and first-occurrence replacement for unanchored patterns. Document the raw-path semantics and the env JSON form. --- config.example.yaml | 13 +- config.managed.yaml | 10 +- crates/aisix-core/src/config.rs | 271 +++++++++++++++++++- crates/aisix-proxy/src/lib.rs | 35 +-- crates/aisix-proxy/src/rewrite.rs | 38 ++- tests/e2e/src/cases/url-rewrite-e2e.test.ts | 6 +- 6 files changed, 328 insertions(+), 45 deletions(-) diff --git a/config.example.yaml b/config.example.yaml index 5f5b1c6b7..c4e8e09c7 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -55,10 +55,15 @@ proxy: # before routing. The first rule whose `match` regex matches the request # path rewrites it (once, no cascading); the request then flows through # the normal endpoint — auth, ACL, quota — as if the client had sent the - # rewritten path. `rewrite` replaces the matched portion of the path - # ($1/${name} expand capture groups); the query string is preserved. An - # invalid regex fails startup. Example: serve per-server MCP URLs like - # /mcp-servers/github/mcp on the /mcp/{server} endpoint. + # rewritten path. `match` runs against the RAW, percent-encoded path (no + # decoding or normalization); `rewrite` replaces the matched portion + # ($1/${name} expand capture groups; use ${1}text to follow a group with + # literal text); the query string is preserved. Invalid regexes, unknown + # group references, and `?`/`#`/whitespace in the template fail startup. + # In env-only deployments set the whole list as one JSON array: + # AISIX_PROXY__URL_REWRITES='[{"match":"...","rewrite":"..."}]' + # Example: serve per-server MCP URLs like /mcp-servers/github/mcp on + # the /mcp/{server} endpoint. # url_rewrites: # - name: per-server-mcp-compat # match: "^/mcp-servers/([^/]+)/mcp$" diff --git a/config.managed.yaml b/config.managed.yaml index 5e35e8770..cd7311ab7 100644 --- a/config.managed.yaml +++ b/config.managed.yaml @@ -44,10 +44,12 @@ proxy: # 0 = no request-body cap (the default); set a value to bound # per-request memory. # request_body_limit_bytes: 0 - # Entry-level URL rewriting (first matching rule wins; `rewrite` - # replaces the matched portion of the path, $1 = capture group; the - # query string is preserved). Lets clients keep legacy URL shapes, - # e.g. per-server MCP paths served on the /mcp/{server} endpoint. + # Entry-level URL rewriting (first matching rule wins; `match` runs on + # the raw, percent-encoded path; `rewrite` replaces the matched portion, + # $1 = capture group; the query string is preserved). Lets clients keep + # legacy URL shapes, e.g. per-server MCP paths served on the + # /mcp/{server} endpoint. When configuring through env vars only, set + # one JSON array: AISIX_PROXY__URL_REWRITES='[{"match":"...","rewrite":"..."}]' # url_rewrites: # - name: per-server-mcp-compat # match: "^/mcp-servers/([^/]+)/mcp$" diff --git a/crates/aisix-core/src/config.rs b/crates/aisix-core/src/config.rs index 27a28fa8e..6fb4814c3 100644 --- a/crates/aisix-core/src/config.rs +++ b/crates/aisix-core/src/config.rs @@ -401,7 +401,12 @@ pub struct ProxyConfig { /// had sent the rewritten path. Lets operators map legacy URL shapes /// onto AISIX endpoints, e.g. per-server MCP paths onto /// `/mcp/{server}`. Empty (the default) = no rewriting. - #[serde(default)] + /// + /// Env-only deployments (the chart injects config purely through + /// `AISIX_*` vars, which cannot express a structured list) set the + /// whole list as one JSON array: + /// `AISIX_PROXY__URL_REWRITES='[{"match":"^/x$","rewrite":"/y"}]'`. + #[serde(default, deserialize_with = "deserialize_url_rewrites")] pub url_rewrites: Vec, } @@ -418,19 +423,111 @@ pub struct UrlRewriteRule { /// Optional name, used in logs when the rule fires. #[serde(default)] pub name: Option, - /// Regex matched against the request path (never the query string). - /// Anchor with `^`/`$` to match the whole path; an unanchored pattern - /// matches anywhere in it. + /// Regex matched against the **raw, percent-encoded** request path + /// (never the query string) — no decoding, no normalization. Anchor + /// with `^`/`$` to match the whole path; an unanchored pattern matches + /// anywhere in it. Must not match the empty string. #[serde(rename = "match")] pub pattern: String, /// Replacement for the matched portion of the path. Capture groups are /// available as `$1`… / `${name}`; use `${1}x` (braced) when a literal - /// character follows a group reference. The query string is preserved - /// as sent. + /// character follows a group reference (`$1x` reads as the group named + /// `1x`). The query string is preserved as sent, so the template must + /// not contain `?`, `#`, whitespace, or control characters. #[serde(rename = "rewrite")] pub replacement: String, } +/// Accept `url_rewrites` either as a structured sequence (config file) or +/// as a JSON array carried in one string — the only shape an env var can +/// hold, and env vars are the sole config channel in chart-driven +/// deployments. +fn deserialize_url_rewrites<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + #[derive(Deserialize)] + #[serde(untagged)] + enum SeqOrJsonString { + Seq(Vec), + JsonString(String), + } + match SeqOrJsonString::deserialize(deserializer)? { + SeqOrJsonString::Seq(rules) => Ok(rules), + SeqOrJsonString::JsonString(raw) => { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Ok(Vec::new()); + } + serde_json::from_str(trimmed) + .map_err(|e| serde::de::Error::custom(format!("url_rewrites JSON string: {e}"))) + } + } +} + +/// Reject a rewrite template that references capture groups its pattern +/// does not define — the regex engine expands unknown references to the +/// empty string, which would silently rewrite traffic to the wrong +/// endpoint. Mirrors the engine's replacement syntax: `$$` is a literal +/// `$`, `${name}` is a braced reference, and a bare `$name` reference +/// spans the longest run of `[0-9A-Za-z_]` (so `$1x` reads as a group +/// named `1x`, not group 1 followed by `x`). +fn validate_rewrite_template_refs(regex: ®ex::Regex, template: &str) -> Result<(), String> { + let names: std::collections::HashSet<&str> = regex.capture_names().flatten().collect(); + let group_count = regex.captures_len(); // includes group 0 (the whole match) + let ref_ok = |name: &str| { + if name.chars().all(|c| c.is_ascii_digit()) { + name.parse::().is_ok_and(|idx| idx < group_count) + } else { + names.contains(name) + } + }; + let bytes = template.as_bytes(); + let mut i = 0; + while i < bytes.len() { + if bytes[i] != b'$' { + i += 1; + continue; + } + if bytes.get(i + 1) == Some(&b'$') { + i += 2; + continue; + } + if bytes.get(i + 1) == Some(&b'{') { + let Some(end) = template[i + 2..].find('}') else { + return Err("rewrite has an unterminated `${…}` group reference".to_string()); + }; + let name = &template[i + 2..i + 2 + end]; + if name.is_empty() || !ref_ok(name) { + return Err(format!( + "rewrite references unknown capture group `${{{name}}}`" + )); + } + i += 2 + end + 1; + continue; + } + let start = i + 1; + let mut end = start; + while end < bytes.len() && (bytes[end].is_ascii_alphanumeric() || bytes[end] == b'_') { + end += 1; + } + if end == start { + // A bare trailing `$`: the engine treats it as a literal. + i += 1; + continue; + } + let name = &template[start..end]; + if !ref_ok(name) { + return Err(format!( + "rewrite references unknown capture group `${name}` \ + (write `${{N}}text` to follow group N with literal text)" + )); + } + i = end; + } + Ok(()) +} + /// nginx `set_real_ip_from` + `real_ip_recursive` equivalent. Resolves /// the downstream client IP for usage logs (#492) from a forwarded /// header, trusting only addresses inside `trusted_proxies`. @@ -1124,19 +1221,53 @@ impl Config { } fn validate(&self) -> Result<(), BootstrapError> { - // Fail fast on a rewrite rule that can never compile: it would - // otherwise surface as every legacy-path request 404ing, which is - // much harder to trace back to a typo in one regex. + // Fail fast on an unusable rewrite rule: a broken rule would + // otherwise surface at runtime as silently mis-routed or 404ing + // legacy traffic, which is much harder to trace back to a typo in + // one line of config. for (i, rule) in self.proxy.url_rewrites.iter().enumerate() { - if let Err(e) = regex::Regex::new(&rule.pattern) { - let ctx = rule - .name + let ctx = || { + rule.name .clone() - .unwrap_or_else(|| format!("proxy.url_rewrites[{i}]")); + .unwrap_or_else(|| format!("proxy.url_rewrites[{i}]")) + }; + let regex = match regex::Regex::new(&rule.pattern) { + Ok(regex) => regex, + Err(e) => { + return Err(BootstrapError::Config(format!( + "{}: invalid match regex: {e}", + ctx() + ))); + } + }; + // A pattern that matches the empty string would fire on every + // request (zero-width match at position 0) and prepend the + // template to every path. + if regex.find("").is_some() { return Err(BootstrapError::Config(format!( - "{ctx}: invalid match regex: {e}" + "{}: match must not match the empty string", + ctx() ))); } + // The template lands inside a URI path; a `?` would absorb the + // caller's query into itself and a `#` would truncate the path + // as a fragment — both silently. Reject them (and unprintables) + // up front; capture-group expansions are safe because a request + // path can never contain these characters raw. + if let Some(bad) = rule + .replacement + .chars() + .find(|c| matches!(c, '?' | '#') || c.is_whitespace() || c.is_control()) + { + return Err(BootstrapError::Config(format!( + "{}: rewrite must not contain {bad:?} (the template is a path; \ + the query string is preserved automatically)", + ctx() + ))); + } + if let Err(e) = validate_rewrite_template_refs(®ex, &rule.replacement) { + return Err(BootstrapError::Config(format!("{}: {e}", ctx()))); + } } if let Some(path) = self.resources_file.as_deref() { // File source selected: exactly one resource source may be @@ -1394,6 +1525,118 @@ admin: ); } + #[test] + fn url_rewrites_accepts_a_json_string_for_env_only_deployments() { + // Chart-driven deployments inject config purely through AISIX_* env + // vars, which cannot express a structured list — the whole list + // rides in one JSON string. A YAML string scalar takes the same + // code path as the env source. + let f = write_yaml( + r#" +etcd: + endpoints: ["http://127.0.0.1:2379"] + prefix: "/aisix" +proxy: + addr: "0.0.0.0:3000" + url_rewrites: '[{"name":"compat","match":"^/mcp-servers/([^/]+)/mcp$","rewrite":"/mcp/$1"}]' +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +"#, + ); + let cfg = Config::load_from_path(Some(f.path())).unwrap(); + assert_eq!(cfg.proxy.url_rewrites.len(), 1); + assert_eq!(cfg.proxy.url_rewrites[0].name.as_deref(), Some("compat")); + + let f = write_yaml( + r#" +etcd: + endpoints: ["http://127.0.0.1:2379"] + prefix: "/aisix" +proxy: + addr: "0.0.0.0:3000" + url_rewrites: '[{"match": broken' +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +"#, + ); + let err = Config::load_from_path(Some(f.path())).unwrap_err(); + assert!( + format!("{err}").contains("url_rewrites"), + "error should name the field: {err}" + ); + } + + #[test] + fn url_rewrites_rejects_silent_template_mistakes() { + let case = |rule_yaml: &str, expect: &str| { + let f = write_yaml(&format!( + r#" +etcd: + endpoints: ["http://127.0.0.1:2379"] + prefix: "/aisix" +proxy: + addr: "0.0.0.0:3000" + url_rewrites: +{rule_yaml} +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +"# + )); + let err = Config::load_from_path(Some(f.path())).unwrap_err(); + assert!( + format!("{err}").contains(expect), + "expected {expect:?} in: {err}" + ); + }; + + // An unknown group reference expands to the empty string at runtime + // — every legacy request would silently land on the wrong endpoint. + case( + " - match: \"^/mcp-servers/([^/]+)/mcp$\"\n rewrite: \"/mcp/$2\"", + "unknown capture group", + ); + case( + " - match: \"^/gw/(?P[^/]+)$\"\n rewrite: \"/mcp/${srv}\"", + "unknown capture group", + ); + // `?` would absorb the caller's query; `#` would truncate the path. + case( + " - match: \"^/a$\"\n rewrite: \"/v1/chat?model=x\"", + "must not contain", + ); + case( + " - match: \"^/a$\"\n rewrite: \"/v1/models#frag\"", + "must not contain", + ); + // A pattern matching the empty string fires on every request. + case( + " - match: \"(x)?\"\n rewrite: \"/y\"", + "empty string", + ); + + // The braced form with literal text after a group is the valid way + // to write what `$1x` cannot mean. + let f = write_yaml( + r#" +etcd: + endpoints: ["http://127.0.0.1:2379"] + prefix: "/aisix" +proxy: + addr: "0.0.0.0:3000" + url_rewrites: + - match: "^/gw/(?P[^/]+)/v(\\d+)$" + rewrite: "/mcp/${server}-v${2}" +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +"#, + ); + Config::load_from_path(Some(f.path())).expect("braced references are valid"); + } + #[test] fn rejects_malformed_trusted_proxy_cidr() { let f = write_yaml( diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index dbdff1117..ac2b2dd36 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -221,18 +221,8 @@ pub fn build_router(state: ProxyState) -> Router { header::SERVER, SERVER_HEADER_VALUE.clone(), )) - // Rewrite the request path per `proxy.url_rewrites` before route - // Outermost: mint the request id into the request extensions - // before any handler/extractor runs, and stamp it onto every - // response (including the short-circuited 4xx from the layers - // above) so the whole proxy family carries `x-aisix-request-id` - // and it equals the telemetry request_id. See request_id.rs. - .layer(middleware::from_fn(request_id::ensure_request_id)) .with_state(state.clone()); - if state.url_rewrites.is_empty() { - return router; - } // Pre-routing URL rewriting (`proxy.url_rewrites`). `Router::layer` // middleware runs AFTER route matching, so a URI rewritten there could // never change which route matches. Wrapping the whole router as the @@ -241,12 +231,25 @@ pub fn build_router(state: ProxyState) -> Router { // it mutates the URI, then the inner router matches on the rewritten // path. Built only when rules are configured, so the default path pays // nothing. See rewrite.rs. - Router::new() - .fallback_service(router) - .layer(middleware::from_fn_with_state( - state, - rewrite::rewrite_request_uri, - )) + let router = if state.url_rewrites.is_empty() { + router + } else { + Router::new() + .fallback_service(router) + .layer(middleware::from_fn_with_state( + state, + rewrite::rewrite_request_uri, + )) + }; + + // Outermost: mint the request id into the request extensions + // before any handler/extractor runs — including the rewrite layer, + // whose fired/failed log lines must carry the request span — and + // stamp it onto every response (including the short-circuited 4xx + // from the layers above) so the whole proxy family carries + // `x-aisix-request-id` and it equals the telemetry request_id. See + // request_id.rs. + router.layer(middleware::from_fn(request_id::ensure_request_id)) } async fn record_in_flight_request( diff --git a/crates/aisix-proxy/src/rewrite.rs b/crates/aisix-proxy/src/rewrite.rs index 462eccee8..d9cfbb446 100644 --- a/crates/aisix-proxy/src/rewrite.rs +++ b/crates/aisix-proxy/src/rewrite.rs @@ -9,8 +9,11 @@ //! //! Replacement substitutes the **matched portion** of the path, with //! `$1`/`${name}` expanding capture groups; the query string is preserved as -//! sent. Rules were syntax-checked at config load (`Config::validate`), so -//! compiling them here cannot fail on operator input. +//! sent. Patterns are matched against the **raw, percent-encoded** request +//! path — no decoding, no normalization — unlike gateways that match a +//! decoded `$uri`. Rules were validated at config load (`Config::validate`: +//! regex syntax, template group references, forbidden template characters), +//! so compiling them here cannot fail on operator input. use std::borrow::Cow; use std::sync::Arc; @@ -70,9 +73,6 @@ pub async fn rewrite_request_uri( mut request: Request, next: Next, ) -> Response { - if state.url_rewrites.is_empty() { - return next.run(request).await; - } let fired = state.url_rewrites.iter().find_map(|rule| { let path = request.uri().path(); rule.apply(path).map(|to| (rule, path.to_owned(), to)) @@ -224,4 +224,32 @@ mod tests { assert_eq!(get(router.clone(), "/livez").await, 200); assert_eq!(get(router, "/legacy/health").await, 404); } + + #[tokio::test] + async fn runtime_invalid_rewrite_leaves_the_request_unrewritten() { + // Bypasses Config::validate on purpose (ProxyState::new takes the + // config directly): a template that assembles an invalid path must + // warn and serve the ORIGINAL path — here /livez still routes. + let router = router_with_rules(vec![rule("^/livez$", "/livez x")]); + assert_eq!(get(router, "/livez").await, 200); + } + + #[tokio::test] + async fn router_matches_the_raw_percent_encoded_path() { + // No decoding, no normalization: an encoded variant of a matching + // path does not match and passes through to a 404. + let router = router_with_rules(vec![rule("^/legacy/health$", "/livez")]); + assert_eq!(get(router.clone(), "/legacy/health").await, 200); + assert_eq!(get(router.clone(), "/legacy%2Fhealth").await, 404); + assert_eq!(get(router, "/legac%79/health").await, 404); + } + + #[test] + fn apply_replaces_only_the_first_occurrence() { + let compiled = compile(&[rule("/legacy", "/cur")]); + assert_eq!( + compiled[0].apply("/legacy/legacy").as_deref(), + Some("/cur/legacy") + ); + } } diff --git a/tests/e2e/src/cases/url-rewrite-e2e.test.ts b/tests/e2e/src/cases/url-rewrite-e2e.test.ts index 874790fae..d46970e14 100644 --- a/tests/e2e/src/cases/url-rewrite-e2e.test.ts +++ b/tests/e2e/src/cases/url-rewrite-e2e.test.ts @@ -194,8 +194,10 @@ describe("url rewrite e2e: proxy.url_rewrites", () => { const res = await fetch(`${app.proxyUrl}/compat/health`); expect(res.status).toBe(200); - // The query string survives the rewrite (the endpoint ignores it, but - // a broken rewrite would 404 or drop it before routing). + // A request carrying a query string still matches and routes (`match` + // never sees the query). Preservation of the query itself is pinned by + // the `with_path_preserves_the_query_string` unit test — no proxy + // endpoint echoes its query back for an e2e-level assertion. const withQuery = await fetch(`${app.proxyUrl}/compat/health?probe=1`); expect(withQuery.status).toBe(200); });