fix(guardrails): bound custom regex calls by the execution deadline

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

@@ -13,6 +13,7 @@ from typing import Final, Literal
 from urllib.parse import urlparse
 
 import httpx
+import regex
 from pydantic import JsonValue
 from typing_extensions import ReadOnly, TypedDict
 
@@ -22,6 +23,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 +113,26 @@ 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),
+)
+
+
+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_match(text: str, pattern: str, flags: int = 0) -> bool:
     """
     Check if a regex pattern matches anywhere in the text.
@@ -123,8 +146,14 @@ 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))
-    except re.error as e:
+        return bool(
+            regex.search(
+                pattern, text, _regex_flags(pattern, flags), timeout=remaining_execution_time(), concurrent=True
+            )
+        )
+    except TimeoutError:
+        raise ExecutionInterrupted from None
+    except (re.error, regex.error) as e:
         verbose_proxy_logger.warning("Starlark regex_match error: %s", e)
         return False
 
@@ -142,8 +171,14 @@ 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))
-    except re.error as e:
+        return bool(
+            regex.fullmatch(
+                pattern, text, _regex_flags(pattern, flags), timeout=remaining_execution_time(), concurrent=True
+            )
+        )
+    except TimeoutError:
+        raise ExecutionInterrupted from None
+    except (re.error, regex.error) as e:
         verbose_proxy_logger.warning("Starlark regex_match_all error: %s", e)
         return False
 
@@ -162,8 +197,17 @@ 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)
-    except re.error as e:
+        return regex.sub(
+            pattern,
+            replacement,
+            text,
+            flags=_regex_flags(pattern, flags),
+            timeout=remaining_execution_time(),
+            concurrent=True,
+        )
+    except TimeoutError:
+        raise ExecutionInterrupted from None
+    except (re.error, regex.error) as e:
         verbose_proxy_logger.warning("Starlark regex_replace error: %s", e)
         return text
 
@@ -181,8 +225,12 @@ def regex_find_all(text: str, pattern: str, flags: int = 0) -> list[str]:
         List of all matches
     """
     try:
-        return re.findall(pattern, text, flags)
-    except re.error as e:
+        return regex.findall(
+            pattern, text, _regex_flags(pattern, flags), timeout=remaining_execution_time(), concurrent=True
+        )
+    except TimeoutError:
+        raise ExecutionInterrupted from None
+    except (re.error, regex.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,6 +1,8 @@
 import asyncio
+import re
 import threading
 import time
+from typing import Final
 
 import pytest
 
@@ -11,6 +13,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 +198,94 @@ 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()
+    try:
+        result: Final = 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() == []
+
+
+@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),
+        ("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,
+    )