From 0eca7d7c5f855f61c9ae990da4273c39e8bd0a66 Mon Sep 17 00:00:00 2001 From: Muhammad Osama Amin Date: Tue, 20 Oct 2026 01:14:43 +0000 Subject: [PATCH] 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, + )