Agentgateway: Missing parameter sanitization in MCP to OpenAPI conversion

@@ -198,6 +198,7 @@ dependencies = [
  "opentelemetry-proto",
  "opentelemetry_sdk",
  "parking_lot",
+ "percent-encoding",
  "phonenumber",
  "pin-project-lite",
  "ppp",
@@ -215,6 +216,7 @@ dependencies = [
  "reqwest",
  "rmcp 0.10.0",
  "rmcp 0.12.0",
+ "rstest",
  "rustls",
  "rustls-native-certs",
  "rustls-pemfile",
@@ -2436,6 +2438,12 @@ version = "0.3.31"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988"
 
+[[package]]
+name = "futures-timer"
+version = "3.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24"
+
 [[package]]
 name = "futures-util"
 version = "0.3.31"
@@ -2497,6 +2505,12 @@ version = "0.32.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7"
 
+[[package]]
+name = "glob"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280"
+
 [[package]]
 name = "go-parse-duration"
 version = "0.1.1"
@@ -4405,6 +4419,15 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "proc-macro-crate"
+version = "3.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "219cb19e96be00ab2e37d6e299658a0cfa83e52429179969b0f0121b4ac46983"
+dependencies = [
+ "toml_edit",
+]
+
 [[package]]
 name = "proc-macro-error-attr2"
 version = "2.0.0"
@@ -5034,6 +5057,12 @@ version = "0.8.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7a2d987857b319362043e95f5353c0535c1f58eec5336fdfcf626430af7def58"
 
+[[package]]
+name = "relative-path"
+version = "1.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2"
+
 [[package]]
 name = "reqwest"
 version = "0.12.28"
@@ -5193,6 +5222,35 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "rstest"
+version = "0.26.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49"
+dependencies = [
+ "futures-timer",
+ "futures-util",
+ "rstest_macros",
+]
+
+[[package]]
+name = "rstest_macros"
+version = "0.26.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0"
+dependencies = [
+ "cfg-if",
+ "glob",
+ "proc-macro-crate",
+ "proc-macro2",
+ "quote",
+ "regex",
+ "relative-path",
+ "rustc_version",
+ "syn",
+ "unicode-ident",
+]
+
 [[package]]
 name = "rustc-demangle"
 version = "0.1.26"
@@ -6290,6 +6348,36 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "toml_datetime"
+version = "0.7.5+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347"
+dependencies = [
+ "serde_core",
+]
+
+[[package]]
+name = "toml_edit"
+version = "0.23.10+spec-1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "84c8b9f757e028cee9fa244aea147aab2a9ec09d5325a9b01e0a49730c2b5269"
+dependencies = [
+ "indexmap 2.12.1",
+ "toml_datetime",
+ "toml_parser",
+ "winnow",
+]
+
+[[package]]
+name = "toml_parser"
+version = "1.0.6+spec-1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a3198b4b0a8e11f09dd03e133c0280504d0801269e9afa46362ffde1cbeebf44"
+dependencies = [
+ "winnow",
+]
+
 [[package]]
 name = "tonic"
 version = "0.14.2"
@@ -7367,6 +7455,15 @@ version = "0.53.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650"
 
+[[package]]
+name = "winnow"
+version = "0.7.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a5364e9d77fcdeeaa6062ced926ee3381faa2ee02d3eb83a5c27a8825540829"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "winreg"
 version = "0.50.0"

@@ -138,6 +138,7 @@ reqwest = { version = "0.12", default-features = false, features = [
     "macos-system-configuration",
     "rustls-tls",
 ] }
+rstest = "0.26"
 rustc_version = "0.4"
 rustls = { version = "0.23", features = ["tls12", "ring"] }
 rustls-native-certs = "0.8"

@@ -86,6 +86,7 @@ opentelemetry-http.workspace = true
 opentelemetry.workspace = true
 opentelemetry_sdk.workspace = true
 parking_lot.workspace = true
+percent-encoding.workspace = true
 phonenumber.workspace = true
 pin-project-lite.workspace = true
 ppp.workspace = true
@@ -167,6 +168,7 @@ prost-wkt-build.workspace = true
 assert_matches.workspace = true
 divan.workspace = true
 insta.workspace = true
+rstest.workspace = true
 tempfile.workspace = true
 tokio = { workspace = true, features = ["test-util"] }
 which.workspace = true

@@ -1,12 +1,16 @@
 use std::borrow::Cow;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
+use tracing::{debug, warn};
 
 use ::http::header::{HeaderName, HeaderValue};
 use headers::HeaderMapExt;
 use http::Method;
-use http::header::{ACCEPT, CONTENT_TYPE};
+use http::header::{ACCEPT, CONTENT_LENGTH, CONTENT_TYPE, HOST, TRANSFER_ENCODING};
+use once_cell::sync::Lazy;
 use openapiv3::{OpenAPI, Parameter, ReferenceOr, RequestBody, Schema, SchemaKind, Type};
+use percent_encoding::{AsciiSet, utf8_percent_encode};
+use regex::{Captures, Regex, Replacer};
 use rmcp::model::{ClientRequest, JsonObject, JsonRpcRequest, Tool};
 use serde::{Deserialize, Serialize};
 use serde_json::{Value, json};
@@ -19,6 +23,7 @@ use crate::mcp::upstream::{IncomingRequestContext, UpstreamError};
 pub struct UpstreamOpenAPICall {
 	pub method: String, /* TODO: Switch to Method, but will require getting rid of Serialize/Deserialize */
 	pub path: String,
+	pub allowed_headers: HashSet<String>,
 	// todo: params
 }
 
@@ -32,7 +37,7 @@ pub enum ParseError {
 	MissingReference(String),
 	#[error("unsupported reference")]
 	UnsupportedReference(String),
-	#[error("information required: {0}")] // Corrected typo from "requireds"
+	#[error("information required: {0}")]
 	InformationRequired(String),
 	#[error("serde error: {0}")]
 	SerdeError(#[from] serde_json::Error),
@@ -333,6 +338,12 @@ pub(crate) fn parse_openapi_schema(
 									}
 								})?;
 
+							// Extract allowed header names before consuming param_schemas
+							let allowed_headers: HashSet<String> = param_schemas
+								.get(&ParameterType::Header)
+								.map(|headers| headers.iter().map(|(name, _, _)| name.clone()).collect())
+								.unwrap_or_default();
+
 							for (param_type, props) in param_schemas {
 								let sub_schema = JsonSchema {
 									required: props
@@ -382,6 +393,7 @@ pub(crate) fn parse_openapi_schema(
 								// method: Method::from_bytes(method.as_ref()).expect("todo"),
 								method: method.to_string(),
 								path: path.clone(),
+								allowed_headers,
 							};
 							Ok((tool, upstream))
 						},
@@ -477,6 +489,39 @@ impl Default for JsonSchema {
 	}
 }
 
+/// Regex to match path template parameters like `{param_name}`.
+static PATH_PARAM_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\{([^}]+)\}").unwrap());
+
+/// Characters that are safe in path segments (RFC 3986 unreserved characters).
+/// All other characters will be percent-encoded to prevent path traversal/injection.
+const PATH_SEGMENT_SAFE: &AsciiSet = &percent_encoding::NON_ALPHANUMERIC
+	.remove(b'-')
+	.remove(b'.')
+	.remove(b'_')
+	.remove(b'~');
+
+/// Replaces path template parameters.
+struct PathParamReplacer(serde_json::Map<String, Value>);
+
+impl Replacer for PathParamReplacer {
+	fn replace_append(&mut self, caps: &Captures<'_>, dst: &mut String) {
+		let param = &caps[1];
+		match self.0.get(param) {
+			Some(Value::Number(n_val)) => return dst.push_str(&n_val.to_string()),
+			Some(Value::String(s_val)) => {
+				return dst.extend(utf8_percent_encode(s_val, PATH_SEGMENT_SAFE));
+			},
+			Some(unexpected) => warn!(
+				"Unexpected parameter '{param}' (value: {:?}), leaving path param",
+				unexpected
+			),
+			_ => {},
+		};
+		// fallback to use path parm
+		dst.push_str(&caps[0]);
+	}
+}
+
 /// Normalizes URL path construction to avoid double slashes
 /// Ensures exactly one slash between prefix and path components
 fn normalize_url_path(prefix: &str, path: &str) -> String {
@@ -495,6 +540,13 @@ fn normalize_url_path(prefix: &str, path: &str) -> String {
 	}
 }
 
+fn encode_query_value(value: &str) -> Cow<'_, str> {
+	Cow::from(utf8_percent_encode(
+		value,
+		percent_encoding::NON_ALPHANUMERIC,
+	))
+}
+
 #[derive(Debug)]
 pub struct Handler {
 	pub prefix: String,
@@ -643,26 +695,8 @@ impl Handler {
 		let body_value = args.get(&*BODY_NAME).cloned();
 
 		// --- URL Construction ---
-		let mut path = info.path.clone();
-		// Substitute path parameters into the path template
-		for (key, value) in &path_params {
-			match value {
-				Value::String(s_val) => {
-					path = path.replace(&format!("{{{key}}}"), s_val);
-				},
-				Value::Number(n_val) => {
-					path = path.replace(&format!("{{{key}}}"), n_val.to_string().as_str());
-				},
-				_ => {
-					tracing::warn!(
-						"Path parameter '{}' for tool '{}' is not a string (value: {:?}), skipping substitution",
-						key,
-						name,
-						value
-					);
-				},
-			}
-		}
+		// Substitute path parameters into the path template in a single pass
+		let path = PATH_PARAM_RE.replace_all(&info.path, PathParamReplacer(path_params));
 
 		// Use normalize_url_path to avoid double slashes
 		let normalized_path = normalize_url_path(&self.prefix, &path);
@@ -684,77 +718,116 @@ impl Handler {
 		})?;
 
 		// Build query string
-		let query_string = if !query_params.is_empty() {
-			let mut pairs = Vec::new();
-			for (k, v) in query_params.iter() {
-				if let Some(s) = v.as_str() {
-					pairs.push(format!("{k}={s}"));
-				} else {
-					tracing::warn!(
-						"Query parameter '{}' for tool '{}' is not a string (value: {:?}), skipping",
-						k,
-						name,
-						v
-					);
-				}
-			}
-			if !pairs.is_empty() {
-				format!("?{}", pairs.join("&"))
-			} else {
-				String::new()
-			}
-		} else {
+		let query_string = if query_params.is_empty() {
 			String::new()
+		} else {
+			let format_pair = |param_name: &str, key: &str, v: &Value| -> Option<String> {
+				match v {
+					Value::Null => Some(key.to_string()),
+					Value::Bool(b) => Some(format!("{key}={b}")),
+					Value::Number(n) => Some(format!("{key}={n}")),
+					Value::String(s) => Some(format!("{key}={}", encode_query_value(s))),
+					_ => {
+						warn!(
+							"Query parameter '{}' for tool '{}' unsupported (value: {:?}), skipping",
+							param_name, name, v
+						);
+						None
+					},
+				}
+			};
+			query_params
+				.iter()
+				.flat_map(|(k, v)| {
+					let key = encode_query_value(k);
+					match v {
+						Value::Array(a) => a
+							.iter()
+							.filter_map(|v| format_pair(k, &key, v))
+							.collect::<Vec<_>>(),
+						_ => format_pair(k, &key, v).into_iter().collect(),
+					}
+				})
+				.fold(String::new(), |mut acc, pair| {
+					acc.push(if acc.is_empty() { '?' } else { '&' });
+					acc.push_str(&pair);
+					acc
+				})
 		};
 
 		let uri = format!("{base_url}{query_string}");
-		let mut rb = http::Request::builder().method(method).uri(uri);
 
-		rb = rb.header(ACCEPT, HeaderValue::from_static("application/json"));
+		let mut rb = http::Request::builder()
+			.method(method)
+			.uri(uri)
+			.header(ACCEPT, HeaderValue::from_static("application/json"));
+
+		// Build request body
+		let body = if let Some(body_val) = body_value {
+			rb = rb.header(CONTENT_TYPE, HeaderValue::from_static("application/json"));
+			serde_json::to_vec(&body_val)?
+		} else {
+			Vec::new()
+		};
+
+		// Build the final request
+		let mut request = rb
+			.body(body.into())
+			.map_err(|e| anyhow::anyhow!("Failed to build request: {}", e))?;
+
+		ctx.apply(&mut request);
+
+		// First header set wins including headers set by ctx.apply or the gateway
 		for (key, value) in &header_params {
+			// Only allow headers defined in the OpenAPI schema
+			if !info.allowed_headers.contains(key) {
+				debug!(
+					"Ignoring header '{}' for tool '{}' not defined in schema",
+					key, name
+				);
+				continue;
+			}
+
 			if let Some(s_val) = value.as_str() {
 				match (
 					HeaderName::from_bytes(key.as_bytes()),
 					HeaderValue::from_str(s_val),
 				) {
-					(Ok(h_name), Ok(h_value)) => {
-						rb = rb.header(h_name, h_value);
+					(Ok(header_name), Ok(header_value)) => {
+						// Ingore if header is protected
+						if header_name == CONTENT_LENGTH
+							|| header_name == CONTENT_TYPE
+							|| header_name == TRANSFER_ENCODING
+							|| header_name == HOST
+						{
+							debug!("Ignoring protected header '{}' for tool '{}'", key, name);
+							continue;
+						}
+
+						// Don't override existing headers
+						if request.headers().contains_key(&header_name) {
+							debug!("Ingoring header '{}' for tool '{}' already set", key, name);
+							continue;
+						}
+
+						request.headers_mut().insert(header_name, header_value);
 					},
-					(Err(_), _) => tracing::warn!(
+					(Err(_), _) => warn!(
 						"Invalid header name '{}' for tool '{}', skipping",
-						key,
-						name
+						key, name
 					),
-					(_, Err(_)) => tracing::warn!(
+					(_, Err(_)) => warn!(
 						"Invalid header value '{}' for header '{}' in tool '{}', skipping",
-						s_val,
-						key,
-						name
+						s_val, key, name
 					),
 				}
 			} else {
-				tracing::warn!(
+				warn!(
 					"Header parameter '{}' for tool '{}' is not a string (value: {:?}), skipping",
-					key,
-					name,
-					value
+					key, name, value
 				);
 			}
 		}
-		// Build request body
-		let body = if let Some(body_val) = body_value {
-			rb = rb.header(CONTENT_TYPE, HeaderValue::from_static("application/json"));
-			serde_json::to_vec(&body_val)?
-		} else {
-			Vec::new()
-		};
-
-		// Build the final request
-		let mut request = rb
-			.body(body.into())
-			.map_err(|e| anyhow::anyhow!("Failed to build request: {}", e))?;
-
-		ctx.apply(&mut request);
 
 		let response = self.http_client.call(request).await?;
 

@@ -5,9 +5,10 @@ use agent_core::{metrics, strng};
 use hickory_resolver::config::{ResolverConfig, ResolverOpts};
 use prometheus_client::registry::Registry;
 use rmcp::model::Tool;
+use rstest::rstest;
 use serde_json::json;
 use wiremock::matchers::{body_json, header, method, path, query_param};
-use wiremock::{Mock, MockServer, ResponseTemplate};
+use wiremock::{Match, Mock, MockServer, Request, ResponseTemplate};
 
 use super::*;
 use crate::client::Client;
@@ -91,6 +92,7 @@ async fn setup() -> (MockServer, Handler) {
 	let upstream_call_get = UpstreamOpenAPICall {
 		method: "GET".to_string(),
 		path: "/users/{user_id}".to_string(),
+		allowed_headers: HashSet::from(["X-Request-ID".to_string()]),
 	};
 
 	let test_tool_post = Tool {
@@ -136,6 +138,7 @@ async fn setup() -> (MockServer, Handler) {
 	let upstream_call_post = UpstreamOpenAPICall {
 		method: "POST".to_string(),
 		path: "/users".to_string(),
+		allowed_headers: HashSet::from(["X-API-Key".to_string()]),
 	};
 
 	let backend = SimpleBackend::Opaque(
@@ -359,7 +362,6 @@ async fn test_call_tool_invalid_header_value() {
 	let (server, handler) = setup().await;
 
 	let user_id = "header-issue";
-	// Mock is set up but won't be hit because header construction fails client-side
 	Mock::given(method("GET"))
 		.and(path(format!("/users/{user_id}")))
 		.respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": user_id })))
@@ -369,21 +371,18 @@ async fn test_call_tool_invalid_header_value() {
 	// Intentionally provide a non-string header value
 	let args = json!({
 			"path": { "user_id": user_id },
-			"header": { "X-Request-ID": 12345 } // Invalid header value (not a string)
+			"header": { "X-Request-ID": 12345 }
 	});
 
-	// We expect the call to succeed, but the invalid header should be skipped (and logged)
-	// The mock doesn't expect the header, so if the request goes through without it, it passes.
 	let result = handler
 		.call_tool(
 			"get_user",
 			Some(args.as_object().unwrap().clone()),
 			&IncomingRequestContext::empty(),
 		)
 		.await;
-	assert!(result.is_ok()); // Check that the call still succeeds despite the bad header
+	assert!(result.is_ok());
 	assert_eq!(result.unwrap(), json!({ "id": user_id }));
-	// We can't easily assert the log message here, but manual inspection of logs would show the warning.
 }
 
 #[tokio::test]
@@ -547,7 +546,6 @@ async fn test_call_tool_response_wrapping() {
 		assert_eq!(result.unwrap(), expected);
 	}
 }
-
 #[tokio::test]
 async fn test_normalize_url_path_empty_prefix() {
 	// Test the fix for double slash issue when prefix is empty (host/port config)
@@ -576,9 +574,237 @@ async fn test_normalize_url_path_path_without_leading_slash() {
 	assert_eq!(result, "/api/v3/pet");
 }
 
+#[rstest]
+#[case::empty_prefix("", "/mqtt/healthcheck", "/mqtt/healthcheck")]
+#[case::with_prefix("/api/v3/", "/pet", "/api/v3/pet")]
+#[case::prefix_no_trailing_slash("/api/v3", "/pet", "/api/v3/pet")]
+#[case::without_leading_slash("/api/v3", "pet", "/api/v3/pet")]
+#[case::empty_prefix_path_without_slash("", "pet", "/pet")]
+fn test_normalize_url_path(#[case] prefix: &str, #[case] path: &str, #[case] expected: &str) {
+	let result = super::normalize_url_path(prefix, path);
+	assert_eq!(result, expected);
+}
+
+#[rstest]
+#[case::empty_string(json!({"verbose": ""}), vec![("verbose", "")])]
+#[case::string_value(json!({"verbose": "true"}), vec![("verbose", "true")])]
+#[case::boolean_true(json!({"verbose": true}), vec![("verbose", "true")])]
+#[case::boolean_false(json!({"verbose": false}), vec![("verbose", "false")])]
+#[case::integer_value(json!({"verbose": "123"}), vec![("verbose", "123")])]
+#[case::special_chars(json!({"verbose": "hello world"}), vec![("verbose", "hello world")])]
+#[case::array_values(json!({"verbose": ["a", "b", "c"]}), vec![("verbose", "a"), ("verbose", "b"), ("verbose", "c")])]
+#[case::ampersand_in_value(json!({"verbose": "foo&admin=true"}), vec![("verbose", "foo&admin=true")])]
+#[case::equals_in_value(json!({"verbose": "foo=bar"}), vec![("verbose", "foo=bar")])]
+#[case::question_mark_in_value(json!({"verbose": "foo?bar"}), vec![("verbose", "foo?bar")])]
+#[case::combined_injection(json!({"verbose": "x&evil=1&admin=true"}), vec![("verbose", "x&evil=1&admin=true")])]
+#[tokio::test]
+async fn test_query_param_types(
+	#[case] query_args: serde_json::Value,
+	#[case] expected_params: Vec<(&str, &str)>,
+) {
+	let (server, handler) = setup().await;
+
+	let user_id = "test-user";
+
+	let mut mock = Mock::given(method("GET")).and(path(format!("/users/{user_id}")));
+	for (key, value) in &expected_params {
+		mock = mock.and(query_param(*key, *value));
+	}
+	mock
+		.respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": user_id })))
+		.expect(1)
+		.mount(&server)
+		.await;
+
+	let args = json!({
+		"path": { "user_id": user_id },
+		"query": query_args
+	});
+
+	let result = handler
+		.call_tool(
+			"get_user",
+			Some(args.as_object().unwrap().clone()),
+			&IncomingRequestContext::empty(),
+		)
+		.await;
+
+	assert!(result.is_ok(), "Expected success, got: {:?}", result.err());
+	assert_eq!(result.unwrap(), json!({ "id": user_id }));
+}
+
+#[rstest]
+#[case::simple_id("123", "/users/123")]
+#[case::numeric_id("456", "/users/456")]
+#[case::spaces("user name", "/users/user%20name")]
+#[case::unicode("user\u{00e9}", "/users/user%C3%A9")]
+#[case::path_traversal("../admin", "/users/..%2Fadmin")]
+#[case::embedded_slashes("user-1/o/er-1001", "/users/user-1%2Fo%2Fer-1001")]
+#[case::query_injection("123?admin=true", "/users/123%3Fadmin%3Dtrue")]
+#[case::query_with_ampersand("123?a=1&b=2", "/users/123%3Fa%3D1%26b%3D2")]
+#[case::hash_fragment("user#section", "/users/user%23section")]
+#[case::ampersand_in_path("user&admin=true", "/users/user%26admin%3Dtrue")]
+#[tokio::test]
+async fn test_path_param_encoding(#[case] user_id: &str, #[case] expected_path: &str) {
+	let (server, handler) = setup().await;
+
+	Mock::given(method("GET"))
+		.and(path(expected_path))
+		.respond_with(ResponseTemplate::new(200).set_body_json(json!({ "id": user_id })))
+		.expect(1)
+		.mount(&server)
+		.await;
+
+	let args = json!({ "path": { "user_id": user_id } });
+
+	let result = handler
+		.call_tool(
+			"get_user",
+			Some(args.as_object().unwrap().clone()),
+			&IncomingRequestContext::empty(),
+		)
+		.await;
+
+	assert!(result.is_ok(), "Expected success, got: {:?}", result.err());
+	assert_eq!(result.unwrap(), json!({ "id": user_id }));
+}
+
+#[tokio::test]
+async fn test_schema_defined_headers_work() {
+	let (server, handler) = setup().await;
+
+	let user_id = "custom-header-test";
+	let expected_response = json!({ "id": user_id });
+
+	// Only X-Request-ID is defined in the schema for get_user tool
+	Mock::given(method("GET"))
+		.and(path(format!("/users/{user_id}")))
+		.and(header("X-Request-ID", "my-request-123"))
+		.respond_with(ResponseTemplate::new(200).set_body_json(&expected_response))
+		.expect(1)
+		.mount(&server)
+		.await;
+
+	let args = json!({
+		"path": { "user_id": user_id },
+		"header": {
+			"X-Request-ID": "my-request-123"
+		}
+	});
+
+	let result = handler
+		.call_tool(
+			"get_user",
+			Some(args.as_object().unwrap().clone()),
+			&IncomingRequestContext::empty(),
+		)
+		.await;
+
+	assert!(
+		result.is_ok(),
+		"Schema-defined headers should work: {:?}",
+		result.err()
+	);
+	assert_eq!(result.unwrap(), expected_response);
+}
+
+// Custom matcher to verify a header is NOT present
+struct HeaderNotPresent {
+	header_name: String,
+}
+
+impl HeaderNotPresent {
+	fn new(header_name: impl Into<String>) -> Self {
+		Self {
+			header_name: header_name.into(),
+		}
+	}
+}
+
+impl Match for HeaderNotPresent {
+	fn matches(&self, request: &Request) -> bool {
+		!request.headers.contains_key(self.header_name.as_str())
+	}
+}
+
 #[tokio::test]
-async fn test_normalize_url_path_empty_prefix_path_without_slash() {
-	// Test edge case: empty prefix and path without leading slash
-	let result = super::normalize_url_path("", "pet");
-	assert_eq!(result, "/pet");
+async fn test_blocked_headers_are_ignored() {
+	let (server, handler) = setup().await;
+
+	let request_body = json!({ "name": "Test User", "email": "test@example.com" });
+	let expected_response = json!({ "id": "new-user", "name": "Test User" });
+
+	Mock::given(method("POST"))
+		.and(path("/users"))
+		.and(header("content-length", "47")) // length of request_body
+		.and(header("content-type", "application/json"))
+		.and(HeaderNotPresent::new("transfer-encoding"))
+		.and(body_json(&request_body))
+		.respond_with(ResponseTemplate::new(201).set_body_json(&expected_response))
+		.expect(1)
+		.mount(&server)
+		.await;
+
+	let args = json!({
+		"body": request_body,
+		"header": {
+			"content-length": "999999999",
+			"content-type": "text/plain",
+			"transfer-encoding": "chunked",
+			"host": "evil.com"
+		}
+	});
+
+	let result = handler
+		.call_tool(
+			"create_user",
+			Some(args.as_object().unwrap().clone()),
+			&IncomingRequestContext::empty(),
+		)
+		.await;
+
+	// The request should succeed with the correct headers (blocked headers ignored)
+	assert!(result.is_ok(), "Request should succeed: {:?}", result.err());
+	assert_eq!(result.unwrap(), expected_response);
+}
+
+#[tokio::test]
+async fn test_headers_not_in_schema_are_ignored() {
+	let (server, handler) = setup().await;
+
+	let user_id = "schema-header-test";
+	let expected_response = json!({ "id": user_id });
+
+	// Only expect X-Request-ID (defined in schema), NOT X-Malicious-Header
+	Mock::given(method("GET"))
+		.and(path(format!("/users/{user_id}")))
+		.and(header("X-Request-ID", "valid-request"))
+		.and(HeaderNotPresent::new("X-Malicious-Header"))
+		.respond_with(ResponseTemplate::new(200).set_body_json(&expected_response))
+		.expect(1)
+		.mount(&server)
+		.await;
+
+	let args = json!({
+		"path": { "user_id": user_id },
+		"header": {
+			"X-Request-ID": "valid-request",
+			"X-Malicious-Header": "should-be-ignored"
+		}
+	});
+
+	let result = handler
+		.call_tool(
+			"get_user",
+			Some(args.as_object().unwrap().clone()),
+			&IncomingRequestContext::empty(),
+		)
+		.await;
+
+	assert!(
+		result.is_ok(),
+		"Request should succeed with schema-defined headers: {:?}",
+		result.err()
+	);
+	assert_eq!(result.unwrap(), expected_response);
 }