Skip to content

Commit 05094fc

Browse files
lucarliggcgoncalves
lucarlig
authored and
gcgoncalves
committed
feat: add experimental rust request logging masking extension (#4030)
* feat: add rust request logging masking sidecar Signed-off-by: lucarlig <luca.carlig@ibm.com> * perf: cache key sensitivity in rust masking sidecar Signed-off-by: lucarlig <luca.carlig@ibm.com> * feat: restore pyo3 request logging path Signed-off-by: lucarlig <luca.carlig@ibm.com> * refactor: rename request logging native extension Signed-off-by: lucarlig <luca.carlig@ibm.com> * perf: add native json request logging fast path Signed-off-by: lucarlig <luca.carlig@ibm.com> * test: cover request logging native fallback paths Signed-off-by: lucarlig <luca.carlig@ibm.com> * feat: generate stubs for request logging native extension Signed-off-by: lucarlig <luca.carlig@ibm.com> * perf: reuse request logging key sensitivity cache Signed-off-by: lucarlig <luca.carlig@ibm.com> * chore: normalize detect-secrets baseline after rebase Signed-off-by: lucarlig <luca.carlig@ibm.com> --------- Signed-off-by: lucarlig <luca.carlig@ibm.com>
1 parent 34e1a9f commit 05094fc

13 files changed

Lines changed: 2172 additions & 75 deletions

File tree

‎.secrets.baseline‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5814,7 +5814,7 @@
58145814
"hashed_secret": "ff37a98a9963d347e9749a5c1b3936a4a245a6ff",
58155815
"is_secret": false,
58165816
"is_verified": false,
5817-
"line_number": 2221,
5817+
"line_number": 2225,
58185818
"type": "Secret Keyword",
58195819
"verified_result": null
58205820
}
@@ -5860,23 +5860,23 @@
58605860
"hashed_secret": "f2b14f68eb995facb3a1c35287b778d5bd785511",
58615861
"is_secret": false,
58625862
"is_verified": false,
5863-
"line_number": 171,
5863+
"line_number": 175,
58645864
"type": "Secret Keyword",
58655865
"verified_result": null
58665866
},
58675867
{
58685868
"hashed_secret": "1073ab6cda4b991cd29f9e83a307f34004ae9327",
58695869
"is_secret": false,
58705870
"is_verified": false,
5871-
"line_number": 177,
5871+
"line_number": 181,
58725872
"type": "Secret Keyword",
58735873
"verified_result": null
58745874
},
58755875
{
58765876
"hashed_secret": "d4fe581561f18ee5006254a7e53f53cbed780bc2",
58775877
"is_secret": false,
58785878
"is_verified": false,
5879-
"line_number": 259,
5879+
"line_number": 268,
58805880
"type": "Secret Keyword",
58815881
"verified_result": null
58825882
}
@@ -7964,23 +7964,23 @@
79647964
"hashed_secret": "1a91d62f7ca67399625a4368a6ab5d4a3baa6073",
79657965
"is_secret": false,
79667966
"is_verified": false,
7967-
"line_number": 98,
7967+
"line_number": 115,
79687968
"type": "Secret Keyword",
79697969
"verified_result": null
79707970
},
79717971
{
79727972
"hashed_secret": "dbdab9be92cacdae6a97e8601332bfaa8545800f",
79737973
"is_secret": false,
79747974
"is_verified": false,
7975-
"line_number": 99,
7975+
"line_number": 116,
79767976
"type": "Secret Keyword",
79777977
"verified_result": null
79787978
},
79797979
{
79807980
"hashed_secret": "4ea8d2335b430796cf3f500368c5b0f5b1dc90f5",
79817981
"is_secret": false,
79827982
"is_verified": false,
7983-
"line_number": 151,
7983+
"line_number": 213,
79847984
"type": "Secret Keyword",
79857985
"verified_result": null
79867986
}

‎mcpgateway/config.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -350,6 +350,10 @@ class Settings(BaseSettings):
350350

351351
# Security Validation & Sanitization
352352
experimental_validate_io: bool = Field(default=False, description="Enable experimental input validation and output sanitization")
353+
experimental_rust_request_logging_masking_enabled: bool = Field(
354+
default=False,
355+
description="Enable experimental Rust native extension for request logging sensitive-data masking",
356+
)
353357
validation_middleware_enabled: bool = Field(default=False, description="Enable validation middleware for all requests")
354358
validation_strict: bool = Field(default=True, description="Strict validation mode - reject on violations")
355359
sanitize_output: bool = Field(default=True, description="Sanitize output to remove control characters")

‎mcpgateway/middleware/request_logging_middleware.py‎

Lines changed: 54 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
"""
4747

4848
# Standard
49+
import importlib
4950
import logging
5051
import re
5152
import secrets
@@ -74,6 +75,9 @@
7475
# Initialize structured logger for gateway boundary logging
7576
structured_logger = get_structured_logger("http_gateway")
7677

78+
_RUST_REQUEST_LOGGING_MODULE = None
79+
_RUST_REQUEST_LOGGING_IMPORT_FAILED = False
80+
7781
SENSITIVE_KEYS = frozenset(
7882
{
7983
"password",
@@ -183,6 +187,11 @@ def mask_sensitive_data(data, max_depth: int = 10):
183187
>>> mask_sensitive_data({"level": {"nested": {}}}, max_depth=1)
184188
{'level': '<nested too deep>'}
185189
"""
190+
if getattr(settings, "experimental_rust_request_logging_masking_enabled", False) is True:
191+
rust_module = _load_rust_request_logging_module()
192+
if rust_module is not None:
193+
return rust_module.mask_sensitive_data(data, max_depth)
194+
186195
if max_depth <= 0:
187196
return "<nested too deep>"
188197

@@ -262,6 +271,11 @@ def mask_sensitive_headers(headers):
262271
>>> "******" in result["Cookie"]
263272
True
264273
"""
274+
if getattr(settings, "experimental_rust_request_logging_masking_enabled", False) is True:
275+
rust_module = _load_rust_request_logging_module()
276+
if rust_module is not None:
277+
return rust_module.mask_sensitive_headers(headers)
278+
265279
masked_headers = {}
266280
for key, value in headers.items():
267281
key_lower = key.lower()
@@ -275,6 +289,43 @@ def mask_sensitive_headers(headers):
275289
return masked_headers
276290

277291

292+
def _mask_json_payload_for_logging(payload: bytes, max_depth: int = 10) -> str:
293+
"""Mask a JSON payload for logging, preferring the native bytes fast path."""
294+
if getattr(settings, "experimental_rust_request_logging_masking_enabled", False) is True:
295+
rust_module = _load_rust_request_logging_module()
296+
if rust_module is not None and hasattr(rust_module, "mask_sensitive_json_bytes"):
297+
try:
298+
masked_payload = rust_module.mask_sensitive_json_bytes(payload, max_depth)
299+
if isinstance(masked_payload, bytes):
300+
return masked_payload.decode("utf-8", errors="ignore")
301+
return str(masked_payload)
302+
except Exception:
303+
pass
304+
305+
json_payload = orjson.loads(payload)
306+
payload_to_log = mask_sensitive_data(json_payload, max_depth)
307+
return orjson.dumps(payload_to_log).decode()
308+
309+
310+
def _load_rust_request_logging_module():
311+
"""Load the experimental Rust masking native extension on demand."""
312+
global _RUST_REQUEST_LOGGING_IMPORT_FAILED, _RUST_REQUEST_LOGGING_MODULE
313+
314+
if _RUST_REQUEST_LOGGING_MODULE is None:
315+
if _RUST_REQUEST_LOGGING_IMPORT_FAILED:
316+
return None
317+
try:
318+
_RUST_REQUEST_LOGGING_MODULE = importlib.import_module("request_logging_masking_native_extension")
319+
except ImportError as exc:
320+
_RUST_REQUEST_LOGGING_IMPORT_FAILED = True
321+
logger.warning(
322+
f"Experimental Rust request logging masking is enabled but the native extension is unavailable; "
323+
f"falling back to Python masking. Install the request logging masking native extension first. ({exc})"
324+
)
325+
return None
326+
return _RUST_REQUEST_LOGGING_MODULE
327+
328+
278329
class RequestLoggingMiddleware(BaseHTTPMiddleware):
279330
"""Middleware for logging HTTP requests with sensitive data masking.
280331
@@ -597,16 +648,13 @@ async def dispatch(self, request: Request, call_next: Callable):
597648
truncated = False
598649
body_to_log = body
599650

600-
payload = body_to_log.decode("utf-8", errors="ignore").strip()
651+
payload = body_to_log.strip()
601652
if payload:
602653
try:
603-
json_payload = orjson.loads(payload)
604-
payload_to_log = mask_sensitive_data(json_payload)
605-
# Use orjson without indent for performance (compact output)
606-
payload_str = orjson.dumps(payload_to_log).decode()
654+
payload_str = _mask_json_payload_for_logging(payload)
607655
except orjson.JSONDecodeError:
608656
# For non-JSON payloads, still mask potential sensitive data
609-
payload_str = payload
657+
payload_str = payload.decode("utf-8", errors="ignore")
610658
for sensitive_key in SENSITIVE_KEYS:
611659
if sensitive_key in payload_str.lower():
612660
payload_str = "<contains sensitive data - masked>"

‎tests/performance/test_request_logging_masking_native_extension_benchmark.py‎

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
# -*- coding: utf-8 -*-
2+
"""Benchmark the request-logging masking Rust native extension against the Python path."""
3+
4+
# Standard
5+
from __future__ import annotations
6+
7+
import importlib
8+
import statistics
9+
import subprocess
10+
import time
11+
from dataclasses import dataclass
12+
from pathlib import Path
13+
from typing import Any, Callable
14+
15+
# First-Party
16+
from mcpgateway.config import settings
17+
from mcpgateway.middleware import request_logging_middleware
18+
from mcpgateway.middleware.request_logging_middleware import mask_sensitive_data, mask_sensitive_headers
19+
20+
REPO_ROOT = Path(__file__).resolve().parents[2]
21+
NATIVE_EXTENSION_MANIFEST = REPO_ROOT / "tools_rust" / "request_logging_masking_native_extension" / "Cargo.toml"
22+
23+
24+
@dataclass(frozen=True)
25+
class BenchmarkScenario:
26+
name: str
27+
python_fn: Callable[[Any], Any]
28+
public_native_fn: Callable[[Any], Any]
29+
direct_native_fn: Callable[[Any], Any]
30+
payload: Any
31+
iterations: int
32+
33+
34+
def _build_scenarios(native_extension: Any) -> list[BenchmarkScenario]:
35+
def python_data(payload: Any) -> Any:
36+
settings.experimental_rust_request_logging_masking_enabled = False
37+
return mask_sensitive_data(payload, 12)
38+
39+
def public_native_data(payload: Any) -> Any:
40+
settings.experimental_rust_request_logging_masking_enabled = True
41+
return mask_sensitive_data(payload, 12)
42+
43+
def direct_native_data(payload: Any) -> Any:
44+
return native_extension.mask_sensitive_data(payload, 12)
45+
46+
def python_headers(payload: Any) -> Any:
47+
settings.experimental_rust_request_logging_masking_enabled = False
48+
return mask_sensitive_headers(payload)
49+
50+
def public_native_headers(payload: Any) -> Any:
51+
settings.experimental_rust_request_logging_masking_enabled = True
52+
return mask_sensitive_headers(payload)
53+
54+
def direct_native_headers(payload: Any) -> Any:
55+
return native_extension.mask_sensitive_headers(payload)
56+
57+
return [
58+
BenchmarkScenario(
59+
name="nested_payload_masking",
60+
python_fn=python_data,
61+
public_native_fn=public_native_data,
62+
direct_native_fn=direct_native_data,
63+
payload={
64+
"events": [
65+
{
66+
"actor": {"userName": f"user-{index}", "sessionToken": f"token-{index}", "sessionCount": index},
67+
"request": {
68+
"clientSecret": f"secret-{index}",
69+
"payload": {"safeField": "value" * 8, "authDevice": f"device-{index}", "auth_count": index},
70+
},
71+
}
72+
for index in range(1024)
73+
]
74+
},
75+
iterations=120,
76+
),
77+
BenchmarkScenario(
78+
name="headers_masking",
79+
python_fn=python_headers,
80+
public_native_fn=public_native_headers,
81+
direct_native_fn=direct_native_headers,
82+
payload={
83+
**{f"X-Custom-{index}": f"value-{index}" for index in range(512)},
84+
**{f"X-Api-Key-{index}": f"secret-{index}" for index in range(256)},
85+
"Cookie": "; ".join([f"jwt_token_{index}=abc{index}" for index in range(128)]),
86+
},
87+
iterations=300,
88+
),
89+
]
90+
91+
92+
def test_build_scenarios_exposes_public_and_direct_native_paths(monkeypatch):
93+
native_extension = type(
94+
"NativeExtension",
95+
(),
96+
{
97+
"mask_sensitive_data": staticmethod(lambda payload, max_depth: {"path": "direct", "payload": payload, "max_depth": max_depth}),
98+
"mask_sensitive_headers": staticmethod(lambda payload: {"path": "direct", "payload": payload}),
99+
},
100+
)()
101+
observed_flags: list[bool] = []
102+
103+
def fake_mask_sensitive_data(payload: Any, max_depth: int = 10) -> Any:
104+
observed_flags.append(settings.experimental_rust_request_logging_masking_enabled)
105+
return {"path": "public" if settings.experimental_rust_request_logging_masking_enabled else "python", "max_depth": max_depth}
106+
107+
def fake_mask_sensitive_headers(payload: Any) -> Any:
108+
observed_flags.append(settings.experimental_rust_request_logging_masking_enabled)
109+
return {"path": "public" if settings.experimental_rust_request_logging_masking_enabled else "python"}
110+
111+
monkeypatch.setattr("tests.performance.test_request_logging_masking_native_extension_benchmark.mask_sensitive_data", fake_mask_sensitive_data)
112+
monkeypatch.setattr("tests.performance.test_request_logging_masking_native_extension_benchmark.mask_sensitive_headers", fake_mask_sensitive_headers)
113+
114+
scenarios = {scenario.name: scenario for scenario in _build_scenarios(native_extension)}
115+
116+
public_data = scenarios["nested_payload_masking"].public_native_fn({"k": "v"})
117+
direct_data = scenarios["nested_payload_masking"].direct_native_fn({"k": "v"})
118+
python_data = scenarios["nested_payload_masking"].python_fn({"k": "v"})
119+
public_headers = scenarios["headers_masking"].public_native_fn({"Authorization": "Bearer x"})
120+
direct_headers = scenarios["headers_masking"].direct_native_fn({"Authorization": "Bearer x"})
121+
122+
assert public_data["path"] == "public"
123+
assert direct_data["path"] == "direct"
124+
assert python_data["path"] == "python"
125+
assert public_headers["path"] == "public"
126+
assert direct_headers["path"] == "direct"
127+
assert observed_flags == [True, False, True]
128+
129+
130+
def _ensure_native_extension_installed() -> Any:
131+
subprocess.run(["uv", "run", "maturin", "develop", "--release", "--manifest-path", str(NATIVE_EXTENSION_MANIFEST)], check=True, cwd=REPO_ROOT)
132+
return importlib.import_module("request_logging_masking_native_extension")
133+
134+
135+
def _measure(label: str, fn: Callable[[Any], Any], payload: Any, iterations: int) -> tuple[float, float]:
136+
samples = []
137+
for _ in range(iterations):
138+
started = time.perf_counter_ns()
139+
fn(payload)
140+
samples.append(time.perf_counter_ns() - started)
141+
142+
median_ms = statistics.median(samples) / 1_000_000
143+
p95_ms = statistics.quantiles(samples, n=100)[94] / 1_000_000
144+
print(f"{label}: median={median_ms:.3f}ms p95={p95_ms:.3f}ms")
145+
return median_ms, p95_ms
146+
147+
148+
def _assert_parity(python_fn: Callable[[Any], Any], rust_fn: Callable[[Any], Any], payloads: list[Any]) -> None:
149+
for payload in payloads:
150+
python_result = python_fn(payload)
151+
rust_result = rust_fn(payload)
152+
if python_result != rust_result:
153+
raise AssertionError(f"Parity mismatch for payload {payload!r}: python={python_result!r} rust={rust_result!r}")
154+
155+
156+
def _prepare_public_native_path(native_extension: Any) -> None:
157+
settings.experimental_rust_request_logging_masking_enabled = True
158+
request_logging_middleware._RUST_REQUEST_LOGGING_MODULE = native_extension
159+
request_logging_middleware._RUST_REQUEST_LOGGING_IMPORT_FAILED = False
160+
161+
162+
def main() -> None:
163+
native_extension = _ensure_native_extension_installed()
164+
_prepare_public_native_path(native_extension)
165+
166+
scenarios = _build_scenarios(native_extension)
167+
168+
_assert_parity(
169+
scenarios[0].python_fn,
170+
scenarios[0].direct_native_fn,
171+
[
172+
{"password": "secret", "nested": {"authToken": "abc", "ok": "value"}},
173+
{"token_count": 3, "tokenizer": "ok", "privateKey": "secret"},
174+
[{"jwt_token": "abc"}, {"normal": "value"}],
175+
],
176+
)
177+
_assert_parity(
178+
scenarios[1].python_fn,
179+
scenarios[1].direct_native_fn,
180+
[
181+
{"Authorization": "Bearer abc", "Cookie": "jwt_token=abc; theme=dark", "X-Trace-Id": "123"},
182+
{"X-Auth-Count": "5", "X-Api-Key": "secret"},
183+
],
184+
)
185+
186+
for scenario in scenarios:
187+
print(f"\n{scenario.name} ({scenario.iterations} iterations)")
188+
python_median, _ = _measure("python", scenario.python_fn, scenario.payload, scenario.iterations)
189+
public_native_median, _ = _measure("public_native", scenario.public_native_fn, scenario.payload, scenario.iterations)
190+
direct_native_median, _ = _measure("direct_native", scenario.direct_native_fn, scenario.payload, scenario.iterations)
191+
print(f"public_speedup={python_median / public_native_median:.2f}x")
192+
print(f"direct_speedup={python_median / direct_native_median:.2f}x")
193+
print(f"public_overhead={(public_native_median / direct_native_median - 1) * 100:.1f}%")
194+
195+
196+
if __name__ == "__main__":
197+
main()

0 commit comments

Comments
 (0)