Merge a6493b5854 into 11779b260e

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/bounded_execution.py

@@ -6,7 +6,7 @@ variable, and the sandbox compiler routes every ``while`` test, ``for`` iteratio
 :func:`budget_ok`, which raises ``ExecutionInterrupted`` once that deadline has passed, whatever the code
 catches around the loop body. As a backstop, a worker thread still running at the deadline has
 ``ExecutionInterrupted`` injected with ``PyThreadState_SetAsyncExc`` and a task still running is cancelled
-repeatedly. A long-running C call (a catastrophic regex, for one) only sees any of this once it returns, so
+repeatedly. A long-running C call without its own timeout only sees any of this once it returns, so
 the caller still gets its timeout on schedule while the worker keeps burning CPU until that call ends.
 """
 
@@ -56,6 +56,16 @@ def _past_deadline() -> bool:
     return deadline is not None and time.monotonic() > deadline
 
 
+def remaining_execution_time() -> float | None:
+    deadline: Final = _deadline.get()
+    if deadline is None:
+        return None
+    remaining: Final = deadline - time.monotonic()
+    if remaining <= 0:
+        raise ExecutionInterrupted
+    return remaining
+
+
 def budget_ok() -> bool:
     """Bound to ``_budget_ok_`` in the sandbox, where every ``while`` test starts with a call to it."""
     if _past_deadline():

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

@@ -8,12 +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
-from pydantic import JsonValue
+from pydantic import JsonValue, TypeAdapter
 from typing_extensions import ReadOnly, TypedDict
 
 import litellm
@@ -22,6 +25,8 @@ from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, vali
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
 from litellm.types.llms.custom_http import httpxSpecialProvider
 
+from .bounded_execution import ExecutionInterrupted, remaining_execution_time
+
 # =============================================================================
 # Result Types - Used by Starlark code to return guardrail decisions
 # =============================================================================
@@ -110,6 +115,46 @@ def modify(
 # =============================================================================
 
 
+_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_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:
     """
     Check if a regex pattern matches anywhere in the text.
@@ -123,7 +168,9 @@ def regex_match(text: str, pattern: str, flags: int = 0) -> bool:
         True if pattern matches, False otherwise
     """
     try:
-        return bool(re.search(pattern, text, flags))
+        return _regex_result("search", text, pattern, flags, _REGEX_MATCH)
+    except subprocess.TimeoutExpired:
+        raise ExecutionInterrupted from None
     except re.error as e:
         verbose_proxy_logger.warning("Starlark regex_match error: %s", e)
         return False
@@ -142,7 +189,9 @@ def regex_match_all(text: str, pattern: str, flags: int = 0) -> bool:
         True if pattern matches entire text, False otherwise
     """
     try:
-        return bool(re.fullmatch(pattern, text, flags))
+        return _regex_result("fullmatch", text, pattern, flags, _REGEX_MATCH)
+    except subprocess.TimeoutExpired:
+        raise ExecutionInterrupted from None
     except re.error as e:
         verbose_proxy_logger.warning("Starlark regex_match_all error: %s", e)
         return False
@@ -162,13 +211,15 @@ def regex_replace(text: str, pattern: str, replacement: str, flags: int = 0) ->
         The text with replacements applied
     """
     try:
-        return re.sub(pattern, replacement, text, flags=flags)
+        return _regex_result("sub", text, pattern, flags, _REGEX_TEXT, replacement)
+    except subprocess.TimeoutExpired:
+        raise ExecutionInterrupted from None
     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.
 
@@ -181,7 +232,9 @@ def regex_find_all(text: str, pattern: str, flags: int = 0) -> list[str]:
         List of all matches
     """
     try:
-        return re.findall(pattern, text, flags)
+        return _regex_result("findall", text, pattern, flags, _REGEX_MATCHES)
+    except subprocess.TimeoutExpired:
+        raise ExecutionInterrupted from None
     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,9 +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,
@@ -11,6 +16,16 @@ from litellm.proxy.guardrails.guardrail_hooks.custom_code.bounded_execution impo
     call_off_loop_with_timeout,
     call_with_timeout,
 )
+from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
+    CustomCodeExecutionError,
+    CustomCodeGuardrail,
+)
+from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import (
+    regex_find_all,
+    regex_match,
+    regex_match_all,
+    regex_replace,
+)
 
 
 def _worker_threads() -> list[str]:
@@ -186,3 +201,127 @@ async def test_await_with_timeout_cancels_the_code_when_the_caller_is_cancelled(
         await waiting
 
     assert await asyncio.wait_for(interrupted.wait(), timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_regex_near_match_does_not_hold_the_execution_timeout() -> None:
+    guardrail: Final = CustomCodeGuardrail(
+        custom_code=(
+            "def apply_guardrail(inputs, request_data, input_type):\n"
+            '    regex_match(inputs["texts"][0], "(a+)+$")\n'
+            "    return allow()\n"
+        ),
+        guardrail_name="regex-near-match",
+        execution_timeout=0.1,
+    )
+    started: Final = time.monotonic()
+    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"
+        )
+    await asyncio.to_thread(_join_worker_threads)
+    assert time.monotonic() - started < 0.8
+    assert _worker_threads() == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize(
+    "operation",
+    [
+        'regex_match(inputs["texts"][0], "(a|aa)+$")',
+        'regex_match_all(inputs["texts"][0], "(a|aa)+$")',
+        'regex_replace(inputs["texts"][0], "(a|aa)+$", "x")',
+        'regex_find_all(inputs["texts"][0], "(a|aa)+$")',
+    ],
+)
+async def test_regex_deadline_cannot_be_swallowed_as_a_failed_match(operation: str, asynchronous: bool) -> None:
+    guardrail: Final = CustomCodeGuardrail(
+        custom_code=(
+            f"{'async ' if asynchronous else ''}def apply_guardrail(inputs, request_data, input_type):\n"
+            "    try:\n"
+            f"        {operation}\n"
+            "    except Exception:\n"
+            "        return allow()\n"
+            "    return allow()\n"
+        ),
+        guardrail_name="regex-deadline",
+        execution_timeout=0.1,
+    )
+    started: Final = time.monotonic()
+    with pytest.raises(CustomCodeExecutionError, match=r"exceeded its 0\.1s execution timeout"):
+        await guardrail.apply_guardrail(
+            inputs={"texts": ["a" * 32 + "!"]}, request_data={"model": "m"}, input_type="request"
+        )
+    await asyncio.to_thread(_join_worker_threads)
+    assert time.monotonic() - started < 0.8
+    assert _worker_threads() == []
+
+
+@pytest.mark.parametrize(
+    ("text", "pattern", "flags"),
+    [
+        ("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),
+        ("hello", "h e l l o", re.VERBOSE),
+        ("hello", "(h)(ello)", 0),
+        ("hello", "(", 0),
+    ],
+)
+def test_regex_primitives_preserve_standard_matching_results(text: str, pattern: str, flags: int) -> None:
+    try:
+        expected_match: Final = bool(re.search(pattern, text, flags))
+        expected_fullmatch: Final = bool(re.fullmatch(pattern, text, flags))
+        expected_replacement: Final = re.sub(pattern, "x", text, flags=flags)
+        expected_matches: Final = re.findall(pattern, text, flags)
+    except re.error:
+        assert (regex_match(text, pattern, flags), regex_match_all(text, pattern, flags)) == (False, False)
+        assert (regex_replace(text, pattern, "x", flags), regex_find_all(text, pattern, flags)) == (text, [])
+        return
+    assert (regex_match(text, pattern, flags), regex_match_all(text, pattern, flags)) == (
+        expected_match,
+        expected_fullmatch,
+    )
+    assert (regex_replace(text, pattern, "x", flags), regex_find_all(text, pattern, flags)) == (
+        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())