fix(guardrails): preserve stdlib regex semantics in isolated workers

litellm/proxy/guardrails/guardrail_hooks/custom_code/_regex_worker.py

@@ -0,0 +1,42 @@
+"""Run stdlib regex in isolation so a caller can terminate pathological matches."""
+
+import json
+import re
+import sys
+from contextlib import redirect_stdout
+from typing import Final, TextIO
+
+
+def _execute(operation: str, text: str, pattern: str, flags: int, replacement: str) -> object:
+    match operation:
+        case "search":
+            return bool(re.search(pattern, text, flags))
+        case "fullmatch":
+            return bool(re.fullmatch(pattern, text, flags))
+        case "sub":
+            return re.sub(pattern, replacement, text, flags=flags)
+        case "findall":
+            return re.findall(pattern, text, flags)
+        case _:
+            raise ValueError("unknown regex operation")
+
+
+def _main(request: str, output: TextIO, errors: TextIO) -> int:
+    payload: Final[object] = json.loads(request)  # pyright: ignore[reportAny]  # stdlib JSON is validated below
+    match payload:
+        case [str(operation), str(text), str(pattern), int(flags), str(replacement)]:
+            pass
+        case _:
+            raise ValueError("invalid regex request")
+    try:
+        with redirect_stdout(errors):
+            result: Final = _execute(operation, text, pattern, flags, replacement)
+    except re.error as error:
+        errors.write(str(error))
+        return 65
+    output.write(json.dumps(result))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(_main(sys.stdin.read(), sys.stdout, sys.stderr))

litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py

@@ -8,13 +8,15 @@ and provide safe, sandboxed functionality for common guardrail operations.
 import asyncio
 import json
 import re
+import subprocess
+import sys
 from collections.abc import Mapping, Sequence
-from typing import Final, Literal
+from pathlib import Path
+from typing import Final, Literal, TypeVar
 from urllib.parse import urlparse
 
 import httpx
-import regex
-from pydantic import JsonValue
+from pydantic import JsonValue, TypeAdapter
 from typing_extensions import ReadOnly, TypedDict
 
 import litellm
@@ -113,24 +115,44 @@ def modify(
 # =============================================================================
 
 
-_REGEX_FLAGS: Final = (
-    (re.ASCII, regex.ASCII),
-    (re.IGNORECASE, regex.IGNORECASE),
-    (re.LOCALE, regex.LOCALE),
-    (re.UNICODE, regex.UNICODE),
-    (re.MULTILINE, regex.MULTILINE),
-    (re.DOTALL, regex.DOTALL),
-    (re.VERBOSE, regex.VERBOSE),
-    (re.TEMPLATE, regex.TEMPLATE),
-    (re.DEBUG, regex.DEBUG),
-)
+_RegexResult: Final = TypeVar("_RegexResult")
+_REGEX_MATCH: Final = TypeAdapter(bool)
+_REGEX_TEXT: Final = TypeAdapter(str)
+_REGEX_MATCHES: Final[TypeAdapter[Sequence[str | tuple[str, ...]]]] = TypeAdapter(Sequence[str | tuple[str, ...]])
+_REGEX_WORKER: Final = str(Path(__file__).with_name("_regex_worker.py"))
 
 
-def _regex_flags(pattern: str, flags: int) -> int:
-    validated: Final = re.compile(pattern, flags).flags
-    return regex.VERSION0 | sum(
-        int(safe_flag) for standard_flag, safe_flag in _REGEX_FLAGS if validated & standard_flag
-    )
+def _regex_result(
+    operation: Literal["search", "fullmatch", "sub", "findall"],
+    text: str,
+    pattern: str,
+    flags: int,
+    result_type: TypeAdapter[_RegexResult],
+    replacement: str = "",
+) -> _RegexResult:
+    remaining_execution_time()
+    with subprocess.Popen(
+        [sys.executable, "-I", "-S", _REGEX_WORKER],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        encoding="utf-8",
+        env={},
+    ) as worker:
+        try:
+            output, error = worker.communicate(
+                json.dumps((operation, text, pattern, flags, replacement)), timeout=remaining_execution_time()
+            )
+        except BaseException:  # noqa: BLE001  # sandbox interruption must also terminate and reap the child
+            worker.kill()
+            worker.communicate()
+            raise
+    if worker.returncode == 65:
+        raise re.error(error)
+    if worker.returncode:
+        raise RuntimeError(f"regex worker failed ({worker.returncode}): {error}")
+    return result_type.validate_json(output)
 
 
 def regex_match(text: str, pattern: str, flags: int = 0) -> bool:
@@ -146,14 +168,10 @@ def regex_match(text: str, pattern: str, flags: int = 0) -> bool:
         True if pattern matches, False otherwise
     """
     try:
-        return bool(
-            regex.search(
-                pattern, text, _regex_flags(pattern, flags), timeout=remaining_execution_time(), concurrent=True
-            )
-        )
-    except TimeoutError:
+        return _regex_result("search", text, pattern, flags, _REGEX_MATCH)
+    except subprocess.TimeoutExpired:
         raise ExecutionInterrupted from None
-    except (re.error, regex.error) as e:
+    except re.error as e:
         verbose_proxy_logger.warning("Starlark regex_match error: %s", e)
         return False
 
@@ -171,14 +189,10 @@ def regex_match_all(text: str, pattern: str, flags: int = 0) -> bool:
         True if pattern matches entire text, False otherwise
     """
     try:
-        return bool(
-            regex.fullmatch(
-                pattern, text, _regex_flags(pattern, flags), timeout=remaining_execution_time(), concurrent=True
-            )
-        )
-    except TimeoutError:
+        return _regex_result("fullmatch", text, pattern, flags, _REGEX_MATCH)
+    except subprocess.TimeoutExpired:
         raise ExecutionInterrupted from None
-    except (re.error, regex.error) as e:
+    except re.error as e:
         verbose_proxy_logger.warning("Starlark regex_match_all error: %s", e)
         return False
 
@@ -197,22 +211,15 @@ def regex_replace(text: str, pattern: str, replacement: str, flags: int = 0) ->
         The text with replacements applied
     """
     try:
-        return regex.sub(
-            pattern,
-            replacement,
-            text,
-            flags=_regex_flags(pattern, flags),
-            timeout=remaining_execution_time(),
-            concurrent=True,
-        )
-    except TimeoutError:
+        return _regex_result("sub", text, pattern, flags, _REGEX_TEXT, replacement)
+    except subprocess.TimeoutExpired:
         raise ExecutionInterrupted from None
-    except (re.error, regex.error) as e:
+    except re.error as e:
         verbose_proxy_logger.warning("Starlark regex_replace error: %s", e)
         return text
 
 
-def regex_find_all(text: str, pattern: str, flags: int = 0) -> list[str]:
+def regex_find_all(text: str, pattern: str, flags: int = 0) -> Sequence[str | tuple[str, ...]]:
     """
     Find all occurrences of a pattern in text.
 
@@ -225,12 +232,10 @@ def regex_find_all(text: str, pattern: str, flags: int = 0) -> list[str]:
         List of all matches
     """
     try:
-        return regex.findall(
-            pattern, text, _regex_flags(pattern, flags), timeout=remaining_execution_time(), concurrent=True
-        )
-    except TimeoutError:
+        return _regex_result("findall", text, pattern, flags, _REGEX_MATCHES)
+    except subprocess.TimeoutExpired:
         raise ExecutionInterrupted from None
-    except (re.error, regex.error) as e:
+    except re.error as e:
         verbose_proxy_logger.warning("Starlark regex_find_all error: %s", e)
         return []
 

tests/unit/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py

@@ -1,11 +1,14 @@
 import asyncio
+import json
 import re
 import threading
 import time
+from io import StringIO
 from typing import Final
 
 import pytest
 
+from litellm.proxy.guardrails.guardrail_hooks.custom_code._regex_worker import _main as regex_worker
 from litellm.proxy.guardrails.guardrail_hooks.custom_code.bounded_execution import (
     ExecutionTimeoutError,
     SandboxExit,
@@ -212,13 +215,10 @@ async def test_regex_near_match_does_not_hold_the_execution_timeout() -> None:
         execution_timeout=0.1,
     )
     started: Final = time.monotonic()
-    try:
-        result: Final = await guardrail.apply_guardrail(
+    with pytest.raises(CustomCodeExecutionError, match=r"exceeded its 0\.1s execution timeout"):
+        await guardrail.apply_guardrail(
             inputs={"texts": ["a" * 26 + "!"]}, request_data={"model": "m"}, input_type="request"
         )
-        assert result == {"texts": ["a" * 26 + "!"]}
-    except CustomCodeExecutionError as error:
-        assert "exceeded its 0.1s execution timeout" in str(error)
     await asyncio.to_thread(_join_worker_threads)
     assert time.monotonic() - started < 0.8
     assert _worker_threads() == []
@@ -263,6 +263,13 @@ async def test_regex_deadline_cannot_be_swallowed_as_a_failed_match(operation: s
     [
         ("café", r"\w+", re.ASCII),
         ("café", r"\w+", re.UNICODE),
+        ("123-45-6789\u200d", r"\b\d{3}-\d{2}-\d{4}\b", 0),
+        ("123-45-6789\u200c", r"\b\d{3}-\d{2}-\d{4}\b", 0),
+        ("a\u0301", r"\w+", 0),
+        ("²", r"\w+", 0),
+        ("\u001c", r"\s+", 0),
+        ("a\u200d", r"a\b\W", 0),
+        ("²", r"[^\W]+", 0),
         ("HELLO", "hello", re.IGNORECASE),
         ("first\nsecond", "^second$", re.MULTILINE),
         ("first\nsecond", ".+", re.DOTALL),
@@ -289,3 +296,32 @@ def test_regex_primitives_preserve_standard_matching_results(text: str, pattern:
         expected_replacement,
         expected_matches,
     )
+
+
+@pytest.mark.parametrize(
+    ("operation", "expected"),
+    [("search", True), ("fullmatch", False), ("sub", "[REDACTED]\u200d"), ("findall", ["123-45-6789"])],
+)
+def test_regex_worker_protocol_preserves_unicode_and_separates_debug_output(operation: str, expected: object) -> None:
+    output: Final = StringIO()
+    errors: Final = StringIO()
+    request: Final = json.dumps((operation, "123-45-6789\u200d", r"\b\d{3}-\d{2}-\d{4}\b", int(re.DEBUG), "[REDACTED]"))
+
+    assert regex_worker(request, output, errors) == 0
+    assert json.loads(output.getvalue()) == expected
+    assert errors.getvalue()
+
+
+def test_regex_worker_reports_invalid_patterns_without_a_success_result() -> None:
+    output: Final = StringIO()
+    errors: Final = StringIO()
+
+    assert regex_worker(json.dumps(("search", "hello", "(", 0, "")), output, errors) == 65
+    assert output.getvalue() == ""
+    assert errors.getvalue()
+
+
+@pytest.mark.parametrize("payload", ["[]", '["unknown", "text", "pattern", 0, ""]'])
+def test_regex_worker_rejects_invalid_requests_instead_of_returning_false(payload: str) -> None:
+    with pytest.raises(ValueError, match=r"invalid regex request|unknown regex operation"):
+        regex_worker(payload, StringIO(), StringIO())