From 0eca7d7c5f855f61c9ae990da4273c39e8bd0a66 Mon Sep 17 00:00:00 2001 From: Muhammad Osama Amin Date: Tue, 20 Oct 2026 01:14:43 +0000 Subject: [PATCH 1/2] fix(guardrails): bound custom regex calls by the execution deadline --- .../custom_code/bounded_execution.py | 12 +- .../guardrail_hooks/custom_code/primitives.py | 64 +++++++++-- .../test_custom_code_bounded_execution.py | 103 ++++++++++++++++++ 3 files changed, 170 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/bounded_execution.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/bounded_execution.py index d926b40da6a..4f3b6160773 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/bounded_execution.py +++ b/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(): diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py index 55f7abd40a8..64a37ae35f8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/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 [] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py index 9114d811c57..9ec681a944a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py +++ b/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, + ) From a6493b5854191c96165408cc686080abc7b9b87d Mon Sep 17 00:00:00 2001 From: Muhammad Osama Amin Date: Thu, 22 Oct 2026 01:06:34 +0000 Subject: [PATCH 2/2] fix(guardrails): preserve stdlib regex semantics in isolated workers --- .../custom_code/_regex_worker.py | 42 +++++++ .../guardrail_hooks/custom_code/primitives.py | 103 +++++++++--------- .../test_custom_code_bounded_execution.py | 46 +++++++- 3 files changed, 137 insertions(+), 54 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/custom_code/_regex_worker.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/_regex_worker.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/_regex_worker.py new file mode 100644 index 00000000000..03c67a4f9cf --- /dev/null +++ b/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)) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py index 64a37ae35f8..adaf50614f5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/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 [] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py index 9ec681a944a..595ad6d84f6 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py +++ b/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())