Skip to content

fix(auth): Fix unnecesary header sanitization in tansport #137

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Agent-Hellboy wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(auth): Fix unnecesary header sanitization in tansport #137

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion mcp_fuzzer/cli/runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,12 @@ def create_transport_with_auth(args, client_args: dict[str, Any]):
logger.debug("No auth headers found for default tool mapping")

factory_kwargs = {"timeout": args.timeout}


# Get safety settings
safety_enabled = client_args.get("safety_enabled", True)
if safety_enabled:
factory_kwargs["safety_enabled"] = safety_enabled
Comment on lines +50 to +54
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Fix the conditional logic for safety_enabled.

The current implementation only adds safety_enabled to factory_kwargs when it's True. If a user explicitly sets safety_enabled=False, it won't be passed to the transport, and the transport will default to True, ignoring the user's preference.

Apply this diff to always pass the safety_enabled value:

     factory_kwargs = {"timeout": args.timeout}
-
-    # Get safety settings
-    safety_enabled = client_args.get("safety_enabled", True)
-    if safety_enabled:
-        factory_kwargs["safety_enabled"] = safety_enabled
-
+    
+    # Pass safety_enabled to transport (defaults to True if not provided)
+    factory_kwargs["safety_enabled"] = client_args.get("safety_enabled", True)
+
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
# Get safety settings
safety_enabled = client_args.get("safety_enabled", True)
if safety_enabled:
factory_kwargs["safety_enabled"] = safety_enabled
factory_kwargs = {"timeout": args.timeout}
# Pass safety_enabled to transport (defaults to True if not provided)
factory_kwargs["safety_enabled"] = client_args.get("safety_enabled", True)
🤖 Prompt for AI Agents 
In mcp_fuzzer/cli/runner.py around lines 50 to 54, the code only sets
factory_kwargs["safety_enabled"] when safety_enabled is True, so an explicit
False is ignored; always assign the retrieved safety_enabled value into
factory_kwargs (i.e., remove the conditional and unconditionally set
factory_kwargs["safety_enabled"] = safety_enabled) so the transport receives the
user's preference.


# Apply auth headers to HTTP-based protocols
if args.protocol in ("http", "https", "streamablehttp", "sse") and auth_headers:
factory_kwargs["auth_headers"] = auth_headers
Expand Down
1 change: 1 addition & 0 deletions mcp_fuzzer/client/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ def __init__(self, protocol, endpoint, timeout):

client_args = {
"auth_manager": config.get("auth_manager"),
"safety_enabled": config.get("safety_enabled", True),
}

transport = create_transport_with_auth(args, client_args)
Expand Down
23 changes: 17 additions & 6 deletions mcp_fuzzer/transport/http.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,19 +48,30 @@ def __init__(
url: str,
timeout: float = 30.0,
auth_headers: dict[str, str | None] | None = None,
safety_enabled: bool = True,
):
self.url = url
self.timeout = timeout
self.safety_enabled = safety_enabled
self.headers = {
"Accept": DEFAULT_HTTP_ACCEPT,
"Content-Type": JSON_CONTENT_TYPE,
}
if auth_headers:
self.headers.update(auth_headers)
self.auth_headers = {k: v for k, v in (auth_headers or {}).items() if v is not None}
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Line exceeds maximum length.

Line 60 exceeds the project's 88-character limit (92 characters). Consider breaking it into multiple lines for better readability.

Apply this diff:

-    self.auth_headers = {k: v for k, v in (auth_headers or {}).items() if v is not None}
+    self.auth_headers = {
+        k: v for k, v in (auth_headers or {}).items() if v is not None
+    }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
self.auth_headers = {k: v for k, v in (auth_headers or {}).items() if v is not None}
self.auth_headers = {
k: v for k, v in (auth_headers or {}).items() if v is not None
}
🧰 Tools
🪛 GitHub Actions: Lint

[error] 60-60: E501 Line too long (92 > 88)

🤖 Prompt for AI Agents 
In mcp_fuzzer/transport/http.py around line 60, the dict comprehension creating
self.auth_headers exceeds the 88-character line limit; break the expression
across multiple lines to stay within the limit while preserving behavior — e.g.,
compute the source mapping on one line (src = auth_headers or {}), iterate or
build the dict comprehension on the next line(s), or split the comprehension
into a multiline dict(...) or a simple for-loop that assigns self.auth_headers =
{ ... } across lines; ensure the filter "if v is not None" remains and no
semantic changes are made.


# Track last activity for process management
self._last_activity = time.time()

def _prepare_headers_with_auth(self, headers: dict[str, str]) -> dict[str, str]:
"""Prepare headers with optional safety sanitization and auth headers."""
if self.safety_enabled:
safe_headers = self._prepare_safe_headers(headers)
else:
safe_headers = headers.copy()
# Add auth headers after sanitization (they are user-configured and safe)
safe_headers.update(self.auth_headers)
return safe_headers

# Initialize process manager for any subprocesses (like proxy servers)
watchdog_config = WatchdogConfig(
check_interval=1.0,
Expand Down Expand Up @@ -118,7 +129,7 @@ async def send_request(

# Use shared network functionality
self._validate_network_request(self.url)
safe_headers = self._prepare_safe_headers(self.headers)
safe_headers = self._prepare_headers_with_auth(self.headers)

async with self._create_http_client(self.timeout) as client:
response = await client.post(self.url, json=payload, headers=safe_headers)
Expand Down Expand Up @@ -162,7 +173,7 @@ async def send_raw(self, payload: dict[str, Any]) -> Any:

# Use shared network functionality
self._validate_network_request(self.url)
safe_headers = self._prepare_safe_headers(self.headers)
safe_headers = self._prepare_headers_with_auth(self.headers)

async with self._create_http_client(self.timeout) as client:
response = await client.post(self.url, json=payload, headers=safe_headers)
Expand Down Expand Up @@ -202,7 +213,7 @@ async def send_notification(

# Use shared network functionality
self._validate_network_request(self.url)
safe_headers = self._prepare_safe_headers(self.headers)
safe_headers = self._prepare_headers_with_auth(self.headers)

async with self._create_http_client(self.timeout) as client:
response = await client.post(self.url, json=payload, headers=safe_headers)
Expand Down Expand Up @@ -242,7 +253,7 @@ async def _stream_request(

# Use shared network functionality
self._validate_network_request(self.url)
safe_headers = self._prepare_safe_headers(self.headers)
safe_headers = self._prepare_headers_with_auth(self.headers)

async with self._create_http_client(self.timeout) as client:
# First request
Expand Down
21 changes: 16 additions & 5 deletions mcp_fuzzer/transport/sse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,26 @@ def __init__(
url: str,
timeout: float = 30.0,
auth_headers: dict[str, str | None] | None = None,
safety_enabled: bool = True,
):
self.url = url
self.timeout = timeout
self.safety_enabled = safety_enabled
self.headers = {
"Accept": "text/event-stream",
"Content-Type": "application/json",
}
if auth_headers:
self.headers.update(auth_headers)
self.auth_headers = {k: v for k, v in (auth_headers or {}).items() if v is not None}
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Line exceeds maximum length.

Line 26 exceeds the project's 88-character limit (92 characters).

Apply this diff:

-    self.auth_headers = {k: v for k, v in (auth_headers or {}).items() if v is not None}
+    self.auth_headers = {
+        k: v for k, v in (auth_headers or {}).items() if v is not None
+    }

Committable suggestion skipped: line range outside the PR's diff.

🧰 Tools
🪛 GitHub Actions: Lint

[error] 26-26: E501 Line too long (92 > 88)

🤖 Prompt for AI Agents 
In mcp_fuzzer/transport/sse.py around line 26, the dict comprehension assigning
self.auth_headers is longer than the 88-character limit; break the expression
across multiple lines so no source line exceeds 88 chars. Split the
comprehension into a multi-line expression (for example, put the opening brace
on the first line, place the k: v for k, v ... clause on the next line, and the
closing brace on its own line) or build the dict using a short call (e.g.,
dict(...) or a small helper) spread across lines to satisfy the line-length rule
while preserving the same filtering logic that excludes None values.


def _prepare_headers_with_auth(self, headers: dict[str, str]) -> dict[str, str]:
"""Prepare headers with optional safety sanitization and auth headers."""
if self.safety_enabled:
safe_headers = sanitize_headers(headers)
else:
safe_headers = headers.copy()
# Add auth headers after sanitization (they are user-configured and safe)
safe_headers.update(self.auth_headers)
return safe_headers

async def send_request(
self, method: str, params: dict[str, Any | None] | None = None
Expand All @@ -42,7 +53,7 @@ async def send_raw(self, payload: dict[str, Any]) -> Any:
"Network to non-local host is disallowed by safety policy",
context={"url": self.url},
)
safe_headers = sanitize_headers(self.headers)
safe_headers = self._prepare_headers_with_auth(self.headers)
response = await client.post(self.url, json=payload, headers=safe_headers)
response.raise_for_status()
buffer: list[str] = []
Expand Down Expand Up @@ -107,7 +118,7 @@ async def send_notification(
"Network to non-local host is disallowed by safety policy",
context={"url": self.url},
)
safe_headers = sanitize_headers(self.headers)
safe_headers = self._prepare_headers_with_auth(self.headers)
response = await client.post(self.url, json=payload, headers=safe_headers)
response.raise_for_status()

Expand All @@ -130,7 +141,7 @@ async def _stream_request(self, payload: dict[str, Any]):
"Network to non-local host is disallowed by safety policy",
context={"url": self.url},
)
safe_headers = sanitize_headers(self.headers)
safe_headers = self._prepare_headers_with_auth(self.headers)
async with client.stream(
"POST",
self.url,
Expand Down
21 changes: 16 additions & 5 deletions mcp_fuzzer/transport/streamable_http.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,17 +51,28 @@ def __init__(
url: str,
timeout: float = DEFAULT_TIMEOUT,
auth_headers: dict[str, str | None] = None,
safety_enabled: bool = True,
):
self.url = url
self.timeout = timeout
self.safety_enabled = safety_enabled
self.headers: dict[str, str] = {
"Accept": DEFAULT_HTTP_ACCEPT,
"Content-Type": JSON_CT,
}
if auth_headers:
self.headers.update(auth_headers)
self.auth_headers = {k: v for k, v in (auth_headers or {}).items() if v is not None}
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Line exceeds maximum length.

Line 63 exceeds the project's 88-character limit (92 characters).

Apply this diff:

-    self.auth_headers = {k: v for k, v in (auth_headers or {}).items() if v is not None}
+    self.auth_headers = {
+        k: v for k, v in (auth_headers or {}).items() if v is not None
+    }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
self.auth_headers = {k: v for k, v in (auth_headers or {}).items() if v is not None}
self.auth_headers = {
k: v for k, v in (auth_headers or {}).items() if v is not None
}
🧰 Tools
🪛 GitHub Actions: Lint

[error] 63-63: E501 Line too long (92 > 88)

🤖 Prompt for AI Agents 
In mcp_fuzzer/transport/streamable_http.py around line 63, the dict
comprehension assignment for self.auth_headers exceeds the 88-character line
length limit; split the expression across multiple lines (for example assign the
source dict to a short local variable or break the dict comprehension into a
multi-line comprehension with each k, v pair on its own line) so the resulting
lines are under 88 characters while preserving the same behavior filtering out
None values.


self._logger = logging.getLogger(__name__)

def _prepare_headers_with_auth(self, headers: dict[str, str]) -> dict[str, str]:
"""Prepare headers with optional safety sanitization and auth headers."""
if self.safety_enabled:
safe_headers = self._prepare_headers_with_auth(headers)
else:
safe_headers = headers.copy()
# Add auth headers after sanitization (they are user-configured and safe)
safe_headers.update(self.auth_headers)
return safe_headers
Comment on lines +67 to +75
Copy link
Contributor

@coderabbitai coderabbitai bot Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Critical: Infinite recursion in _prepare_headers_with_auth.

Line 70 calls self._prepare_headers_with_auth(headers) recursively within the same method when safety_enabled is True, causing infinite recursion and a stack overflow.

Apply this diff to call sanitize_headers instead (matching the pattern in sse.py and http.py):

 def _prepare_headers_with_auth(self, headers: dict[str, str]) -> dict[str, str]:
     """Prepare headers with optional safety sanitization and auth headers."""
     if self.safety_enabled:
-        safe_headers = self._prepare_headers_with_auth(headers)
+        safe_headers = sanitize_headers(headers)
     else:
         safe_headers = headers.copy()
     # Add auth headers after sanitization (they are user-configured and safe)
     safe_headers.update(self.auth_headers)
     return safe_headers
-    self.session_id: str | None = None
-    self.protocol_version: str | None = None
-    self._initialized: bool = False
-    self._init_lock: asyncio.Lock = asyncio.Lock()
-    self._initializing: bool = False

Note: The code snippet also shows lines 76-80 that appear to be misplaced (they should be in __init__). Please verify the file structure.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
def _prepare_headers_with_auth(self, headers: dict[str, str]) -> dict[str, str]:
"""Prepare headers with optional safety sanitization and auth headers."""
if self.safety_enabled:
safe_headers = self._prepare_headers_with_auth(headers)
else:
safe_headers = headers.copy()
# Add auth headers after sanitization (they are user-configured and safe)
safe_headers.update(self.auth_headers)
return safe_headers
def _prepare_headers_with_auth(self, headers: dict[str, str]) -> dict[str, str]:
"""Prepare headers with optional safety sanitization and auth headers."""
if self.safety_enabled:
safe_headers = sanitize_headers(headers)
else:
safe_headers = headers.copy()
# Add auth headers after sanitization (they are user-configured and safe)
safe_headers.update(self.auth_headers)
return safe_headers
🤖 Prompt for AI Agents 
In mcp_fuzzer/transport/streamable_http.py around lines 67 to 75, the method
_prepare_headers_with_auth currently calls itself when self.safety_enabled is
True which causes infinite recursion; replace that recursive call with a call to
the header sanitizer (e.g., self.sanitize_headers(headers)) so headers are
sanitized instead of recursing, then update safe_headers accordingly and merge
auth headers as before; additionally inspect lines ~76-80 which look like
constructor-only initialization and move any misplaced initialization into
__init__ so file structure and class initialization are correct.

self.session_id: str | None = None
self.protocol_version: str | None = None
self._initialized: bool = False
Expand Down Expand Up @@ -211,7 +222,7 @@ async def send_raw(self, payload: dict[str, Any]) -> Any:
) as client:
self._ensure_host_allowed()
response = await self._post_with_retries(
client, self.url, payload, sanitize_headers(headers)
client, self.url, payload, self._prepare_headers_with_auth(headers)
)
# Handle redirect by retrying once with provided Location or trailing slash
redirect_url = self._resolve_redirect(response)
Expand Down Expand Up @@ -314,7 +325,7 @@ async def send_notification(
timeout=self.timeout, follow_redirects=False, trust_env=False
) as client:
self._ensure_host_allowed()
safe_headers = sanitize_headers(headers)
safe_headers = self._prepare_headers_with_auth(headers)
response = await self._post_with_retries(
client, self.url, payload, safe_headers
)
Expand Down Expand Up @@ -422,7 +433,7 @@ async def _stream_request(self, payload: dict[str, Any]):
timeout=self.timeout, follow_redirects=False, trust_env=False
) as client:
self._ensure_host_allowed()
safe_headers = sanitize_headers(headers)
safe_headers = self._prepare_headers_with_auth(headers)
response = await client.stream(
"POST", self.url, json=payload, headers=safe_headers
)
Expand Down
Loading