From 3dab898976563d9a6ad0d1837f2478123752ba99 Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Thu, 2 Jul 2026 14:33:56 +0800 Subject: [PATCH 1/2] feat(mcp): api_key and OAuth2 client-credentials upstream auth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two new upstream auth types for registered MCP servers, alongside the existing none/bearer: - api_key: the stored secret is sent as `x-api-key` on every upstream request (via the transport's custom headers; the value is marked sensitive so header-map debug formatting can't echo it; a non-header- safe key fails with a clean error). - oauth2: OAuth 2.0 client credentials (RFC 6749 §4.4). New optional McpServer fields `client_id`, `token_url`, `scopes`; the existing `secret` carries the client secret. The gateway mints its own access token at the token endpoint and presents `Authorization: Bearer` — the caller's key is never forwarded upstream (the MCP authorization spec's no-passthrough requirement). Tokens are cached process-globally per (token_url, client_id, sha256(client_secret)) until 60s before expiry — a rotated secret can never reuse the old token — and the entry is invalidated when the upstream rejects with 401 (both rmcp 401 shapes matched, prefix-anchored so a response body can't spoof it) so the next attempt re-mints. Token minting is bounded by the same per-upstream deadline as the handshake. Schema stays flat and permissive on the cross-field coupling (no oneOf restructure); the runtime degrades a mis-configured oauth2 server gracefully — its tools drop out of tools/list, the failure is logged, nothing leaks to the caller — while this data plane's own Admin API enforces the coupling at write time (api_key ⇒ secret; oauth2 ⇒ client_id + secret + token_url). Secrets never appear in Debug output, logs, or error messages (token-endpoint errors report status only, so an identity provider echoing request parameters can't leak the secret through us). New fields are optional and skipped when unset, so existing stored resources and older control-plane payloads keep deserializing. Note the rmcp reqwest split: rmcp 1.8.0 links reqwest 0.13 while the workspace pins 0.12 — the 401 downcast names rmcp's own reqwest via a renamed dep (error type only); the live 401 tests fail loudly if the unification ever stops holding. Control-plane pairing (required for users to reach this): cp-admin.yaml enum + fields, Go model/validation/projection, and the dashboard auth subforms ship as the paired AISIX-Cloud PR (deploy the data plane first: its resource schema must know the new fields before the control plane projects them). Refs AISIX-Cloud#894 Revision C. --- Cargo.lock | 5 + .../aisix-admin/src/mcp_servers_handlers.rs | 94 +++- crates/aisix-admin/src/openapi.rs | 4 +- crates/aisix-core/src/models/mcp_server.rs | 91 +++- crates/aisix-core/src/models/schema.rs | 80 ++- crates/aisix-mcp/Cargo.toml | 16 + crates/aisix-mcp/src/bridge.rs | 171 +++++- crates/aisix-mcp/src/lib.rs | 3 +- crates/aisix-mcp/src/oauth.rs | 487 ++++++++++++++++++ crates/aisix-mcp/tests/gateway_aggregation.rs | 119 +++++ crates/aisix-mcp/tests/upstream_roundtrip.rs | 255 ++++++++- schemas/resources/mcp_server.schema.json | 42 +- 12 files changed, 1331 insertions(+), 36 deletions(-) create mode 100644 crates/aisix-mcp/src/oauth.rs diff --git a/Cargo.lock b/Cargo.lock index eab90297..898606ea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -161,9 +161,14 @@ dependencies = [ "async-trait", "axum", "futures", + "hex", + "http 1.4.0", + "reqwest 0.12.28", + "reqwest 0.13.4", "rmcp", "serde", "serde_json", + "sha2 0.10.9", "thiserror 1.0.69", "tokio", "tracing", diff --git a/crates/aisix-admin/src/mcp_servers_handlers.rs b/crates/aisix-admin/src/mcp_servers_handlers.rs index 23a5e725..d692b798 100644 --- a/crates/aisix-admin/src/mcp_servers_handlers.rs +++ b/crates/aisix-admin/src/mcp_servers_handlers.rs @@ -102,12 +102,34 @@ fn decode(raw: &Value) -> Result { "display_name must not contain the reserved separator `{TOOL_NAMESPACE_SEPARATOR}`" ))); } - if matches!(server.auth_type, McpAuthType::Bearer) - && server.secret.as_deref().unwrap_or_default().is_empty() - { - return Err(AdminError::BadRequest( - "secret is required and must be non-empty when auth_type is `bearer`".to_string(), - )); + // Per-auth_type credential coupling. The JSON schema stays flat and + // permissive on this (see the note on the McpServer struct); the write + // path is where an incomplete credential set is rejected outright. + let has_secret = !server.secret.as_deref().unwrap_or_default().is_empty(); + match server.auth_type { + McpAuthType::None => {} + McpAuthType::Bearer if !has_secret => { + return Err(AdminError::BadRequest( + "secret is required and must be non-empty when auth_type is `bearer`".to_string(), + )); + } + McpAuthType::ApiKey if !has_secret => { + return Err(AdminError::BadRequest( + "secret is required and must be non-empty when auth_type is `api_key`".to_string(), + )); + } + McpAuthType::OAuth2 => { + let has_client_id = !server.client_id.as_deref().unwrap_or_default().is_empty(); + let has_token_url = !server.token_url.as_deref().unwrap_or_default().is_empty(); + if !has_secret || !has_client_id || !has_token_url { + return Err(AdminError::BadRequest( + "client_id, token_url, and secret (the OAuth client secret) are required \ + and must be non-empty when auth_type is `oauth2`" + .to_string(), + )); + } + } + McpAuthType::Bearer | McpAuthType::ApiKey => {} } Ok(server) } @@ -148,6 +170,66 @@ mod tests { assert!(matches!(err, AdminError::BadRequest(_))); } + #[test] + fn decode_rejects_api_key_without_secret() { + let err = decode(&json!({ + "display_name": "gh", + "url": "https://x/mcp", + "auth_type": "api_key" + })) + .expect_err("api_key auth without a secret must be rejected"); + assert!(matches!(err, AdminError::BadRequest(_))); + } + + #[test] + fn decode_rejects_incomplete_oauth2() { + // Each of client_id / token_url / secret is individually required. + for missing in ["client_id", "token_url", "secret"] { + let mut v = json!({ + "display_name": "gh", + "url": "https://x/mcp", + "auth_type": "oauth2", + "client_id": "cid", + "token_url": "https://auth.example.com/oauth/token", + "secret": "cs" + }); + v.as_object_mut().unwrap().remove(missing); + let err = decode(&v).unwrap_err(); + assert!( + matches!(err, AdminError::BadRequest(_)), + "oauth2 without `{missing}` must be a BadRequest" + ); + } + } + + #[test] + fn decode_accepts_api_key_and_oauth2_servers() { + let api_key = decode(&json!({ + "display_name": "gh", + "url": "https://x/mcp", + "auth_type": "api_key", + "secret": "k-1" + })) + .expect("valid api_key server should decode"); + assert_eq!(api_key.secret.as_deref(), Some("k-1")); + + let oauth2 = decode(&json!({ + "display_name": "gh2", + "url": "https://x/mcp", + "auth_type": "oauth2", + "client_id": "cid", + "token_url": "https://auth.example.com/oauth/token", + "secret": "cs", + "scopes": ["read"] + })) + .expect("valid oauth2 server should decode"); + assert_eq!(oauth2.client_id.as_deref(), Some("cid")); + assert_eq!( + oauth2.token_url.as_deref(), + Some("https://auth.example.com/oauth/token") + ); + } + #[test] fn decode_accepts_valid_server() { let server = decode(&json!({ diff --git a/crates/aisix-admin/src/openapi.rs b/crates/aisix-admin/src/openapi.rs index 0b3148b7..7337e9d0 100644 --- a/crates/aisix-admin/src/openapi.rs +++ b/crates/aisix-admin/src/openapi.rs @@ -1417,7 +1417,7 @@ const OPENAPI_JSON_BASE: &str = r##"{ } }, "400": { - "description": "Schema validation failed, the JSON body is malformed, `display_name` contains the reserved `__` separator, or bearer auth is missing `secret`", + "description": "Schema validation failed, the JSON body is malformed, `display_name` contains the reserved `__` separator, or the credentials required by `auth_type` are missing (`secret` for `bearer`/`api_key`; `client_id`, `token_url`, and `secret` for `oauth2`)", "content": { "application/json": { "schema": { @@ -1575,7 +1575,7 @@ const OPENAPI_JSON_BASE: &str = r##"{ } }, "400": { - "description": "Schema validation failed, the JSON body is malformed, `display_name` contains the reserved `__` separator, or bearer auth is missing `secret`", + "description": "Schema validation failed, the JSON body is malformed, `display_name` contains the reserved `__` separator, or the credentials required by `auth_type` are missing (`secret` for `bearer`/`api_key`; `client_id`, `token_url`, and `secret` for `oauth2`)", "content": { "application/json": { "schema": { diff --git a/crates/aisix-core/src/models/mcp_server.rs b/crates/aisix-core/src/models/mcp_server.rs index 991a8848..e0104fd3 100644 --- a/crates/aisix-core/src/models/mcp_server.rs +++ b/crates/aisix-core/src/models/mcp_server.rs @@ -38,13 +38,39 @@ pub struct McpServer { #[serde(default)] pub auth_type: McpAuthType, - /// Authentication credential for the upstream server. Required when - /// `auth_type` is `bearer`, where it is sent as `Authorization: Bearer - /// ` on every upstream request. Leave unset when `auth_type` is - /// `none`. + /// Authentication credential for the upstream server. Its meaning follows + /// `auth_type`: the bearer token when `auth_type` is `bearer` (sent as + /// `Authorization: Bearer `), the API key when `auth_type` is + /// `api_key` (sent as `x-api-key: `), or the OAuth client secret + /// when `auth_type` is `oauth2`. Leave unset when `auth_type` is `none`. #[serde(default, skip_serializing_if = "Option::is_none")] pub secret: Option, + // Cross-field coupling (`oauth2` requires `client_id` + `secret` + + // `token_url`; `bearer`/`api_key` require `secret`) is deliberately NOT + // expressed in this flat schema — that would force restructuring the + // resource into a oneOf. The control plane enforces the coupling strictly + // at write time, this gateway's own Admin API re-checks it on write, and + // the runtime degrades gracefully when a snapshot-loaded server is + // mis-configured: its credential exchange fails, its tools become + // unavailable, and the failure is logged like any other upstream failure. + /// OAuth client identifier used for the OAuth 2.0 client credentials + /// grant. Required when `auth_type` is `oauth2`; ignored otherwise. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub client_id: Option, + + /// OAuth token endpoint URL where the gateway exchanges the client + /// credentials for an access token, such as + /// `https://auth.example.com/oauth/token`. Required when `auth_type` is + /// `oauth2`; ignored otherwise. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub token_url: Option, + + /// OAuth scopes to request. Joined with spaces into the `scope` parameter + /// of the token request. Only used when `auth_type` is `oauth2`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scopes: Option>, + /// Maximum time, in milliseconds, to wait for a single upstream operation /// (establishing the session, listing tools, or calling a tool). Must be at /// least `1` when set. When omitted, the gateway applies a built-in default. @@ -90,6 +116,16 @@ pub enum McpAuthType { /// Bearer token authentication. The token is supplied in `secret` and sent /// as `Authorization: Bearer `. Bearer, + /// API key authentication. The key is supplied in `secret` and sent as an + /// `x-api-key: ` header on every upstream request. + ApiKey, + /// OAuth 2.0 client credentials grant. The gateway exchanges `client_id`, + /// the client secret in `secret`, and the optional `scopes` at `token_url` + /// for an access token, and sends it as `Authorization: Bearer + /// ` on every upstream request. Access tokens are cached + /// until shortly before their reported expiry. + #[serde(rename = "oauth2")] + OAuth2, } impl Resource for McpServer { @@ -122,6 +158,9 @@ mod tests { assert_eq!(s.transport, McpTransport::StreamableHttp); assert_eq!(s.auth_type, McpAuthType::None); assert!(s.secret.is_none()); + assert!(s.client_id.is_none()); + assert!(s.token_url.is_none()); + assert!(s.scopes.is_none()); assert!(s.timeout_ms.is_none()); assert!(s.enabled); } @@ -138,6 +177,47 @@ mod tests { assert!(!s.enabled); } + #[test] + fn deserialises_with_api_key_auth() { + let s: McpServer = serde_json::from_str( + r#"{"display_name":"gh","url":"https://x/mcp","auth_type":"api_key","secret":"k-1"}"#, + ) + .unwrap(); + assert_eq!(s.auth_type, McpAuthType::ApiKey); + assert_eq!(s.secret.as_deref(), Some("k-1")); + } + + #[test] + fn deserialises_with_oauth2_auth() { + let s: McpServer = serde_json::from_str( + r#"{"display_name":"gh","url":"https://x/mcp","auth_type":"oauth2","secret":"cs-1","client_id":"cid","token_url":"https://auth/x/token","scopes":["read","write"]}"#, + ) + .unwrap(); + assert_eq!(s.auth_type, McpAuthType::OAuth2); + assert_eq!(s.secret.as_deref(), Some("cs-1")); + assert_eq!(s.client_id.as_deref(), Some("cid")); + assert_eq!(s.token_url.as_deref(), Some("https://auth/x/token")); + assert_eq!( + s.scopes.as_deref(), + Some(&["read".to_string(), "write".to_string()][..]) + ); + } + + #[test] + fn oauth2_round_trips_and_omits_unset_optionals() { + let original: McpServer = serde_json::from_str( + r#"{"display_name":"gh","url":"https://x/mcp","auth_type":"oauth2","secret":"cs","client_id":"cid","token_url":"https://auth/token"}"#, + ) + .unwrap(); + let s = serde_json::to_string(&original).unwrap(); + // The oauth2 tag serialises as `oauth2` (not a snake_cased `o_auth2`) + // and unset optionals (`scopes` here) are omitted entirely. + assert!(s.contains(r#""auth_type":"oauth2""#), "got: {s}"); + assert!(!s.contains("scopes"), "unset scopes must be omitted: {s}"); + let back: McpServer = serde_json::from_str(&s).unwrap(); + assert_eq!(original, back); + } + #[test] fn rejects_unknown_fields() { let r: Result = @@ -175,6 +255,9 @@ mod tests { transport: McpTransport::StreamableHttp, auth_type: McpAuthType::None, secret: None, + client_id: None, + token_url: None, + scopes: None, timeout_ms: None, enabled: true, runtime_id: String::new(), diff --git a/crates/aisix-core/src/models/schema.rs b/crates/aisix-core/src/models/schema.rs index d75b9254..8a11cb0d 100644 --- a/crates/aisix-core/src/models/schema.rs +++ b/crates/aisix-core/src/models/schema.rs @@ -183,18 +183,26 @@ pub fn provider_key_root_schema() -> Value { /// Canonical JSON Schema for the `mcp_server` resource, derived from the /// [`McpServer`](crate::models::McpServer) struct. Uses the nullable `Option` -/// representation (`true`) so the optional `secret` / `timeout_ms` fields accept -/// an explicit `null` as well as being absent, matching the resource's wire -/// contract. The `transport` / `auth_type` closed sets come from the +/// representation (`true`) so the optional fields (`secret`, `client_id`, +/// `token_url`, `scopes`, `timeout_ms`) accept an explicit `null` as well as +/// being absent, matching the resource's wire contract. The `transport` / +/// `auth_type` closed sets come from the /// [`McpTransport`](crate::models::McpTransport) / -/// [`McpAuthType`](crate::models::McpAuthType) enums. +/// [`McpAuthType`](crate::models::McpAuthType) enums. The per-`auth_type` +/// credential coupling is intentionally not encoded here (see the note on the +/// struct); the schema stays permissive and write paths enforce it. pub fn mcp_server_root_schema() -> Value { let mut schema = struct_root_schema::(true); if let Some(Value::Object(defs)) = schema.get_mut("definitions") { title_single_value_enum_variants( defs, "McpAuthType", - &[("none", "No authentication"), ("bearer", "Bearer token")], + &[ + ("none", "No authentication"), + ("bearer", "Bearer token"), + ("api_key", "API key"), + ("oauth2", "OAuth 2.0 client credentials"), + ], ); title_single_value_enum_variants( defs, @@ -2163,6 +2171,68 @@ mod tests { validate_mcp_server(&v).unwrap(); } + #[test] + fn mcp_server_accepts_api_key_auth() { + let v = json!({ + "display_name": "github", + "url": "https://api.example.com/mcp", + "auth_type": "api_key", + "secret": "k-123" + }); + validate_mcp_server(&v).unwrap(); + } + + #[test] + fn mcp_server_accepts_oauth2_auth_with_client_fields() { + let v = json!({ + "display_name": "github", + "url": "https://api.example.com/mcp", + "auth_type": "oauth2", + "secret": "client-secret", + "client_id": "cid", + "token_url": "https://auth.example.com/oauth/token", + "scopes": ["read", "write"] + }); + validate_mcp_server(&v).unwrap(); + } + + #[test] + fn mcp_server_rejects_unknown_auth_type_and_bad_scopes_shape() { + // The `auth_type` set is closed: near-misses like `oauth` must fail. + let v = json!({ + "display_name": "x", + "url": "https://x/mcp", + "auth_type": "oauth" + }); + assert!(validate_mcp_server(&v).is_err()); + + // `scopes` is an array of strings, not a single space-joined string. + let v = json!({ + "display_name": "x", + "url": "https://x/mcp", + "auth_type": "oauth2", + "secret": "s", + "client_id": "cid", + "token_url": "https://auth/token", + "scopes": "read write" + }); + assert!(validate_mcp_server(&v).is_err()); + } + + #[test] + fn mcp_server_schema_stays_permissive_on_credential_coupling() { + // The per-`auth_type` credential coupling (oauth2 ⇒ client_id + + // secret + token_url) is enforced by write paths, not this schema — + // an incomplete oauth2 row must still validate so the snapshot loader + // keeps it (the runtime degrades that server gracefully instead). + let v = json!({ + "display_name": "x", + "url": "https://x/mcp", + "auth_type": "oauth2" + }); + validate_mcp_server(&v).unwrap(); + } + #[test] fn mcp_server_rejects_zero_timeout_ms() { // A zero deadline times out every upstream op instantly and silently diff --git a/crates/aisix-mcp/Cargo.toml b/crates/aisix-mcp/Cargo.toml index d0ef00dd..fa420b2a 100644 --- a/crates/aisix-mcp/Cargo.toml +++ b/crates/aisix-mcp/Cargo.toml @@ -17,6 +17,22 @@ serde.workspace = true serde_json.workspace = true thiserror.workspace = true tracing.workspace = true +# Upstream OAuth 2.0 client-credentials token minting (`crate::oauth`) plus +# header types for the `api_key` transport headers. +http.workspace = true +reqwest.workspace = true +# rmcp 1.8.0 builds its transport on the reqwest 0.13 line (the workspace is +# on 0.12) and does not re-export it. This renamed dep exists ONLY to name +# `StreamableHttpError` when downcasting a connect failure to +# detect an upstream 401 — it resolves to the same crate version rmcp links, +# so the downcast sees the identical type. Kept `default-features = false`: +# only the `Error` type is used, never a client. The 401-invalidation +# integration test fails loudly if these ever stop unifying. +rmcp-reqwest = { package = "reqwest", version = "0.13", default-features = false } +# Token-cache keying: the client secret is folded into the cache key as a +# SHA-256 digest so a rotated secret never reuses the previous secret's token. +sha2.workspace = true +hex.workspace = true # Official MCP Rust SDK. Pinned exactly: rmcp is <16 months old and still # ships breaking changes on a roughly-monthly cadence, so we hold a fixed diff --git a/crates/aisix-mcp/src/bridge.rs b/crates/aisix-mcp/src/bridge.rs index 59114a6d..e95b5188 100644 --- a/crates/aisix-mcp/src/bridge.rs +++ b/crates/aisix-mcp/src/bridge.rs @@ -11,13 +11,17 @@ //! the rest of the data plane never depends on the SDK directly. That keeps //! rmcp's still-moving API contained to this file. +use std::collections::HashMap; use std::time::Duration; use aisix_core::{McpAuthType, McpServer}; use async_trait::async_trait; +use http::{HeaderName, HeaderValue}; use rmcp::model::CallToolRequestParams; -use rmcp::service::{RoleClient, RunningService}; -use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig; +use rmcp::service::{ClientInitializeError, RoleClient, RunningService}; +use rmcp::transport::streamable_http_client::{ + StreamableHttpClientTransportConfig, StreamableHttpError, +}; use rmcp::transport::StreamableHttpClientTransport; use rmcp::ServiceExt; @@ -29,11 +33,15 @@ use crate::error::McpError; /// indefinitely. Overridable per upstream via [`McpUpstream::with_timeout`]. pub const DEFAULT_UPSTREAM_TIMEOUT: Duration = Duration::from_secs(30); +/// Header carrying the gateway-held key for `api_key` upstream auth. +const API_KEY_HEADER: &str = "x-api-key"; + /// How the gateway authenticates to an upstream MCP server. The credential is /// held here on the gateway side and is never exposed to the calling agent — /// the agent presents only its AISIX key. The MCP authorization spec /// (2025-11-25) also requires that a downstream client token is never passed -/// through to the upstream; a Bearer set here is a distinct, gateway-held +/// through to the upstream; every credential set here — a Bearer, an API key, +/// or an OAuth token the gateway mints itself — is a distinct, gateway-held /// credential. #[derive(Clone)] pub enum McpAuth { @@ -42,6 +50,41 @@ pub enum McpAuth { /// Send `Authorization: Bearer ` on every upstream request. The /// token is the raw value, without the `Bearer ` prefix. Bearer(String), + /// Send `x-api-key: ` on every upstream request. + ApiKey(String), + /// OAuth 2.0 client credentials (RFC 6749 §4.4): mint an access token at + /// the configured token endpoint and send it as `Authorization: Bearer + /// `. The token is gateway-minted and gateway-held — never + /// the caller's credential (see [`crate::oauth`]). + OAuth2(OAuthClientConfig), +} + +/// Client-credentials parameters for [`McpAuth::OAuth2`]. +#[derive(Clone)] +pub struct OAuthClientConfig { + /// OAuth client identifier (non-secret). + pub client_id: String, + /// OAuth client secret. Redacted from `Debug` like every other + /// gateway-held credential in this module. + pub client_secret: String, + /// Token endpoint URL the credentials are exchanged at (non-secret). + pub token_url: String, + /// Scopes to request, joined with spaces into the `scope` parameter. + pub scopes: Vec, +} + +// Hand-written for the same reason as `McpAuth`'s: the client secret must +// never land in logs via `{:?}`. The non-secret fields stay visible — they +// are what an operator needs to identify the token exchange being logged. +impl std::fmt::Debug for OAuthClientConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OAuthClientConfig") + .field("client_id", &self.client_id) + .field("client_secret", &"***redacted***") + .field("token_url", &self.token_url) + .field("scopes", &self.scopes) + .finish() + } } // Hand-written so the gateway-held token never lands in logs via `{:?}`. This @@ -52,6 +95,9 @@ impl std::fmt::Debug for McpAuth { match self { McpAuth::None => f.write_str("None"), McpAuth::Bearer(_) => f.write_str("Bearer(***redacted***)"), + McpAuth::ApiKey(_) => f.write_str("ApiKey(***redacted***)"), + // Delegates to the redacting `OAuthClientConfig` impl above. + McpAuth::OAuth2(cfg) => f.debug_tuple("OAuth2").field(cfg).finish(), } } } @@ -96,6 +142,18 @@ impl McpUpstream { self } + /// Set API-key auth (sent as `x-api-key: `). + pub fn with_api_key(mut self, key: impl Into) -> Self { + self.auth = McpAuth::ApiKey(key.into()); + self + } + + /// Set OAuth 2.0 client-credentials auth. + pub fn with_oauth2(mut self, config: OAuthClientConfig) -> Self { + self.auth = McpAuth::OAuth2(config); + self + } + /// Override the per-operation deadline. pub fn with_timeout(mut self, timeout: Duration) -> Self { self.timeout = timeout; @@ -159,20 +217,58 @@ pub struct RmcpBridge { impl RmcpBridge { /// Open a session to `upstream`: build the Streamable HTTP transport - /// (injecting gateway-held auth) and run the `initialize` handshake, - /// bounded by the upstream's timeout. + /// (injecting gateway-held auth — for `oauth2` this mints or reuses a + /// cached access token first) and run the `initialize` handshake. The + /// whole sequence, token minting included, is bounded by the upstream's + /// timeout. pub async fn connect(upstream: &McpUpstream) -> Result { - let transport = match &upstream.auth { - McpAuth::None => StreamableHttpClientTransport::from_uri(upstream.url.clone()), - McpAuth::Bearer(token) => StreamableHttpClientTransport::from_config( - StreamableHttpClientTransportConfig::with_uri(upstream.url.clone()) - .auth_header(token.clone()), - ), + let establish = async { + let transport = match &upstream.auth { + McpAuth::None => StreamableHttpClientTransport::from_uri(upstream.url.clone()), + McpAuth::Bearer(token) => StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(upstream.url.clone()) + .auth_header(token.clone()), + ), + McpAuth::ApiKey(key) => { + // A key with non-header-safe bytes is a clean config error, + // not a panic — and the key itself never enters the message. + let mut value = HeaderValue::from_str(key).map_err(|_| { + McpError::Connect( + "upstream API key is not a valid HTTP header value".to_string(), + ) + })?; + // Marks the value opaque to `Debug` formatting of the + // header map, mirroring this module's redaction posture. + value.set_sensitive(true); + let headers = HashMap::from([(HeaderName::from_static(API_KEY_HEADER), value)]); + StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(upstream.url.clone()) + .custom_headers(headers), + ) + } + McpAuth::OAuth2(cfg) => { + let token = crate::oauth::get_or_fetch(cfg).await?; + StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(upstream.url.clone()) + .auth_header(token), + ) + } + }; + ().serve(transport).await.map_err(|e| { + // An upstream 401 against a minted token means the token was + // revoked or expired earlier than promised: drop the cache + // entry so the next attempt re-mints instead of replaying it. + if let McpAuth::OAuth2(cfg) = &upstream.auth { + if init_error_is_unauthorized(&e) { + crate::oauth::invalidate(cfg); + } + } + McpError::Connect(e.to_string()) + }) }; - let running = tokio::time::timeout(upstream.timeout, ().serve(transport)) + let running = tokio::time::timeout(upstream.timeout, establish) .await - .map_err(|_| McpError::Connect("upstream MCP connect timed out".to_string()))? - .map_err(|e| McpError::Connect(e.to_string()))?; + .map_err(|_| McpError::Connect("upstream MCP connect timed out".to_string()))??; Ok(Self { running, timeout: upstream.timeout, @@ -180,6 +276,36 @@ impl RmcpBridge { } } +/// Whether a failed `initialize` handshake was an upstream `401 Unauthorized`. +/// +/// The reqwest transport surfaces a 401 in one of two stable shapes (rmcp is +/// pinned exactly, so these cannot drift silently): a 401 carrying a +/// `WWW-Authenticate` header becomes `StreamableHttpError::AuthRequired`, and +/// any other non-success status becomes +/// `UnexpectedServerResponse("HTTP : …")`. Both arrive here inside +/// `ClientInitializeError::TransportError` as the type-erased transport error; +/// the downcast names rmcp's own reqwest (`rmcp_reqwest`, the 0.13 line — not +/// the workspace 0.12) so the types match. Post-handshake operations don't +/// need this: [`EphemeralBridge`] reconnects per operation, so every request +/// replays the handshake and a rejected token always surfaces on this path. +fn init_error_is_unauthorized(error: &ClientInitializeError) -> bool { + let ClientInitializeError::TransportError { error, .. } = error else { + return false; + }; + match error + .error + .downcast_ref::>() + { + Some(StreamableHttpError::AuthRequired(_)) => true, + Some(StreamableHttpError::UnexpectedServerResponse(message)) => { + // Anchored to the prefix: the format is `HTTP : `, + // so an upstream body can never fake a 401 here. + message.starts_with("HTTP 401") + } + _ => false, + } +} + #[async_trait] impl McpBridge for RmcpBridge { async fn list_tools(&self) -> Result, McpError> { @@ -229,12 +355,25 @@ fn into_mcp_tool(tool: rmcp::model::Tool) -> McpTool { } /// Build the connection parameters for an upstream from its registered -/// [`McpServer`] resource: maps `auth_type`/`secret` to [`McpAuth`] and -/// `timeout_ms` to the per-operation deadline. +/// [`McpServer`] resource: maps `auth_type` and its credential fields to +/// [`McpAuth`] and `timeout_ms` to the per-operation deadline. +/// +/// Stays permissive on purpose: fields a mis-configured resource left unset +/// map to empty strings rather than erroring here. The credential exchange +/// then fails cleanly at connect time and that server degrades like any +/// unreachable upstream (its tools drop out of `tools/list`, the failure is +/// logged), instead of one bad row poisoning snapshot loading. pub fn upstream_from_mcp_server(server: &McpServer) -> McpUpstream { let auth = match server.auth_type { McpAuthType::None => McpAuth::None, McpAuthType::Bearer => McpAuth::Bearer(server.secret.clone().unwrap_or_default()), + McpAuthType::ApiKey => McpAuth::ApiKey(server.secret.clone().unwrap_or_default()), + McpAuthType::OAuth2 => McpAuth::OAuth2(OAuthClientConfig { + client_id: server.client_id.clone().unwrap_or_default(), + client_secret: server.secret.clone().unwrap_or_default(), + token_url: server.token_url.clone().unwrap_or_default(), + scopes: server.scopes.clone().unwrap_or_default(), + }), }; let timeout = server .timeout_ms diff --git a/crates/aisix-mcp/src/lib.rs b/crates/aisix-mcp/src/lib.rs index 2dbbe296..3a346048 100644 --- a/crates/aisix-mcp/src/lib.rs +++ b/crates/aisix-mcp/src/lib.rs @@ -16,10 +16,11 @@ pub mod bridge; pub mod error; pub mod gateway; +mod oauth; pub use bridge::{ upstream_from_mcp_server, EphemeralBridge, McpAuth, McpBridge, McpTool, McpToolResult, - McpUpstream, RmcpBridge, + McpUpstream, OAuthClientConfig, RmcpBridge, }; pub use error::McpError; pub use gateway::{streamable_http_service, McpGateway, ToolAcl, TOOL_NAMESPACE_SEPARATOR}; diff --git a/crates/aisix-mcp/src/oauth.rs b/crates/aisix-mcp/src/oauth.rs new file mode 100644 index 00000000..b5453a68 --- /dev/null +++ b/crates/aisix-mcp/src/oauth.rs @@ -0,0 +1,487 @@ +//! OAuth 2.0 client-credentials token minting for upstream MCP servers. +//! +//! When an upstream is registered with `auth_type: oauth2`, the gateway mints +//! its own access token at the server's token endpoint (RFC 6749 §4.4, the +//! machine-to-machine grant) and presents it as `Authorization: Bearer +//! `. The token is a gateway-held credential: the calling +//! agent's AISIX key is never forwarded upstream, in line with the MCP +//! authorization spec's no-token-passthrough requirement. +//! +//! Tokens are cached process-globally per `(token_url, client_id, +//! client_secret)` triple and reused until shortly before their reported +//! expiry; [`invalidate`] drops an entry early when the upstream rejects the +//! token so the next attempt re-mints. + +use std::collections::HashMap; +use std::sync::{OnceLock, RwLock}; +use std::time::{Duration, Instant}; + +use serde::Deserialize; +use sha2::{Digest, Sha256}; + +use crate::bridge::{OAuthClientConfig, DEFAULT_UPSTREAM_TIMEOUT}; +use crate::error::McpError; + +/// Stop reusing a cached token this long before its reported expiry, so an +/// upstream operation never starts with a token about to lapse mid-flight. +const EXPIRY_SKEW: Duration = Duration::from_secs(60); + +/// Assumed lifetime when the token endpoint omits `expires_in` (RFC 6749 §5.1 +/// only recommends it). +const DEFAULT_TOKEN_LIFETIME: Duration = Duration::from_secs(3600); + +/// One minted upstream access token. Never printed: the struct deliberately +/// has no `Debug` impl, and nothing in this module logs the token value. +struct CachedToken { + access_token: String, + expires_at: Instant, +} + +/// Process-global token cache. Guards are held only for map lookups/inserts, +/// never across an await; concurrent misses for the same key may fetch in +/// parallel (each minted token is valid — last insert wins). +fn cache() -> &'static RwLock> { + static CACHE: OnceLock>> = OnceLock::new(); + CACHE.get_or_init(|| RwLock::new(HashMap::new())) +} + +/// Shared HTTP client for token fetches, bounded per request by the same +/// default deadline as any other upstream MCP operation. +fn http_client() -> &'static reqwest::Client { + static CLIENT: OnceLock = OnceLock::new(); + CLIENT.get_or_init(|| { + reqwest::Client::builder() + .timeout(DEFAULT_UPSTREAM_TIMEOUT) + .build() + // Building a client with only a timeout set cannot fail on any + // supported platform; fall back to the default client if it does. + .unwrap_or_default() + }) +} + +/// Cache key: `token_url \x1f client_id \x1f sha256(client_secret)`. The +/// secret is folded in as a digest — never in plaintext — so a rotated secret +/// can never reuse the previous secret's token, while the key itself stays +/// safe to hold alongside the non-secret fields. +fn cache_key(cfg: &OAuthClientConfig) -> String { + let secret_digest = hex::encode(Sha256::digest(cfg.client_secret.as_bytes())); + format!( + "{}\x1f{}\x1f{}", + cfg.token_url, cfg.client_id, secret_digest + ) +} + +/// Return the cached access token for `cfg`, minting a fresh one at the token +/// endpoint on a miss or when the cached token is within [`EXPIRY_SKEW`] of +/// expiry. +/// +/// A config missing `client_id`, the client secret, or `token_url` fails here +/// with a clean error (no panic): the mis-configured server simply becomes +/// unavailable, like any upstream that cannot be reached. +pub(crate) async fn get_or_fetch(cfg: &OAuthClientConfig) -> Result { + if cfg.token_url.is_empty() || cfg.client_id.is_empty() || cfg.client_secret.is_empty() { + return Err(McpError::Connect( + "oauth2 upstream auth requires client_id, secret (the OAuth client secret), \ + and token_url" + .to_string(), + )); + } + + let key = cache_key(cfg); + if let Some(token) = cached_token(&key) { + return Ok(token); + } + + let (access_token, lifetime) = fetch_token(cfg).await?; + let entry = CachedToken { + access_token: access_token.clone(), + expires_at: Instant::now() + lifetime, + }; + cache() + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .insert(key, entry); + Ok(access_token) +} + +/// Drop the cached token for `cfg`, forcing the next [`get_or_fetch`] to +/// re-mint. Called when the upstream rejects the presented token with `401` +/// (revoked, or expired earlier than `expires_in` promised). +pub(crate) fn invalidate(cfg: &OAuthClientConfig) { + cache() + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .remove(&cache_key(cfg)); +} + +fn cached_token(key: &str) -> Option { + let cache = cache() + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let entry = cache.get(key)?; + if entry.expires_at <= Instant::now() + EXPIRY_SKEW { + return None; + } + Some(entry.access_token.clone()) +} + +/// Token endpoint success payload (RFC 6749 §5.1). Unknown fields (`scope`, +/// `refresh_token`, …) are ignored; `token_type` is not needed — the gateway +/// always presents the token as `Bearer`. +#[derive(Deserialize)] +struct TokenResponse { + access_token: String, + #[serde(default)] + expires_in: Option, +} + +/// POST the client-credentials grant to the token endpoint and return the +/// minted `(access_token, lifetime)`. +/// +/// Error hygiene: the request carries the client secret, so the form body is +/// never logged or echoed into errors; error responses are reported by status +/// only (identity providers may reflect request parameters into their error +/// payloads). +async fn fetch_token(cfg: &OAuthClientConfig) -> Result<(String, Duration), McpError> { + let mut form: Vec<(&str, &str)> = vec![ + ("grant_type", "client_credentials"), + ("client_id", cfg.client_id.as_str()), + ("client_secret", cfg.client_secret.as_str()), + ]; + let joined_scopes; + if !cfg.scopes.is_empty() { + joined_scopes = cfg.scopes.join(" "); + form.push(("scope", joined_scopes.as_str())); + } + + let response = http_client() + .post(&cfg.token_url) + .form(&form) + .send() + .await + .map_err(|e| { + // reqwest transport errors describe the connection (and may name + // the non-secret token_url), never the form body. + McpError::Connect(format!("upstream OAuth token request failed: {e}")) + })?; + + let status = response.status(); + if !status.is_success() { + return Err(McpError::Connect(format!( + "upstream OAuth token endpoint returned HTTP {status}" + ))); + } + + let token: TokenResponse = response.json().await.map_err(|_| { + McpError::Connect( + "upstream OAuth token endpoint returned a malformed token response".to_string(), + ) + })?; + if token.access_token.is_empty() { + return Err(McpError::Connect( + "upstream OAuth token endpoint returned an empty access_token".to_string(), + )); + } + let lifetime = token + .expires_in + .map(Duration::from_secs) + .unwrap_or(DEFAULT_TOKEN_LIFETIME); + Ok((token.access_token, lifetime)) +} + +#[cfg(test)] +mod tests { + use std::net::SocketAddr; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex}; + + use super::*; + + /// How a test token endpoint answers. + #[derive(Clone)] + enum TokenEndpointBehavior { + /// `200` with `access_token: "tok-"` (`n` = 1-based request count) + /// and this `expires_in` (`None` → field omitted). + Mint { expires_in: Option }, + /// A fixed status + fixed body. + Static { + status: axum::http::StatusCode, + body: &'static str, + }, + } + + #[derive(Clone)] + struct TokenEndpointState { + behavior: TokenEndpointBehavior, + hits: Arc, + /// Every decoded form body, for asserting the request shape. + requests: Arc>>>, + } + + struct TokenEndpoint { + addr: SocketAddr, + hits: Arc, + requests: Arc>>>, + } + + impl TokenEndpoint { + fn url(&self) -> String { + format!("http://{}/oauth/token", self.addr) + } + + fn hits(&self) -> usize { + self.hits.load(Ordering::SeqCst) + } + + fn request(&self, index: usize) -> HashMap { + self.requests.lock().expect("requests lock")[index].clone() + } + } + + async fn handle_token_request( + axum::extract::State(state): axum::extract::State, + axum::extract::Form(form): axum::extract::Form>, + ) -> axum::response::Response { + use axum::response::IntoResponse; + let n = state.hits.fetch_add(1, Ordering::SeqCst) + 1; + state.requests.lock().expect("requests lock").push(form); + match &state.behavior { + TokenEndpointBehavior::Mint { expires_in } => { + let mut body = serde_json::json!({ + "access_token": format!("tok-{n}"), + "token_type": "Bearer", + }); + if let Some(expires_in) = expires_in { + body["expires_in"] = (*expires_in).into(); + } + axum::Json(body).into_response() + } + TokenEndpointBehavior::Static { status, body } => (*status, *body).into_response(), + } + } + + /// Stand up a real token endpoint on an ephemeral port. Each test gets + /// its own port, so `token_url`-keyed cache entries never collide across + /// tests sharing the process-global cache. + async fn spawn_token_endpoint(behavior: TokenEndpointBehavior) -> TokenEndpoint { + let state = TokenEndpointState { + behavior, + hits: Arc::new(AtomicUsize::new(0)), + requests: Arc::new(Mutex::new(Vec::new())), + }; + let hits = state.hits.clone(); + let requests = state.requests.clone(); + let app = axum::Router::new() + .route("/oauth/token", axum::routing::post(handle_token_request)) + .with_state(state); + 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"); + }); + TokenEndpoint { + addr, + hits, + requests, + } + } + + fn config(token_url: String) -> OAuthClientConfig { + OAuthClientConfig { + client_id: "cid".to_string(), + client_secret: "s3cret".to_string(), + token_url, + scopes: Vec::new(), + } + } + + #[tokio::test] + async fn mints_token_with_client_credentials_form() { + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Mint { + expires_in: Some(3600), + }) + .await; + let cfg = config(endpoint.url()); + + let token = get_or_fetch(&cfg).await.expect("mint token"); + assert_eq!(token, "tok-1"); + + // The grant is the urlencoded client-credentials shape (RFC 6749 + // §4.4.2), with no `scope` parameter when no scopes are configured. + let form = endpoint.request(0); + assert_eq!( + form.get("grant_type").map(String::as_str), + Some("client_credentials") + ); + assert_eq!(form.get("client_id").map(String::as_str), Some("cid")); + assert_eq!( + form.get("client_secret").map(String::as_str), + Some("s3cret") + ); + assert!( + !form.contains_key("scope"), + "no scope param when scopes are empty" + ); + } + + #[tokio::test] + async fn second_call_within_expiry_hits_the_cache() { + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Mint { + expires_in: Some(3600), + }) + .await; + let cfg = config(endpoint.url()); + + assert_eq!(get_or_fetch(&cfg).await.expect("first"), "tok-1"); + assert_eq!(get_or_fetch(&cfg).await.expect("second"), "tok-1"); + assert_eq!(endpoint.hits(), 1, "second call must be served from cache"); + } + + #[tokio::test] + async fn missing_expires_in_defaults_to_a_cacheable_lifetime() { + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Mint { expires_in: None }).await; + let cfg = config(endpoint.url()); + + assert_eq!(get_or_fetch(&cfg).await.expect("first"), "tok-1"); + assert_eq!(get_or_fetch(&cfg).await.expect("second"), "tok-1"); + assert_eq!(endpoint.hits(), 1, "default lifetime must allow caching"); + } + + #[tokio::test] + async fn token_expiring_within_the_skew_is_refetched() { + // `expires_in` below EXPIRY_SKEW: valid to hand out now, but stale + // for reuse — the second call must mint a fresh token (no sleeping). + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Mint { + expires_in: Some(10), + }) + .await; + let cfg = config(endpoint.url()); + + assert_eq!(get_or_fetch(&cfg).await.expect("first"), "tok-1"); + assert_eq!(get_or_fetch(&cfg).await.expect("second"), "tok-2"); + assert_eq!(endpoint.hits(), 2, "near-expiry token must be refetched"); + } + + #[tokio::test] + async fn rotated_client_secret_never_reuses_the_old_token() { + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Mint { + expires_in: Some(3600), + }) + .await; + let cfg = config(endpoint.url()); + assert_eq!(get_or_fetch(&cfg).await.expect("first"), "tok-1"); + + let rotated = OAuthClientConfig { + client_secret: "rotated".to_string(), + ..cfg.clone() + }; + assert_eq!( + get_or_fetch(&rotated).await.expect("rotated"), + "tok-2", + "a rotated secret is a different cache key and must re-mint" + ); + assert_eq!(endpoint.hits(), 2); + + // The original secret's entry is untouched. + assert_eq!(get_or_fetch(&cfg).await.expect("original"), "tok-1"); + assert_eq!(endpoint.hits(), 2); + } + + #[tokio::test] + async fn invalidate_forces_a_refetch() { + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Mint { + expires_in: Some(3600), + }) + .await; + let cfg = config(endpoint.url()); + + assert_eq!(get_or_fetch(&cfg).await.expect("first"), "tok-1"); + invalidate(&cfg); + assert_eq!(get_or_fetch(&cfg).await.expect("after invalidate"), "tok-2"); + assert_eq!(endpoint.hits(), 2); + } + + #[tokio::test] + async fn scopes_are_joined_with_spaces() { + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Mint { + expires_in: Some(3600), + }) + .await; + let cfg = OAuthClientConfig { + scopes: vec!["read".to_string(), "write".to_string()], + ..config(endpoint.url()) + }; + + get_or_fetch(&cfg).await.expect("mint token"); + let form = endpoint.request(0); + assert_eq!(form.get("scope").map(String::as_str), Some("read write")); + } + + #[tokio::test] + async fn error_statuses_and_malformed_bodies_fail_cleanly() { + // 500 with a body that must never be echoed into the error. + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Static { + status: axum::http::StatusCode::INTERNAL_SERVER_ERROR, + body: "identity provider exploded: client_secret=echoed-back", + }) + .await; + let cfg = config(endpoint.url()); + let err = get_or_fetch(&cfg).await.expect_err("500 must fail"); + let msg = err.to_string(); + assert!(msg.contains("HTTP 500"), "status should be reported: {msg}"); + assert!( + !msg.contains("echoed-back") && !msg.contains("s3cret"), + "neither the response body nor the secret may leak: {msg}" + ); + + // 200 but not JSON. + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Static { + status: axum::http::StatusCode::OK, + body: "not json", + }) + .await; + let err = get_or_fetch(&config(endpoint.url())) + .await + .expect_err("malformed body must fail"); + assert!(err.to_string().contains("malformed"), "got: {err}"); + + // 200 JSON but no access_token. + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Static { + status: axum::http::StatusCode::OK, + body: r#"{"token_type":"Bearer"}"#, + }) + .await; + let err = get_or_fetch(&config(endpoint.url())) + .await + .expect_err("missing access_token must fail"); + assert!(err.to_string().contains("malformed"), "got: {err}"); + } + + #[tokio::test] + async fn incomplete_config_fails_without_contacting_anything() { + for broken in [ + OAuthClientConfig { + token_url: String::new(), + ..config("ignored".to_string()) + }, + OAuthClientConfig { + client_id: String::new(), + ..config("http://127.0.0.1:9/oauth/token".to_string()) + }, + OAuthClientConfig { + client_secret: String::new(), + ..config("http://127.0.0.1:9/oauth/token".to_string()) + }, + ] { + let err = get_or_fetch(&broken) + .await + .expect_err("incomplete oauth2 config must fail cleanly"); + assert!( + err.to_string().contains("oauth2"), + "error should name the misconfiguration: {err}" + ); + } + } +} diff --git a/crates/aisix-mcp/tests/gateway_aggregation.rs b/crates/aisix-mcp/tests/gateway_aggregation.rs index d529409e..5652d60f 100644 --- a/crates/aisix-mcp/tests/gateway_aggregation.rs +++ b/crates/aisix-mcp/tests/gateway_aggregation.rs @@ -315,6 +315,59 @@ fn upstream_from_mcp_server_maps_auth_and_timeout() { )); } +#[test] +fn upstream_from_mcp_server_maps_api_key_and_oauth2() { + let api_key: McpServer = serde_json::from_value(serde_json::json!({ + "display_name": "gh", + "url": "https://api.example.com/mcp", + "auth_type": "api_key", + "secret": "k-1" + })) + .unwrap(); + assert!(matches!( + upstream_from_mcp_server(&api_key).auth, + McpAuth::ApiKey(ref k) if k == "k-1" + )); + + let oauth: McpServer = serde_json::from_value(serde_json::json!({ + "display_name": "gh2", + "url": "https://api.example.com/mcp", + "auth_type": "oauth2", + "secret": "cs", + "client_id": "cid", + "token_url": "https://auth.example.com/token", + "scopes": ["read", "write"] + })) + .unwrap(); + match upstream_from_mcp_server(&oauth).auth { + McpAuth::OAuth2(cfg) => { + assert_eq!(cfg.client_id, "cid"); + assert_eq!(cfg.client_secret, "cs"); + assert_eq!(cfg.token_url, "https://auth.example.com/token"); + assert_eq!(cfg.scopes, vec!["read".to_string(), "write".to_string()]); + } + other => panic!("expected OAuth2 auth, got {other:?}"), + } + + // Missing oauth2 fields map to empty strings — the mapping never errors; + // the token fetch fails cleanly at connect time instead. + let incomplete: McpServer = serde_json::from_value(serde_json::json!({ + "display_name": "gh3", + "url": "https://api.example.com/mcp", + "auth_type": "oauth2" + })) + .unwrap(); + match upstream_from_mcp_server(&incomplete).auth { + McpAuth::OAuth2(cfg) => { + assert!(cfg.client_id.is_empty()); + assert!(cfg.client_secret.is_empty()); + assert!(cfg.token_url.is_empty()); + assert!(cfg.scopes.is_empty()); + } + other => panic!("expected OAuth2 auth, got {other:?}"), + } +} + /// 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!({ @@ -365,6 +418,72 @@ async fn from_snapshot_sources_only_enabled_upstreams() { assert_eq!(first_text(&result), "alpha:hi"); } +#[tokio::test] +async fn from_snapshot_degrades_misconfigured_oauth2_upstream_gracefully() { + let alpha = spawn_upstream("alpha").await; + + let snapshot = AisixSnapshot::new(); + snapshot + .mcp_servers + .insert(mcp_entry("e1", "alpha", &alpha, true)); + // An `oauth2` server missing its `token_url`: schema-valid and loadable + // (the flat schema stays permissive on credential coupling), but its + // token fetch can never succeed. + let broken: McpServer = serde_json::from_value(serde_json::json!({ + "display_name": "broken", + "url": format!("http://{alpha}/mcp"), + "auth_type": "oauth2", + "client_id": "cid", + "secret": "top-secret-cs", + "enabled": true + })) + .unwrap(); + snapshot + .mcp_servers + .insert(ResourceEntry::new("e2", broken, 1)); + + let gw_addr = spawn_gateway(McpGateway::from_snapshot(&snapshot)).await; + let client = () + .serve(StreamableHttpClientTransport::from_uri(format!( + "http://{gw_addr}/mcp" + ))) + .await + .expect("connect"); + + // The mis-configured upstream's tools are simply absent (its failure is + // logged server-side); the healthy upstream keeps serving. + 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"], + "misconfigured oauth2 upstream must be skipped, not fatal to the aggregate" + ); + + // Calling into it fails with the generic upstream-failure message — no + // hang, and nothing about its credentials leaks to the agent. + let err = client + .call_tool(call("broken__echo", "hi")) + .await + .expect_err("a call through the misconfigured upstream must fail"); + let msg = format!("{err:?}"); + assert!( + msg.contains("failed to call tool"), + "expected the generic upstream-failure error, got: {msg}" + ); + assert!( + !msg.contains("top-secret-cs") && !msg.contains("oauth"), + "credential/config detail must not reach the agent: {msg}" + ); + + // The healthy upstream still routes. + let ok = client + .call_tool(call("alpha__echo", "hi")) + .await + .expect("healthy upstream keeps serving"); + assert_eq!(first_text(&ok), "alpha:hi"); +} + #[test] fn tool_acl_from_allowed_semantics() { // No allowed_tools / empty → deny all. diff --git a/crates/aisix-mcp/tests/upstream_roundtrip.rs b/crates/aisix-mcp/tests/upstream_roundtrip.rs index 172ecc42..d77bce15 100644 --- a/crates/aisix-mcp/tests/upstream_roundtrip.rs +++ b/crates/aisix-mcp/tests/upstream_roundtrip.rs @@ -9,9 +9,10 @@ //! rejected), per the MCP authorization no-passthrough model. use std::net::SocketAddr; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; -use aisix_mcp::{McpBridge, McpUpstream, RmcpBridge}; +use aisix_mcp::{EphemeralBridge, McpBridge, McpUpstream, OAuthClientConfig, RmcpBridge}; use rmcp::model::{ CallToolRequestParams, CallToolResult, Content, ListToolsResult, PaginatedRequestParams, ServerCapabilities, ServerInfo, Tool, @@ -94,6 +95,94 @@ async fn require_bearer( next.run(request).await } +/// Exact-header gate for the richer auth types: reject any request that does +/// not carry `name: expected` with a `401`. `www_authenticate` toggles the +/// `WWW-Authenticate` challenge on the rejection — the rmcp client maps a 401 +/// with and without the challenge to two different error shapes, and the +/// token-invalidation path must handle both. +#[derive(Clone)] +struct RequiredHeader { + name: &'static str, + expected: String, + www_authenticate: bool, +} + +async fn require_exact_header( + axum::extract::State(required): axum::extract::State, + request: axum::extract::Request, + next: axum::middleware::Next, +) -> axum::response::Response { + use axum::response::IntoResponse; + let presented = request + .headers() + .get(required.name) + .and_then(|v| v.to_str().ok()); + if presented != Some(required.expected.as_str()) { + let mut response = + (axum::http::StatusCode::UNAUTHORIZED, "credential rejected").into_response(); + if required.www_authenticate { + response.headers_mut().insert( + axum::http::header::WWW_AUTHENTICATE, + axum::http::HeaderValue::from_static("Bearer realm=\"mcp\""), + ); + } + return response; + } + next.run(request).await +} + +/// Start an echo server that requires an exact header on every request. +async fn spawn_echo_server_requiring_header(required: RequiredHeader) -> SocketAddr { + let service = StreamableHttpService::new( + || Ok(EchoServer), + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default(), + ); + let app = axum::Router::new().nest_service("/mcp", service).layer( + axum::middleware::from_fn_with_state(required, require_exact_header), + ); + 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 +} + +/// A minimal OAuth token endpoint minting `tok-` (`n` = 1-based hit count) +/// with a long `expires_in`. Returns its URL and the hit counter, so tests can +/// assert exactly how many times a token was (re-)minted. +async fn spawn_token_endpoint() -> (String, Arc) { + use axum::response::IntoResponse; + let hits = Arc::new(AtomicUsize::new(0)); + let counter = hits.clone(); + let app = axum::Router::new().route( + "/oauth/token", + axum::routing::post(move || { + let counter = counter.clone(); + async move { + let n = counter.fetch_add(1, Ordering::SeqCst) + 1; + axum::Json(serde_json::json!({ + "access_token": format!("tok-{n}"), + "token_type": "Bearer", + "expires_in": 3600, + })) + .into_response() + } + }), + ); + 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"); + }); + (format!("http://{addr}/oauth/token"), hits) +} + /// Start the echo server on an ephemeral port; return its bound address. async fn spawn_echo_server(require_bearer_token: Option<&str>) -> SocketAddr { let service = StreamableHttpService::new( @@ -181,6 +270,170 @@ async fn forwards_gateway_held_bearer_to_upstream() { assert_eq!(tools[0].name, "echo"); } +#[tokio::test] +async fn forwards_gateway_held_api_key_to_upstream() { + let addr = spawn_echo_server_requiring_header(RequiredHeader { + name: "x-api-key", + expected: "k-123".to_string(), + www_authenticate: false, + }) + .await; + let url = format!("http://{addr}/mcp"); + + // Without the gateway-held key, the upstream rejects the session. + let unauth = RmcpBridge::connect(&McpUpstream::new(url.clone())).await; + assert!( + unauth.is_err(), + "connect without the API key must fail against a key-required upstream" + ); + + // With it, the session establishes — proving `x-api-key` is sent on + // every upstream request. + let upstream = McpUpstream::new(url).with_api_key("k-123"); + let bridge = RmcpBridge::connect(&upstream) + .await + .expect("connect with gateway-held API key"); + let tools = bridge.list_tools().await.expect("list tools"); + assert_eq!(tools.len(), 1); + assert_eq!(tools[0].name, "echo"); +} + +#[tokio::test] +async fn api_key_with_invalid_header_bytes_fails_cleanly() { + let addr = spawn_echo_server(None).await; + // A newline is not a valid HTTP header byte: the connect must return a + // clean config error (no panic) and must not echo the key material. + let upstream = McpUpstream::new(format!("http://{addr}/mcp")).with_api_key("bad\nkey"); + let err = match RmcpBridge::connect(&upstream).await { + Ok(_) => panic!("invalid header bytes must fail cleanly, not connect"), + Err(err) => err, + }; + let msg = err.to_string(); + assert!( + msg.contains("not a valid HTTP header value"), + "expected the config-error message, got: {msg}" + ); + assert!(!msg.contains("bad\n"), "key material must not leak: {msg}"); +} + +#[tokio::test] +async fn oauth2_mints_token_and_reuses_it_across_operations() { + let (token_url, mints) = spawn_token_endpoint().await; + // The upstream accepts exactly the first minted token — so a passing + // list/call proves the gateway attached `Authorization: Bearer tok-1`. + let addr = spawn_echo_server(Some("tok-1")).await; + + let upstream = McpUpstream::new(format!("http://{addr}/mcp")).with_oauth2(OAuthClientConfig { + client_id: "cid-roundtrip".to_string(), + client_secret: "cs".to_string(), + token_url, + scopes: Vec::new(), + }); + + // EphemeralBridge reconnects per operation: the second operation must + // reuse the cached token instead of minting a new one. + let bridge = EphemeralBridge::new(upstream); + let tools = bridge.list_tools().await.expect("list via minted token"); + assert_eq!(tools.len(), 1); + assert_eq!(tools[0].name, "echo"); + let result = bridge + .call_tool("echo", serde_json::json!({ "text": "hello oauth" })) + .await + .expect("call via cached token"); + assert_eq!(result.content[0]["text"], "hello oauth"); + assert_eq!( + mints.load(Ordering::SeqCst), + 1, + "the token must be minted once and then served from the cache" + ); +} + +#[tokio::test] +async fn upstream_401_with_challenge_invalidates_the_cached_token() { + let (token_url, mints) = spawn_token_endpoint().await; + // This upstream never accepts our minted tokens: every request gets a + // 401 WITH a `WWW-Authenticate` challenge (rmcp's `AuthRequired` shape). + let addr = spawn_echo_server_requiring_header(RequiredHeader { + name: "authorization", + expected: "Bearer some-other-token".to_string(), + www_authenticate: true, + }) + .await; + + let upstream = McpUpstream::new(format!("http://{addr}/mcp")).with_oauth2(OAuthClientConfig { + client_id: "cid-401-challenge".to_string(), + client_secret: "cs".to_string(), + token_url, + scopes: Vec::new(), + }); + let bridge = EphemeralBridge::new(upstream); + + assert!(bridge.list_tools().await.is_err()); + assert!(bridge.list_tools().await.is_err()); + assert_eq!( + mints.load(Ordering::SeqCst), + 2, + "each upstream 401 must invalidate the cached token so the next attempt re-mints" + ); +} + +#[tokio::test] +async fn upstream_401_without_challenge_also_invalidates_the_cached_token() { + let (token_url, mints) = spawn_token_endpoint().await; + // Same rejection, but WITHOUT `WWW-Authenticate` — rmcp surfaces this as + // a generic `HTTP 401` response error, the other shape the invalidation + // path must recognise. + let addr = spawn_echo_server_requiring_header(RequiredHeader { + name: "authorization", + expected: "Bearer some-other-token".to_string(), + www_authenticate: false, + }) + .await; + + let upstream = McpUpstream::new(format!("http://{addr}/mcp")).with_oauth2(OAuthClientConfig { + client_id: "cid-401-bare".to_string(), + client_secret: "cs".to_string(), + token_url, + scopes: Vec::new(), + }); + let bridge = EphemeralBridge::new(upstream); + + assert!(bridge.list_tools().await.is_err()); + assert!(bridge.list_tools().await.is_err()); + assert_eq!( + mints.load(Ordering::SeqCst), + 2, + "a bare upstream 401 must also invalidate the cached token" + ); +} + +#[tokio::test] +async fn misconfigured_oauth2_upstream_fails_cleanly_without_leaking() { + let addr = spawn_echo_server(None).await; + // `token_url` missing — the canonical mis-configured row the flat schema + // deliberately lets through. The operation must fail with a clean error, + // not a panic, and never surface the client secret. + let upstream = McpUpstream::new(format!("http://{addr}/mcp")).with_oauth2(OAuthClientConfig { + client_id: "cid-misconfig".to_string(), + client_secret: "super-secret".to_string(), + token_url: String::new(), + scopes: Vec::new(), + }); + let err = EphemeralBridge::new(upstream) + .list_tools() + .await + .expect_err("a token fetch with no token_url must fail"); + let msg = err.to_string(); + assert!( + msg.contains("oauth2"), + "error should name the misconfiguration: {msg}" + ); + assert!( + !msg.contains("super-secret"), + "the client secret must never leak: {msg}" + ); +} + #[tokio::test] async fn upstream_call_times_out_instead_of_hanging() { let addr = spawn_echo_server(None).await; diff --git a/schemas/resources/mcp_server.schema.json b/schemas/resources/mcp_server.schema.json index 2fc64479..60a19adf 100644 --- a/schemas/resources/mcp_server.schema.json +++ b/schemas/resources/mcp_server.schema.json @@ -20,6 +20,22 @@ ], "title": "Bearer token", "type": "string" + }, + { + "description": "API key authentication. The key is supplied in `secret` and sent as an `x-api-key: ` header on every upstream request.", + "enum": [ + "api_key" + ], + "title": "API key", + "type": "string" + }, + { + "description": "OAuth 2.0 client credentials grant. The gateway exchanges `client_id`, the client secret in `secret`, and the optional `scopes` at `token_url` for an access token, and sends it as `Authorization: Bearer ` on every upstream request. Access tokens are cached until shortly before their reported expiry.", + "enum": [ + "oauth2" + ], + "title": "OAuth 2.0 client credentials", + "type": "string" } ] }, @@ -47,6 +63,13 @@ "default": "none", "description": "How the gateway authenticates to the upstream server. The credential is held by the gateway and is never forwarded from or exposed to the calling client." }, + "client_id": { + "description": "OAuth client identifier used for the OAuth 2.0 client credentials grant. Required when `auth_type` is `oauth2`; ignored otherwise.", + "type": [ + "string", + "null" + ] + }, "display_name": { "description": "Operator-facing label, unique within the gateway. It is used as the namespace prefix for this server's tools, which are exposed to clients as `__`, so it must not contain the reserved separator `__`.", "minLength": 1, @@ -57,8 +80,18 @@ "description": "Whether this server is active. When `false`, its tools are not listed and cannot be called.", "type": "boolean" }, + "scopes": { + "description": "OAuth scopes to request. Joined with spaces into the `scope` parameter of the token request. Only used when `auth_type` is `oauth2`.", + "items": { + "type": "string" + }, + "type": [ + "array", + "null" + ] + }, "secret": { - "description": "Authentication credential for the upstream server. Required when `auth_type` is `bearer`, where it is sent as `Authorization: Bearer ` on every upstream request. Leave unset when `auth_type` is `none`.", + "description": "Authentication credential for the upstream server. Its meaning follows `auth_type`: the bearer token when `auth_type` is `bearer` (sent as `Authorization: Bearer `), the API key when `auth_type` is `api_key` (sent as `x-api-key: `), or the OAuth client secret when `auth_type` is `oauth2`. Leave unset when `auth_type` is `none`.", "type": [ "string", "null" @@ -73,6 +106,13 @@ "null" ] }, + "token_url": { + "description": "OAuth token endpoint URL where the gateway exchanges the client credentials for an access token, such as `https://auth.example.com/oauth/token`. Required when `auth_type` is `oauth2`; ignored otherwise.", + "type": [ + "string", + "null" + ] + }, "transport": { "allOf": [ { From 32a0a954bc38ba6b6e83e3b1f62c28f13ec8b7b0 Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Thu, 2 Jul 2026 14:54:37 +0800 Subject: [PATCH 2/2] fix(mcp): harden the oauth2 token lifecycle per audit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Independent audit findings (no HIGH; three MEDIUMs, all fixed): - Clamp a reported expires_in to 30 days: a hostile or broken identity provider returning u64::MAX would overflow Instant + Duration and panic the connect task. Test pins u64::MAX minting + caching cleanly. - Fold the requested scopes into the token-cache key: two servers sharing one OAuth client but requesting different scopes must never share a token (one would be presented a token minted for the other's scopes — silent over-scoping, or invalidation thrash). Components are now length-prefixed too, so no crafted field value can shift a key boundary. Test pins same-client/different-scopes isolation. - Refuse redirects on the token fetch: a token endpoint never legitimately redirects, and following one re-POSTs the secret-bearing form to wherever it points (307/308 keep the body; a chain may even downgrade to plain HTTP). Test pins that a 307 fails without being followed. LOWs: pin the Debug redaction of every credential variant with a test (it was the only guard and untested); sanitize the connect error text (strip control characters, truncate) so an upstream that reflects request headers into a 401 body can't plant them verbatim in gateway logs; document the third 401 shape (JSON-RPC-parseable body arrives as JsonRpcError, availability-only, bounded by expiry skew) and the cache's no-eviction bound. client_secret_basic interop filed as a follow-up issue. --- crates/aisix-mcp/src/bridge.rs | 77 +++++++++++++++++++++- crates/aisix-mcp/src/oauth.rs | 114 +++++++++++++++++++++++++++++---- 2 files changed, 179 insertions(+), 12 deletions(-) diff --git a/crates/aisix-mcp/src/bridge.rs b/crates/aisix-mcp/src/bridge.rs index e95b5188..01dbc4c2 100644 --- a/crates/aisix-mcp/src/bridge.rs +++ b/crates/aisix-mcp/src/bridge.rs @@ -263,7 +263,12 @@ impl RmcpBridge { crate::oauth::invalidate(cfg); } } - McpError::Connect(e.to_string()) + // Bound + sanitize: a bare-401 shape embeds the upstream's + // response body in the error text, which lands in gateway + // logs. An upstream that reflects request headers into its + // error body (or emits control characters for log injection) + // must not get either past this point verbatim. + McpError::Connect(sanitize_error_message(&e.to_string())) }) }; let running = tokio::time::timeout(upstream.timeout, establish) @@ -276,6 +281,23 @@ impl RmcpBridge { } } +/// Bound an upstream-derived error message for logging: control characters +/// (log-injection vectors) are stripped and the text is truncated, since a +/// bare non-success response embeds the upstream's body verbatim. +fn sanitize_error_message(message: &str) -> String { + const MAX_LEN: usize = 256; + let cleaned: String = message + .chars() + .map(|c| if c.is_control() { ' ' } else { c }) + .collect(); + if cleaned.chars().count() <= MAX_LEN { + cleaned + } else { + let truncated: String = cleaned.chars().take(MAX_LEN).collect(); + format!("{truncated}…") + } +} + /// Whether a failed `initialize` handshake was an upstream `401 Unauthorized`. /// /// The reqwest transport surfaces a 401 in one of two stable shapes (rmcp is @@ -288,6 +310,13 @@ impl RmcpBridge { /// the workspace 0.12) so the types match. Post-handshake operations don't /// need this: [`EphemeralBridge`] reconnects per operation, so every request /// replays the handshake and a rejected token always surfaces on this path. +/// +/// Known gap (availability-only): a bare 401 whose body parses as a JSON-RPC +/// error arrives as `ClientInitializeError::JsonRpcError` — not a transport +/// error — and is not recognized here, so a revoked token is replayed until +/// the cache's expiry skew retires it (at most ~59 minutes at the default +/// lifetime). Spec-conforming servers answer 401 with `WWW-Authenticate`, +/// which IS recognized. fn init_error_is_unauthorized(error: &ClientInitializeError) -> bool { let ClientInitializeError::TransportError { error, .. } = error else { return false; @@ -425,3 +454,49 @@ impl McpBridge for EphemeralBridge { .await } } + +#[cfg(test)] +mod tests { + use super::*; + + /// The hand-written Debug impls are the only guard between a credential + /// and the logs; pin them so a `#[derive(Debug)]` regression fails loudly + /// for every secret-bearing variant. + #[test] + fn debug_redacts_every_credential_variant() { + let oauth = McpAuth::OAuth2(OAuthClientConfig { + client_id: "cid".into(), + client_secret: "cs-LEAK".into(), + token_url: "https://idp.example.com/token".into(), + scopes: vec!["read".into()], + }); + let rendered = format!( + "{oauth:?} {:?} {:?}", + McpAuth::ApiKey("key-LEAK".into()), + McpAuth::Bearer("tok-LEAK".into()) + ); + assert!( + !rendered.contains("LEAK"), + "credential leaked into Debug output: {rendered}" + ); + // The non-secret fields stay visible for operability. + assert!(rendered.contains("idp.example.com")); + assert!(rendered.contains("cid")); + } + + #[test] + fn sanitize_error_message_strips_controls_and_truncates() { + let injected = "HTTP 401: bad\r\n[FAKE LOG LINE] evil"; + let cleaned = sanitize_error_message(injected); + assert!( + !cleaned.contains('\n') && !cleaned.contains('\r'), + "{cleaned}" + ); + assert!(cleaned.starts_with("HTTP 401")); + + let long = format!("HTTP 502: {}", "x".repeat(1000)); + let truncated = sanitize_error_message(&long); + assert!(truncated.chars().count() <= 257, "bounded output"); + assert!(truncated.ends_with('…')); + } +} diff --git a/crates/aisix-mcp/src/oauth.rs b/crates/aisix-mcp/src/oauth.rs index b5453a68..8f81a363 100644 --- a/crates/aisix-mcp/src/oauth.rs +++ b/crates/aisix-mcp/src/oauth.rs @@ -30,6 +30,12 @@ const EXPIRY_SKEW: Duration = Duration::from_secs(60); /// only recommends it). const DEFAULT_TOKEN_LIFETIME: Duration = Duration::from_secs(3600); +/// Cap on a reported `expires_in`. A hostile or broken identity provider can +/// return any u64; an absurd value would overflow `Instant + Duration` and +/// panic the connect task. Thirty days is beyond any sane token lifetime and +/// far below the overflow horizon. +const MAX_TOKEN_LIFETIME: Duration = Duration::from_secs(30 * 24 * 3600); + /// One minted upstream access token. Never printed: the struct deliberately /// has no `Debug` impl, and nothing in this module logs the token value. struct CachedToken { @@ -39,35 +45,54 @@ struct CachedToken { /// Process-global token cache. Guards are held only for map lookups/inserts, /// never across an await; concurrent misses for the same key may fetch in -/// parallel (each minted token is valid — last insert wins). +/// parallel (each minted token is valid — last insert wins). Entries are +/// evicted only by 401-invalidation or key rotation, so tokens for deleted +/// servers linger until restart — the map is bounded by the set of distinct +/// configs ever seen (~100 bytes each), which is acceptable. fn cache() -> &'static RwLock> { static CACHE: OnceLock>> = OnceLock::new(); CACHE.get_or_init(|| RwLock::new(HashMap::new())) } /// Shared HTTP client for token fetches, bounded per request by the same -/// default deadline as any other upstream MCP operation. +/// default deadline as any other upstream MCP operation. Redirects are +/// disabled: a token endpoint never legitimately redirects, and following one +/// would re-POST the secret-bearing form to wherever it points (307/308 keep +/// the body; a chain may even downgrade to plain HTTP). fn http_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); CLIENT.get_or_init(|| { reqwest::Client::builder() .timeout(DEFAULT_UPSTREAM_TIMEOUT) + .redirect(reqwest::redirect::Policy::none()) .build() - // Building a client with only a timeout set cannot fail on any - // supported platform; fall back to the default client if it does. + // Building a client with only a timeout and redirect policy set + // cannot fail on any supported platform; fall back to the default + // client if it does. .unwrap_or_default() }) } -/// Cache key: `token_url \x1f client_id \x1f sha256(client_secret)`. The -/// secret is folded in as a digest — never in plaintext — so a rotated secret -/// can never reuse the previous secret's token, while the key itself stays -/// safe to hold alongside the non-secret fields. +/// Cache key over every field that shapes the minted token: token endpoint, +/// client identity, client secret, and the requested scopes. Two servers +/// sharing one OAuth client but requesting different scopes must never share +/// a token (one would be presented a token minted for the other's scopes). +/// The secret is folded in as a digest — never in plaintext — so a rotated +/// secret can never reuse the previous secret's token. Each component is +/// length-prefixed so no crafted field value can shift a boundary and make +/// two distinct configs collide. fn cache_key(cfg: &OAuthClientConfig) -> String { let secret_digest = hex::encode(Sha256::digest(cfg.client_secret.as_bytes())); + let joined_scopes = cfg.scopes.join(" "); format!( - "{}\x1f{}\x1f{}", - cfg.token_url, cfg.client_id, secret_digest + "{}:{}\x1f{}:{}\x1f{}\x1f{}:{}", + cfg.token_url.len(), + cfg.token_url, + cfg.client_id.len(), + cfg.client_id, + secret_digest, + joined_scopes.len(), + joined_scopes ) } @@ -185,7 +210,8 @@ async fn fetch_token(cfg: &OAuthClientConfig) -> Result<(String, Duration), McpE let lifetime = token .expires_in .map(Duration::from_secs) - .unwrap_or(DEFAULT_TOKEN_LIFETIME); + .unwrap_or(DEFAULT_TOKEN_LIFETIME) + .min(MAX_TOKEN_LIFETIME); Ok((token.access_token, lifetime)) } @@ -459,6 +485,72 @@ mod tests { assert!(err.to_string().contains("malformed"), "got: {err}"); } + #[tokio::test] + async fn same_client_different_scopes_mint_distinct_tokens() { + // Two servers sharing one OAuth client (same IdP, client_id, secret) + // but requesting different scopes must never share a cached token — + // one would be presented a token minted for the other's scopes. + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Mint { + expires_in: Some(3600), + }) + .await; + let read_cfg = OAuthClientConfig { + scopes: vec!["read".to_string()], + ..config(endpoint.url()) + }; + let write_cfg = OAuthClientConfig { + scopes: vec!["write".to_string()], + ..config(endpoint.url()) + }; + + assert_eq!(get_or_fetch(&read_cfg).await.expect("read"), "tok-1"); + assert_eq!(get_or_fetch(&write_cfg).await.expect("write"), "tok-2"); + assert_eq!(endpoint.hits(), 2, "different scopes must mint separately"); + assert_eq!( + endpoint.request(1).get("scope").map(String::as_str), + Some("write") + ); + // Each scope set keeps its own cached token. + assert_eq!(get_or_fetch(&read_cfg).await.expect("read again"), "tok-1"); + assert_eq!(endpoint.hits(), 2); + } + + #[tokio::test] + async fn absurd_expires_in_is_clamped_not_panicking() { + // A hostile IdP can return any u64; an unclamped + // `Instant::now() + Duration::from_secs(u64::MAX)` overflows and + // panics the connect task. + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Mint { + expires_in: Some(u64::MAX), + }) + .await; + let cfg = config(endpoint.url()); + assert_eq!(get_or_fetch(&cfg).await.expect("clamped mint"), "tok-1"); + // Still cached (clamped lifetime is far above the skew). + assert_eq!(get_or_fetch(&cfg).await.expect("cached"), "tok-1"); + assert_eq!(endpoint.hits(), 1); + } + + #[tokio::test] + async fn token_endpoint_redirects_are_refused() { + // A token endpoint never legitimately redirects; following one would + // re-POST the secret-bearing form to wherever it points. The client + // must refuse and fail with the redirect status, not follow it. + let endpoint = spawn_token_endpoint(TokenEndpointBehavior::Static { + status: axum::http::StatusCode::TEMPORARY_REDIRECT, + body: "", + }) + .await; + let err = get_or_fetch(&config(endpoint.url())) + .await + .expect_err("redirect must fail"); + assert!( + err.to_string().contains("HTTP 307"), + "redirect should surface as its status: {err}" + ); + assert_eq!(endpoint.hits(), 1, "the redirect must not be followed"); + } + #[tokio::test] async fn incomplete_config_fails_without_contacting_anything() { for broken in [