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())