diff --git a/litellm/proxy/guardrails/guardrail_hooks/ztds.py b/litellm/proxy/guardrails/guardrail_hooks/ztds.py index 691003727d7..73f4b38b5d6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ztds.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ztds.py @@ -36,16 +36,20 @@ class ZTDSGuardrail(CustomGuardrail): and reverses tokens on completion return without external network egress. """ + TOKEN_PATTERN: ClassVar[re.Pattern] = re.compile(r"\[[A-Z_]+_TOKEN_[a-zA-Z0-9_-]+\]") + # Comprehensive zero-egress regex patterns for sensitive identifiers PATTERNS: ClassVar[dict[str, re.Pattern]] = { - "EMAIL": re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,24}\b"), + "EMAIL": re.compile( + r"\b[A-Za-z0-9._%+-]{1,64}@[A-Za-z0-9-]{1,63}(?:\.[A-Za-z0-9-]{1,63})*\.[A-Za-z]{2,24}\b" + ), "IPV4": re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"), "IBAN": re.compile(r"\b[A-Z]{2}[0-9]{2}[A-Z0-9]{4}[0-9]{7}([A-Z0-9]?){0,16}\b"), "CREDIT_CARD": re.compile(r"\b(?:\d{4}[-\s]?){3}\d{4}\b"), "SSN": re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), "PHONE": re.compile(r"\b(?:\+?\d{1,3}[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b"), "API_SECRET": re.compile( - r"\b(?:sk-[a-zA-Z0-9]{20,}|ghp_[a-zA-Z0-9]{20,}|eyJ[a-zA-Z0-9_-]{20,}\.[a-zA-Z0-9_-]{20,}\.[a-zA-Z0-9_-]{20,})\b" + r"\b(?:sk-(?:proj-)?[A-Za-z0-9_-]{20,}|ghp_[A-Za-z0-9]{20,}|eyJ[A-Za-z0-9_-]{20,}\.[A-Za-z0-9_-]{20,}\.[A-Za-z0-9_-]{20,})\b" ), } @@ -90,6 +94,9 @@ class ZTDSGuardrail(CustomGuardrail): sanitized = text for entity_type in self.enabled_entities: + if entity_type == "EMAIL" and "@" not in sanitized: + continue + pattern = self.PATTERNS.get(entity_type) if not pattern: continue @@ -123,7 +130,7 @@ class ZTDSGuardrail(CustomGuardrail): def restore_text(self, text: str, session_id: str) -> str: """ - Restores deterministic surrogates back to original cleartext. + Restores deterministic surrogates back to original cleartext via single-pass token dispatch. Enforces provenance isolation: only restores tokens that originated from caller-visible fields. Hidden/system prompt secrets are never reversed in caller output. """ @@ -132,12 +139,13 @@ class ZTDSGuardrail(CustomGuardrail): if not token_map: return text - restored = text - for token in sorted(token_map.keys(), key=len, reverse=True): - # Only restore if token was authorized from caller-visible inputs - if token in caller_tokens: - restored = restored.replace(token, token_map[token]) - return restored + def _replace_token(match: re.Match) -> str: + tok = match.group(0) + if tok in token_map and tok in caller_tokens: + return token_map[tok] + return tok + + return self.TOKEN_PATTERN.sub(_replace_token, text) def zeroize_session(self, session_id: str) -> None: """ diff --git a/tests/guardrails_tests/test_ztds_guardrail.py b/tests/guardrails_tests/test_ztds_guardrail.py index 2961807a665..38c3643717b 100644 --- a/tests/guardrails_tests/test_ztds_guardrail.py +++ b/tests/guardrails_tests/test_ztds_guardrail.py @@ -222,6 +222,30 @@ class TestZTDSLiteLLMGuardrail(unittest.IsolatedAsyncioTestCase): self.assertNotIn(system_secret, caller_visible_output) self.assertIn("[API_SECRET_TOKEN_1]", caller_visible_output) + def test_redos_resistance(self): + """Verify that email pattern does not cause catastrophic backtracking on adversarial inputs.""" + import time + + session_id = "test-session-redos" + payload = "a." * 16000 # 32 KB adversarial payload + t0 = time.perf_counter() + sanitized, _ = self.guardrail.sanitize_text(payload, session_id) + elapsed = time.perf_counter() - t0 + + # Must execute sub-second without blocking event loop (typically < 0.02s) + self.assertLess(elapsed, 0.1, f"ReDoS vulnerability detected: execution took {elapsed:.4f}s") + self.assertEqual(sanitized, payload) + + def test_modern_openai_project_keys(self): + """Verify detection and sanitization of modern OpenAI sk-proj- and hyphenated API tokens.""" + session_id = "test-session-keys" + secret = "sk-proj-abc-123_45678901234567890" + raw = f"Use OpenAI project key {secret} for deployment." + sanitized, _ = self.guardrail.sanitize_text(raw, session_id) + + self.assertNotIn(secret, sanitized) + self.assertIn("[API_SECRET_TOKEN_1]", sanitized) + def test_token_collision_avoidance(self): """Literal surrogate tokens in input text must not collide with generated tokens.""" session_id = "test-session-collision"