From d23722e8f42f81e24b464af21db7b8e4e9329356 Mon Sep 17 00:00:00 2001 From: reidliu41 Date: Sun, 8 Mar 2026 09:23:47 +0800 Subject: [PATCH] feat(mcp): support custom HTTP headers for non-OAuth MCP servers (#639) --- src/app.rs | 5 +- src/cli/mcp.rs | 60 ++++++++++-- src/extensions/manager.rs | 2 +- src/tools/mcp/client.rs | 190 +++++++++++++++++++++++++++++++++++++- src/tools/mcp/config.rs | 149 ++++++++++++++++++++++++++++++ 5 files changed, 394 insertions(+), 12 deletions(-) diff --git a/src/app.rs b/src/app.rs index 766aff30d77..bedc0827b43 100644 --- a/src/app.rs +++ b/src/app.rs @@ -537,7 +537,10 @@ impl AppBuilder { let has_tokens = is_authenticated(&server, &secrets, "default").await; - let client = if has_tokens || server.requires_auth() { + let client = if has_tokens + || server.requires_auth() + || server.has_custom_headers() + { McpClient::new_authenticated( server, mcp_sm, secrets, "default", ) diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs index 5e2f4deaf5a..b5f5a1bfa20 100644 --- a/src/cli/mcp.rs +++ b/src/cli/mcp.rs @@ -44,6 +44,13 @@ pub enum McpCommand { #[arg(long)] scopes: Option, + /// Custom HTTP header (repeatable, format: "Name:Value") + /// + /// For MCP servers that use header-based auth instead of OAuth. + /// Example: --header "X-API-Key:sk-abc123" + #[arg(long = "header", value_name = "NAME:VALUE")] + headers: Vec, + /// Server description #[arg(long)] description: Option, @@ -107,6 +114,7 @@ pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> { auth_url, token_url, scopes, + headers, description, } => { add_server( @@ -116,6 +124,7 @@ pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> { auth_url, token_url, scopes, + headers, description, ) .await @@ -133,6 +142,7 @@ pub async fn run_mcp_command(cmd: McpCommand) -> anyhow::Result<()> { } /// Add a new MCP server. +#[allow(clippy::too_many_arguments)] async fn add_server( name: String, url: String, @@ -140,6 +150,7 @@ async fn add_server( auth_url: Option, token_url: Option, scopes: Option, + headers: Vec, description: Option, ) -> anyhow::Result<()> { let mut config = McpServerConfig::new(&name, &url); @@ -148,8 +159,8 @@ async fn add_server( config = config.with_description(desc); } - // Track if auth is required - let requires_auth = client_id.is_some(); + // Track if OAuth auth is required + let requires_oauth = client_id.is_some(); // Set up OAuth if client_id is provided if let Some(client_id) = client_id { @@ -170,7 +181,24 @@ async fn add_server( config = config.with_oauth(oauth); } - // Validate + // Parse custom headers (format: "Name:Value") + if !headers.is_empty() { + let mut header_map = std::collections::HashMap::new(); + for raw in &headers { + let (key, value) = raw.split_once(':').ok_or_else(|| { + anyhow::anyhow!("Invalid header format '{}'. Expected 'Name:Value'.", raw) + })?; + let key = key.trim().to_string(); + let value = value.trim().to_string(); + if key.is_empty() { + anyhow::bail!("Header name cannot be empty in '{}'", raw); + } + header_map.insert(key, value); + } + config = config.with_headers(header_map); + } + + // Validate (includes header name/value safety checks) config.validate()?; // Save (DB if available, else disk) @@ -183,9 +211,11 @@ async fn add_server( println!(" ✓ Added MCP server '{}'", name); println!(" URL: {}", url); - if requires_auth { + if requires_oauth { println!(); println!(" Run 'ironclaw mcp auth {}' to authenticate.", name); + } else if !headers.is_empty() { + println!(" Auth: custom headers ({} configured)", headers.len()); } println!(); @@ -230,7 +260,9 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> { for server in &servers.servers { let status = if server.enabled { "●" } else { "○" }; - let auth_status = if server.requires_auth() { + let auth_status = if server.has_custom_headers() { + " (header auth)" + } else if server.requires_auth() { " (auth required)" } else { "" @@ -248,6 +280,17 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> { println!(" Scopes: {}", oauth.scopes.join(", ")); } } + if let Some(ref headers) = server.headers { + for (name, value) in headers { + // Mask header values to avoid leaking secrets in logs/terminal + let masked = if value.len() <= 4 { + "****".to_string() + } else { + format!("{}...{}", &value[..2], &value[value.len() - 2..]) + }; + println!(" Header: {}: {}", name, masked); + } + } println!(); } else { println!( @@ -360,8 +403,9 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> { let secrets = get_secrets_store().await?; let has_tokens = is_authenticated(&server, &secrets, &user_id).await; - let client = if has_tokens { - // We have stored tokens, use authenticated client + let client = if has_tokens || server.has_custom_headers() { + // We have stored tokens or custom headers — use authenticated client + // (custom headers are injected via server_config in send_request) McpClient::new_authenticated(server.clone(), session_manager, secrets, user_id) } else if server.requires_auth() { // OAuth configured but no tokens - need to authenticate @@ -373,7 +417,7 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> { println!(); return Ok(()); } else { - // No OAuth and no tokens - try unauthenticated + // No OAuth, no tokens, no custom headers - try unauthenticated McpClient::new_with_name(&server.name, &server.url) }; diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index 6ef47d22303..fb5a7e2325e 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -2500,7 +2500,7 @@ impl ExtensionManager { let has_tokens = is_authenticated(&server, &self.secrets, &self.user_id).await; - let client = if has_tokens || server.requires_auth() { + let client = if has_tokens || server.requires_auth() || server.has_custom_headers() { McpClient::new_authenticated( server.clone(), Arc::clone(&self.mcp_session_manager), diff --git a/src/tools/mcp/client.rs b/src/tools/mcp/client.rs index 316851fc31b..bbe6f364393 100644 --- a/src/tools/mcp/client.rs +++ b/src/tools/mcp/client.rs @@ -178,8 +178,26 @@ impl McpClient { .header("Content-Type", "application/json") .json(&request); - // Add Authorization header if we have a token - if let Some(token) = self.get_access_token().await? { + // Add custom headers from config (if any) + let has_custom_auth = if let Some(ref config) = self.server_config + && config.has_custom_headers() + { + let has_auth = config.has_custom_auth_header(); + if let Some(ref headers) = config.headers { + for (name, value) in headers { + // Headers were validated at config time (CRLF-safe) + req_builder = req_builder.header(name.as_str(), value.as_str()); + } + } + has_auth + } else { + false + }; + + // Add OAuth Bearer token if we have one and the user hasn't + // supplied a custom Authorization header (case-insensitive check + // already done by has_custom_auth_header). + if !has_custom_auth && let Some(token) = self.get_access_token().await? { req_builder = req_builder.header("Authorization", format!("Bearer {}", token)); } @@ -740,4 +758,172 @@ mod tests { }; assert!(!tool.requires_approval()); } + + #[test] + fn test_custom_auth_header_skips_bearer() { + // When config has a custom Authorization header, + // has_custom_auth_header() should return true so send_request + // skips auto Bearer injection. + let mut headers = std::collections::HashMap::new(); + headers.insert("Authorization".to_string(), "Bearer static-tok".to_string()); + let config = McpServerConfig::new("test", "https://mcp.example.com").with_headers(headers); + assert!(config.has_custom_auth_header()); + } + + #[test] + fn test_non_auth_headers_allow_bearer() { + // Custom headers that aren't Authorization should NOT block Bearer + let mut headers = std::collections::HashMap::new(); + headers.insert("X-API-Key".to_string(), "sk-abc".to_string()); + let config = McpServerConfig::new("test", "https://mcp.example.com").with_headers(headers); + assert!(!config.has_custom_auth_header()); + } + + /// Spin up a lightweight axum server that echoes back all received request + /// headers as a JSON-RPC result, so we can assert on the actual wire-level + /// headers sent by `send_request`. + async fn mock_mcp_echo_server() -> (String, tokio::task::JoinHandle<()>) { + use axum::http::StatusCode; + use axum::{Router, extract::Request, routing::post}; + + let app = Router::new().route( + "/", + post(|req: Request| async move { + // Collect headers into a map for easy assertion. + let headers: std::collections::HashMap = req + .headers() + .iter() + .map(|(k, v)| (k.as_str().to_string(), v.to_str().unwrap_or("").to_string())) + .collect(); + let body = serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "result": headers + }); + (StatusCode::OK, axum::Json(body)) + }), + ); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let url = format!("http://127.0.0.1:{}", addr.port()); + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + (url, handle) + } + + /// Helper: call send_request on a client and return the echoed headers. + async fn send_and_get_headers(client: &McpClient) -> std::collections::HashMap { + let req = McpRequest::initialize(1); + let resp = client.send_request(req).await.unwrap(); + serde_json::from_value(resp.result.unwrap()).unwrap() + } + + #[tokio::test] + async fn test_wire_custom_headers_injected() { + // Custom headers should appear in the outbound request. + let (url, handle) = mock_mcp_echo_server().await; + let mut hdrs = std::collections::HashMap::new(); + hdrs.insert("X-Api-Key".to_string(), "sk-test-123".to_string()); + let config = McpServerConfig::new("test", &url).with_headers(hdrs); + let client = McpClient::new_with_name("test", &url); + // Inject config so send_request reads custom headers + let client = McpClient { + server_config: Some(config), + ..client + }; + let headers = send_and_get_headers(&client).await; + assert_eq!(headers.get("x-api-key").unwrap(), "sk-test-123"); + handle.abort(); + } + + #[tokio::test] + async fn test_wire_custom_auth_suppresses_bearer() { + // When config has a custom Authorization header, auto Bearer must NOT appear. + let (url, handle) = mock_mcp_echo_server().await; + let mut hdrs = std::collections::HashMap::new(); + hdrs.insert( + "Authorization".to_string(), + "Token my-static-token".to_string(), + ); + let config = McpServerConfig::new("test", &url).with_headers(hdrs); + let client = McpClient::new_with_name("test", &url); + let client = McpClient { + server_config: Some(config), + ..client + }; + let headers = send_and_get_headers(&client).await; + assert_eq!( + headers.get("authorization").unwrap(), + "Token my-static-token" + ); + handle.abort(); + } + + #[tokio::test] + async fn test_wire_non_auth_headers_allow_bearer() { + // Custom non-Authorization headers must NOT suppress auto Bearer. + // Use a real InMemorySecretsStore with a stored token to prove Bearer + // IS injected alongside the custom header. + use crate::secrets::{CreateSecretParams, InMemorySecretsStore, SecretsCrypto}; + use secrecy::SecretString; + + let (url, handle) = mock_mcp_echo_server().await; + + // Set up secrets store with a token + let crypto = Arc::new( + SecretsCrypto::new(SecretString::from("test-master-key-32chars!!!!!!!!xx")).unwrap(), + ); + let secrets: Arc = + Arc::new(InMemorySecretsStore::new(crypto)); + let mut hdrs = std::collections::HashMap::new(); + hdrs.insert("X-Custom".to_string(), "value".to_string()); + let config = McpServerConfig::new("test", &url).with_headers(hdrs); + + // Store a token under the expected secret name + let token_name = config.token_secret_name(); // "mcp_test_access_token" + secrets + .create( + "default", + CreateSecretParams::new(&token_name, "oauth-bearer-tok"), + ) + .await + .unwrap(); + + let session_manager = Arc::new(McpSessionManager::new()); + let client = McpClient::new_authenticated(config, session_manager, secrets, "default"); + + let headers = send_and_get_headers(&client).await; + // Custom header present + assert_eq!(headers.get("x-custom").unwrap(), "value"); + // Bearer token also present (not suppressed by non-auth custom header) + let auth = headers.get("authorization").unwrap(); + assert!( + auth.starts_with("Bearer "), + "expected Bearer token, got: {}", + auth + ); + handle.abort(); + } + + #[test] + fn test_client_creation_branches() { + // No headers, no OAuth, localhost → new_with_name path + let config = McpServerConfig::new("local", "http://localhost:8080"); + assert!(!config.requires_auth()); + assert!(!config.has_custom_headers()); + + // Headers only, localhost → has_custom_headers triggers authenticated path + let mut headers = std::collections::HashMap::new(); + headers.insert("X-Key".to_string(), "val".to_string()); + let config = McpServerConfig::new("local", "http://localhost:8080").with_headers(headers); + assert!(!config.requires_auth()); + assert!(config.has_custom_headers()); + + // Remote HTTPS, no headers → requires_auth triggers authenticated path + let config = McpServerConfig::new("remote", "https://mcp.example.com"); + assert!(config.requires_auth()); + assert!(!config.has_custom_headers()); + } } diff --git a/src/tools/mcp/config.rs b/src/tools/mcp/config.rs index 784f0aa25c1..b1d3e7782b0 100644 --- a/src/tools/mcp/config.rs +++ b/src/tools/mcp/config.rs @@ -25,6 +25,13 @@ pub struct McpServerConfig { #[serde(skip_serializing_if = "Option::is_none")] pub oauth: Option, + /// Custom HTTP headers to send with every request. + /// + /// Used for MCP servers that require non-OAuth authentication + /// (e.g., `X-API-Key`, `Authorization: Bearer `). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub headers: Option>, + /// Whether this server is enabled. #[serde(default = "default_true")] pub enabled: bool, @@ -45,6 +52,7 @@ impl McpServerConfig { name: name.into(), url: url.into(), oauth: None, + headers: None, enabled: true, description: None, } @@ -62,6 +70,28 @@ impl McpServerConfig { self } + /// Set custom HTTP headers. + pub fn with_headers(mut self, headers: HashMap) -> Self { + if headers.is_empty() { + self.headers = None; + } else { + self.headers = Some(headers); + } + self + } + + /// Check if this server has custom headers configured. + pub fn has_custom_headers(&self) -> bool { + self.headers.as_ref().is_some_and(|h| !h.is_empty()) + } + + /// Check if custom headers include an Authorization header (case-insensitive). + pub fn has_custom_auth_header(&self) -> bool { + self.headers + .as_ref() + .is_some_and(|h| h.keys().any(|k| k.eq_ignore_ascii_case("authorization"))) + } + /// Validate the server configuration. pub fn validate(&self) -> Result<(), ConfigError> { if self.name.is_empty() { @@ -85,6 +115,25 @@ impl McpServerConfig { }); } + // Validate custom headers (reject invalid header names/values to prevent CRLF injection) + if let Some(ref headers) = self.headers { + for (name, value) in headers { + reqwest::header::HeaderName::from_bytes(name.as_bytes()).map_err(|_| { + ConfigError::InvalidConfig { + reason: format!("Invalid header name: '{}'", name), + } + })?; + reqwest::header::HeaderValue::from_str(value).map_err(|_| { + ConfigError::InvalidConfig { + reason: format!( + "Invalid header value for '{}' (must be visible ASCII, no CRLF)", + name + ), + } + })?; + } + } + Ok(()) } @@ -593,4 +642,104 @@ mod tests { let config = McpServerConfig::new("bad", "http://mcp.example.com"); assert!(!config.requires_auth()); } + + #[test] + fn test_has_custom_headers() { + let config = McpServerConfig::new("test", "https://mcp.example.com"); + assert!(!config.has_custom_headers()); + + let mut headers = HashMap::new(); + headers.insert("X-API-Key".to_string(), "sk-abc".to_string()); + let config = config.with_headers(headers); + assert!(config.has_custom_headers()); + } + + #[test] + fn test_has_custom_auth_header_case_insensitive() { + let mut headers = HashMap::new(); + headers.insert("authorization".to_string(), "Bearer token".to_string()); + let config = McpServerConfig::new("test", "https://mcp.example.com").with_headers(headers); + assert!(config.has_custom_auth_header()); + + let mut headers = HashMap::new(); + headers.insert("Authorization".to_string(), "Bearer token".to_string()); + let config = McpServerConfig::new("test", "https://mcp.example.com").with_headers(headers); + assert!(config.has_custom_auth_header()); + + let mut headers = HashMap::new(); + headers.insert("X-API-Key".to_string(), "key".to_string()); + let config = McpServerConfig::new("test", "https://mcp.example.com").with_headers(headers); + assert!(!config.has_custom_auth_header()); + } + + #[test] + fn test_requires_auth_unchanged_by_headers() { + // Custom headers should NOT change requires_auth() semantics + let mut headers = HashMap::new(); + headers.insert("X-API-Key".to_string(), "sk-abc".to_string()); + let config = McpServerConfig::new("local", "http://localhost:8080").with_headers(headers); + assert!(!config.requires_auth()); + assert!(config.has_custom_headers()); + } + + #[test] + fn test_validate_rejects_invalid_header_name() { + let mut headers = HashMap::new(); + headers.insert("Bad\r\nHeader".to_string(), "value".to_string()); + let config = McpServerConfig::new("test", "http://localhost:8080").with_headers(headers); + assert!(config.validate().is_err()); + } + + #[test] + fn test_validate_rejects_crlf_header_value() { + let mut headers = HashMap::new(); + headers.insert("X-Key".to_string(), "val\r\nInjected: bad".to_string()); + let config = McpServerConfig::new("test", "http://localhost:8080").with_headers(headers); + assert!(config.validate().is_err()); + } + + #[test] + fn test_validate_accepts_valid_headers() { + let mut headers = HashMap::new(); + headers.insert("X-API-Key".to_string(), "sk-abc123".to_string()); + headers.insert("Authorization".to_string(), "Bearer my-token".to_string()); + let config = McpServerConfig::new("test", "http://localhost:8080").with_headers(headers); + assert!(config.validate().is_ok()); + } + + #[test] + fn test_with_headers_empty_clears() { + let mut headers = HashMap::new(); + headers.insert("X-Key".to_string(), "val".to_string()); + let config = McpServerConfig::new("test", "http://localhost:8080").with_headers(headers); + assert!(config.has_custom_headers()); + + let config = config.with_headers(HashMap::new()); + assert!(!config.has_custom_headers()); + } + + #[test] + fn test_headers_serialization_roundtrip() { + let mut headers = HashMap::new(); + headers.insert("X-API-Key".to_string(), "sk-abc".to_string()); + let config = McpServerConfig::new("test", "https://mcp.example.com").with_headers(headers); + + let json = serde_json::to_string(&config).unwrap(); + assert!(json.contains("X-API-Key")); + + let restored: McpServerConfig = serde_json::from_str(&json).unwrap(); + assert!(restored.has_custom_headers()); + assert_eq!( + restored.headers.unwrap().get("X-API-Key").unwrap(), + "sk-abc" + ); + } + + #[test] + fn test_headers_deserialization_missing_field() { + // Old config without headers field should deserialize fine + let json = r#"{"name":"test","url":"https://mcp.example.com","enabled":true}"#; + let config: McpServerConfig = serde_json::from_str(json).unwrap(); + assert!(!config.has_custom_headers()); + } }