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/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..adaf50614f5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/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 [] 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..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,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())