perf: tighten rust masking parity and hot path

Signed-off-by: lucarlig <luca.carlig@ibm.com>

Author: lucarlig

.secrets.baseline

@@ -3,7 +3,7 @@
     "files": null,
     "lines": null
   },
-  "generated_at": "2026-04-03T14:47:08Z",
+  "generated_at": "2026-04-03T15:25:49Z",
   "plugins_used": [
     {
       "name": "AWSKeyDetector"
@@ -22032,7 +22032,7 @@
         "hashed_secret": "4ea8d2335b430796cf3f500368c5b0f5b1dc90f5",
         "is_secret": false,
         "is_verified": false,
-        "line_number": 207,
+        "line_number": 221,
         "type": "Secret Keyword",
         "verified_result": null
       }

tests/performance/test_request_logging_masking_sidecar_benchmark.py

@@ -76,9 +76,9 @@ def _measure_once(label: str, fn: Callable[[Any], Any], payload: Any) -> float:
     return elapsed_ms
 
 
-def _run_with_rust_masking_enabled(enabled: bool, fn: Callable[[Any], Any], payload: Any) -> Any:
+def _run_with_rust_masking_enabled(enabled: bool, fn: Callable[..., Any], *args: Any) -> Any:
     with _rust_masking_enabled(enabled):
-        return fn(payload)
+        return fn(*args)
 
 
 def _measure_with_rust_masking_enabled(enabled: bool, label: str, fn: Callable[[Any], Any], payload: Any, iterations: int) -> tuple[float, float]:
@@ -94,6 +94,14 @@ def _assert_parity(python_fn: Callable[[Any], Any], rust_fn: Callable[[Any], Any
             raise AssertionError(f"Parity mismatch for payload {payload!r}: python={python_result!r} rust={rust_result!r}")
 
 
+def _assert_depth_parity(python_fn: Callable[[Any, int], Any], rust_fn: Callable[[Any, int], Any], cases: list[tuple[Any, int]]) -> None:
+    for payload, max_depth in cases:
+        python_result = python_fn(payload, max_depth)
+        rust_result = rust_fn(payload, max_depth)
+        if python_result != rust_result:
+            raise AssertionError(f"Depth parity mismatch for payload {payload!r} at max_depth={max_depth}: " f"python={python_result!r} rust={rust_result!r}")
+
+
 def main() -> None:
     sidecar = _ensure_sidecar_installed()
     request_logging_middleware._RUST_REQUEST_LOGGING_MODULE = sidecar
@@ -104,6 +112,12 @@ def python_data(payload: Any) -> Any:
     def rust_data(payload: Any) -> Any:
         return mask_sensitive_data(payload, 12)
 
+    def python_data_with_depth(payload: Any, max_depth: int) -> Any:
+        return mask_sensitive_data(payload, max_depth)
+
+    def rust_data_with_depth(payload: Any, max_depth: int) -> Any:
+        return mask_sensitive_data(payload, max_depth)
+
     def python_headers(payload: Any) -> Any:
         return mask_sensitive_headers(payload)
 
@@ -121,6 +135,17 @@ def rust_headers(payload: Any) -> Any:
             {"password": "secret", "nested": {"authToken": "abc", "ok": "value"}},
             {"token_count": 3, "tokenizer": "ok", "privateKey": "secret"},
             [{"jwt_token": "abc"}, {"normal": "value"}],
+            {"outer": [{1: "numeric-key", "authToken": "abc"}, {"nested": {("tuple", 2): "tuple-key", "clientSecret": "hidden"}}]},  # pragma: allowlist secret
+            {"outer": [{"leaf": "value"}, {"nested": ["safe", {"secret": "x"}]}]},
+        ],
+    )
+    _assert_depth_parity(
+        lambda payload, max_depth: _run_with_rust_masking_enabled(False, python_data_with_depth, payload, max_depth),
+        lambda payload, max_depth: _run_with_rust_masking_enabled(True, rust_data_with_depth, payload, max_depth),
+        [
+            ({"level": "x"}, 1),
+            ([{"leaf": "x"}], 1),
+            ({"outer": [{"leaf": "value"}, {"nested": ["safe", {"secret": "x"}]}]}, 2),
         ],
     )
     _assert_parity(

tests/unit/mcpgateway/middleware/test_request_logging_middleware.py

@@ -125,6 +125,20 @@ def test_mask_sensitive_data_ignores_empty_normalized_keys():
     assert masked["password"] == "******"
 
 
+def test_mask_sensitive_data_preserves_non_string_keys_in_nested_structures():
+    data = {
+        "outer": [
+            {1: "numeric-key", "authToken": "secret"},
+            {"nested": {("tuple", 2): "tuple-key", "clientSecret": "hidden"}},  # pragma: allowlist secret
+        ]
+    }
+    masked = mask_sensitive_data(data)
+    assert masked["outer"][0][1] == "numeric-key"
+    assert masked["outer"][0]["authToken"] == "******"
+    assert masked["outer"][1]["nested"][("tuple", 2)] == "tuple-key"
+    assert masked["outer"][1]["nested"]["clientSecret"] == "******"
+
+
 def test_mask_sensitive_data_uses_rust_sidecar_when_enabled(monkeypatch):
     rust_module = MagicMock()
     rust_module.mask_sensitive_data.return_value = {"password": "******", "username": "user"}  # pragma: allowlist secret

tools_rust/request_logging_masking_sidecar/src/lib.rs

@@ -121,6 +121,15 @@ fn is_sensitive_key_cached(key: &str, cache: &mut HashMap<String, bool>) -> bool
     result
 }
 
+fn py_key_is_sensitive(
+    key: &Bound<'_, PyAny>,
+    cache: &mut HashMap<String, bool>,
+) -> PyResult<bool> {
+    let key_string = key.str()?;
+    let key_text = key_string.to_string_lossy();
+    Ok(is_sensitive_key_cached(key_text.as_ref(), cache))
+}
+
 fn mask_cookie_header(cookie_header: &str) -> String {
     let mut masked = Vec::new();
 
@@ -159,8 +168,7 @@ fn mask_sensitive_data_inner(
     if let Ok(dict) = data.cast::<PyDict>() {
         let masked = PyDict::new(py);
         for (key, value) in dict.iter() {
-            let key_string = key.str()?.to_string_lossy().into_owned();
-            if is_sensitive_key_cached(&key_string, key_cache) {
+            if py_key_is_sensitive(&key, key_cache)? {
                 masked.set_item(key, MASKED_VALUE)?;
             } else {
                 masked.set_item(
@@ -173,15 +181,11 @@ fn mask_sensitive_data_inner(
     }
 
     if let Ok(list) = data.cast::<PyList>() {
-        let masked = PyList::empty(py);
-        for item in list.iter() {
-            masked.append(mask_sensitive_data_inner(
-                py,
-                &item,
-                max_depth - 1,
-                key_cache,
-            )?)?;
-        }
+        let items = list
+            .iter()
+            .map(|item| mask_sensitive_data_inner(py, &item, max_depth - 1, key_cache))
+            .collect::<PyResult<Vec<_>>>()?;
+        let masked = PyList::new(py, items)?;
         return Ok(masked.into_any().unbind());
     }
 
@@ -194,7 +198,7 @@ fn mask_sensitive_data(
     data: &Bound<'_, PyAny>,
     max_depth: Option<i32>,
 ) -> PyResult<Py<PyAny>> {
-    let mut key_cache = HashMap::new();
+    let mut key_cache = HashMap::with_capacity(32);
     mask_sensitive_data_inner(py, data, max_depth.unwrap_or(10), &mut key_cache)
 }
 
@@ -205,13 +209,15 @@ fn mask_sensitive_headers(py: Python<'_>, headers: &Bound<'_, PyAny>) -> PyResul
     let mut key_cache = HashMap::with_capacity(source.len());
 
     for (key, value) in source.iter() {
-        let key_string = key.str()?.to_string_lossy().into_owned();
-        if is_sensitive_key_cached(&key_string, &mut key_cache) {
+        let key_string = key.str()?;
+        let key_text = key_string.to_string_lossy();
+
+        if is_sensitive_key_cached(key_text.as_ref(), &mut key_cache) {
             masked.set_item(key, MASKED_VALUE)?;
             continue;
         }
 
-        if key_string.eq_ignore_ascii_case("cookie") && value.is_instance_of::<PyString>() {
+        if key_text.eq_ignore_ascii_case("cookie") && value.is_instance_of::<PyString>() {
             let cookie_value = value.cast::<PyString>()?.to_str()?;
             masked.set_item(key, mask_cookie_header(cookie_value))?;
             continue;