fix(guardrails): replace custom_code sandbox with RestrictedPython
This commit is contained in:
parent
72a461ba4a
commit
0a1b4427a6
@ -18,12 +18,9 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.code_validator import (
|
||||
CustomCodeValidationError,
|
||||
validate_custom_code,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import (
|
||||
get_custom_code_primitives,
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import (
|
||||
build_sandbox_globals,
|
||||
compile_sandboxed,
|
||||
)
|
||||
from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry
|
||||
from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router
|
||||
@ -575,7 +572,9 @@ class GuardrailSubmissionItem(BaseModel):
|
||||
guardrail_name: str
|
||||
status: str # pending_review | active | rejected
|
||||
team_id: Optional[str] = None
|
||||
team_guardrail: bool = False # True when submitted via team (team_id set); use to distinguish team vs regular guardrails
|
||||
team_guardrail: bool = (
|
||||
False # True when submitted via team (team_id set); use to distinguish team vs regular guardrails
|
||||
)
|
||||
litellm_params: Optional[Dict[str, Any]] = None
|
||||
guardrail_info: Optional[Dict[str, Any]] = None
|
||||
submitted_by_user_id: Optional[str] = None
|
||||
@ -682,9 +681,9 @@ async def register_guardrail(
|
||||
guardrail_info = dict(request.guardrail_info or {})
|
||||
guardrail_info["submitted_by_user_id"] = user_api_key_dict.user_id
|
||||
guardrail_info["submitted_by_email"] = user_api_key_dict.user_email
|
||||
guardrail_info[
|
||||
"team_guardrail"
|
||||
] = True # Mark as team submission for filtering/display
|
||||
guardrail_info["team_guardrail"] = (
|
||||
True # Mark as team submission for filtering/display
|
||||
)
|
||||
guardrail_info_str = safe_dumps(guardrail_info)
|
||||
|
||||
try:
|
||||
@ -1879,9 +1878,9 @@ async def get_provider_specific_params():
|
||||
lakera_v2_fields = _get_fields_from_model(LakeraV2GuardrailConfigModel)
|
||||
tool_permission_fields = _get_fields_from_model(ToolPermissionGuardrailConfigModel)
|
||||
|
||||
tool_permission_fields[
|
||||
"ui_friendly_name"
|
||||
] = ToolPermissionGuardrailConfigModel.ui_friendly_name()
|
||||
tool_permission_fields["ui_friendly_name"] = (
|
||||
ToolPermissionGuardrailConfigModel.ui_friendly_name()
|
||||
)
|
||||
|
||||
# Return the provider-specific parameters
|
||||
provider_params = {
|
||||
@ -2029,25 +2028,11 @@ async def test_custom_code_guardrail(
|
||||
EXECUTION_TIMEOUT_SECONDS = 5
|
||||
|
||||
try:
|
||||
# Step 0: Security validation - check for forbidden patterns
|
||||
exec_globals = build_sandbox_globals()
|
||||
|
||||
try:
|
||||
validate_custom_code(request.custom_code)
|
||||
except CustomCodeValidationError as e:
|
||||
return TestCustomCodeGuardrailResponse(
|
||||
success=False,
|
||||
error=str(e),
|
||||
error_type="compilation",
|
||||
)
|
||||
|
||||
# Step 1: Compile the custom code with restricted environment
|
||||
exec_globals = get_custom_code_primitives().copy()
|
||||
|
||||
# Remove access to builtins to prevent escape
|
||||
exec_globals["__builtins__"] = {}
|
||||
|
||||
try:
|
||||
exec(compile(request.custom_code, "<guardrail>", "exec"), exec_globals)
|
||||
compiled = compile_sandboxed(request.custom_code)
|
||||
exec(compiled, exec_globals) # noqa: S102
|
||||
except SyntaxError as e:
|
||||
return TestCustomCodeGuardrailResponse(
|
||||
success=False,
|
||||
@ -2154,10 +2139,10 @@ async def apply_guardrail(
|
||||
from litellm.proxy.utils import handle_exception_on_proxy
|
||||
|
||||
try:
|
||||
active_guardrail: Optional[
|
||||
CustomGuardrail
|
||||
] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(
|
||||
guardrail_name=request.guardrail_name
|
||||
active_guardrail: Optional[CustomGuardrail] = (
|
||||
GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(
|
||||
guardrail_name=request.guardrail_name
|
||||
)
|
||||
)
|
||||
if active_guardrail is None:
|
||||
raise HTTPException(
|
||||
|
||||
@ -1,63 +0,0 @@
|
||||
import re
|
||||
from typing import List, Tuple
|
||||
|
||||
# Security validation patterns
|
||||
FORBIDDEN_PATTERNS: List[Tuple[str, str]] = [
|
||||
# Import statements
|
||||
(r"\bimport\s+", "import statements are not allowed"),
|
||||
(r"\bfrom\s+\w+\s+import\b", "from...import statements are not allowed"),
|
||||
(r"__import__\s*\(", "__import__() is not allowed"),
|
||||
# Dangerous builtins
|
||||
(r"\bexec\s*\(", "exec() is not allowed"),
|
||||
(r"\beval\s*\(", "eval() is not allowed"),
|
||||
(r"\bcompile\s*\(", "compile() is not allowed"),
|
||||
(r"\bopen\s*\(", "open() is not allowed"),
|
||||
(r"\bgetattr\s*\(", "getattr() is not allowed"),
|
||||
(r"\bsetattr\s*\(", "setattr() is not allowed"),
|
||||
(r"\bdelattr\s*\(", "delattr() is not allowed"),
|
||||
(r"\bglobals\s*\(", "globals() is not allowed"),
|
||||
(r"\blocals\s*\(", "locals() is not allowed"),
|
||||
(r"\bvars\s*\(", "vars() is not allowed"),
|
||||
(r"\bdir\s*\(", "dir() is not allowed"),
|
||||
(r"\bbreakpoint\s*\(", "breakpoint() is not allowed"),
|
||||
(r"\binput\s*\(", "input() is not allowed"),
|
||||
# Dangerous dunder access
|
||||
(r"__builtins__", "__builtins__ access is not allowed"),
|
||||
(r"__globals__", "__globals__ access is not allowed"),
|
||||
(r"__code__", "__code__ access is not allowed"),
|
||||
(r"__subclasses__", "__subclasses__ access is not allowed"),
|
||||
(r"__bases__", "__bases__ access is not allowed"),
|
||||
(r"__mro__", "__mro__ access is not allowed"),
|
||||
(r"__class__", "__class__ access is not allowed"),
|
||||
(r"__dict__", "__dict__ access is not allowed"),
|
||||
(r"__getattribute__", "__getattribute__ access is not allowed"),
|
||||
(r"__reduce__", "__reduce__ access is not allowed"),
|
||||
(r"__reduce_ex__", "__reduce_ex__ access is not allowed"),
|
||||
# OS/system access
|
||||
(r"\bos\.", "os module access is not allowed"),
|
||||
(r"\bsys\.", "sys module access is not allowed"),
|
||||
(r"\bsubprocess\.", "subprocess module access is not allowed"),
|
||||
(r"\bshutil\.", "shutil module access is not allowed"),
|
||||
(r"\bctypes\.", "ctypes module access is not allowed"),
|
||||
(r"\bsocket\.", "socket module access is not allowed"),
|
||||
(r"\bpickle\.", "pickle module access is not allowed"),
|
||||
]
|
||||
|
||||
|
||||
class CustomCodeValidationError(Exception):
|
||||
"""Raised when custom code fails security validation."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def validate_custom_code(code: str) -> None:
|
||||
"""
|
||||
Validate custom code against forbidden patterns.
|
||||
|
||||
Raises CustomCodeValidationError if any forbidden pattern is found.
|
||||
"""
|
||||
if not code:
|
||||
return
|
||||
for pattern, error_msg in FORBIDDEN_PATTERNS:
|
||||
if re.search(pattern, code):
|
||||
raise CustomCodeValidationError(f"Security violation: {error_msg}")
|
||||
@ -49,8 +49,7 @@ from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
from .code_validator import CustomCodeValidationError, validate_custom_code
|
||||
from .primitives import get_custom_code_primitives
|
||||
from .sandbox import build_sandbox_globals, compile_sandboxed
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
@ -146,17 +145,10 @@ class CustomCodeGuardrail(CustomGuardrail):
|
||||
|
||||
def _do_compile(self) -> None:
|
||||
"""Internal compilation method without lock. Expected to run inside _compile_lock."""
|
||||
# Create a restricted execution environment
|
||||
# Only include our safe primitives
|
||||
exec_globals = get_custom_code_primitives().copy()
|
||||
exec_globals = build_sandbox_globals()
|
||||
compiled = compile_sandboxed(self.custom_code)
|
||||
exec(compiled, exec_globals) # noqa: S102
|
||||
|
||||
# CRITICAL: Restrict __builtins__ to prevent sandbox escape
|
||||
exec_globals["__builtins__"] = {}
|
||||
|
||||
# Execute the user code in the restricted environment
|
||||
exec(compile(self.custom_code, "<guardrail>", "exec"), exec_globals)
|
||||
|
||||
# Extract the apply_guardrail function
|
||||
if "apply_guardrail" not in exec_globals:
|
||||
raise CustomCodeCompilationError(
|
||||
"Custom code must define an 'apply_guardrail' function. "
|
||||
@ -182,13 +174,6 @@ class CustomCodeGuardrail(CustomGuardrail):
|
||||
return
|
||||
|
||||
try:
|
||||
# Step 1: Security validation — forbidden pattern check
|
||||
try:
|
||||
validate_custom_code(self.custom_code)
|
||||
except CustomCodeValidationError as e:
|
||||
raise CustomCodeCompilationError(str(e)) from e
|
||||
|
||||
# Step 2: Compile logic
|
||||
self._do_compile()
|
||||
verbose_proxy_logger.debug(
|
||||
f"Custom code guardrail '{self.guardrail_name}' compiled successfully"
|
||||
@ -405,12 +390,6 @@ class CustomCodeGuardrail(CustomGuardrail):
|
||||
Raises:
|
||||
CustomCodeCompilationError: If the new code fails to compile
|
||||
"""
|
||||
# Validate BEFORE acquiring lock / resetting state
|
||||
try:
|
||||
validate_custom_code(new_code)
|
||||
except CustomCodeValidationError as e:
|
||||
raise CustomCodeCompilationError(str(e)) from e
|
||||
|
||||
with self._compile_lock:
|
||||
# Reset state
|
||||
old_function = self._compiled_function
|
||||
|
||||
100
litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py
Normal file
100
litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py
Normal file
@ -0,0 +1,100 @@
|
||||
"""
|
||||
RestrictedPython-based sandbox for custom code guardrails.
|
||||
|
||||
User-supplied guardrail source is compiled with ``compile_restricted`` and
|
||||
executed with curated globals. All attribute access, subscripting, iteration,
|
||||
and assignment is mediated by RestrictedPython guards, which block dunder
|
||||
access (``__globals__``, ``__code__``, ``__class__``, ``__setattr__``, etc.)
|
||||
and reject dangerous AST constructs (``import``, ``exec``, ``eval``,
|
||||
``compile``, class definitions, etc.) at compile time.
|
||||
|
||||
The default ``RestrictingNodeTransformer`` denies ``async def``/``await``,
|
||||
which breaks the documented async guardrail pattern (``await http_get(...)``).
|
||||
We subclass it to permit those specific nodes, while keeping every other
|
||||
restriction intact.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
from RestrictedPython import (
|
||||
RestrictingNodeTransformer,
|
||||
compile_restricted,
|
||||
limited_builtins,
|
||||
safe_builtins,
|
||||
utility_builtins,
|
||||
)
|
||||
from RestrictedPython.Eval import default_guarded_getitem, default_guarded_getiter
|
||||
from RestrictedPython.Guards import (
|
||||
full_write_guard,
|
||||
guarded_iter_unpack_sequence,
|
||||
safer_getattr,
|
||||
)
|
||||
|
||||
from .primitives import get_custom_code_primitives
|
||||
|
||||
|
||||
class AsyncAwareTransformer(RestrictingNodeTransformer):
|
||||
"""Extend the default transformer to allow ``async def`` and ``await``.
|
||||
|
||||
The base class rejects every async AST node outright. ``AsyncFunctionDef``
|
||||
has the same ``_fields`` as ``FunctionDef`` and the same security
|
||||
semantics, so we delegate to ``visit_FunctionDef`` — name check, argument
|
||||
check, print-scope wrapping, and any future additions to that method are
|
||||
inherited automatically. ``AsyncFor``/``AsyncWith``/``Await`` delegate to
|
||||
``node_contents_visit`` so their children still get transformed.
|
||||
"""
|
||||
|
||||
def visit_AsyncFunctionDef(self, node: Any) -> Any:
|
||||
return self.visit_FunctionDef(node)
|
||||
|
||||
def visit_AsyncFor(self, node: Any) -> Any:
|
||||
return self.node_contents_visit(node)
|
||||
|
||||
def visit_AsyncWith(self, node: Any) -> Any:
|
||||
return self.node_contents_visit(node)
|
||||
|
||||
def visit_Await(self, node: Any) -> Any:
|
||||
return self.node_contents_visit(node)
|
||||
|
||||
|
||||
def _build_sandbox_builtins() -> Dict[str, Any]:
|
||||
# ``limited_builtins`` overrides ``list``/``tuple``/``range`` from
|
||||
# ``safe_builtins`` with bounds-checking variants (e.g. ``limited_range``
|
||||
# rejects ``range(10**18)``). ``utility_builtins`` adds ``set``,
|
||||
# ``frozenset``, ``math``, ``random``, and a filtered ``string`` delegator.
|
||||
return {
|
||||
**safe_builtins,
|
||||
**limited_builtins,
|
||||
**utility_builtins,
|
||||
}
|
||||
|
||||
|
||||
def build_sandbox_globals() -> Dict[str, Any]:
|
||||
"""Assemble the globals dict for executing guardrail code.
|
||||
|
||||
Includes the LiteLLM-provided primitives (``regex_match``, ``http_get``,
|
||||
``allow``/``block``/``modify``, etc.) plus the RestrictedPython guards
|
||||
that the compiled bytecode expects to find by name.
|
||||
"""
|
||||
sandbox: Dict[str, Any] = get_custom_code_primitives().copy()
|
||||
sandbox["__builtins__"] = _build_sandbox_builtins()
|
||||
sandbox["_getattr_"] = safer_getattr
|
||||
sandbox["_getitem_"] = default_guarded_getitem
|
||||
sandbox["_getiter_"] = default_guarded_getiter
|
||||
sandbox["_iter_unpack_sequence_"] = guarded_iter_unpack_sequence
|
||||
sandbox["_write_"] = full_write_guard
|
||||
return sandbox
|
||||
|
||||
|
||||
def compile_sandboxed(source: str, filename: str = "<guardrail>") -> Any:
|
||||
"""Compile guardrail source with RestrictedPython's AST transformer.
|
||||
|
||||
Raises ``SyntaxError`` on either a Python syntax error or a restricted
|
||||
construct (import, exec, dunder name, etc.).
|
||||
"""
|
||||
return compile_restricted(
|
||||
source=source,
|
||||
filename=filename,
|
||||
mode="exec",
|
||||
policy=AsyncAwareTransformer,
|
||||
)
|
||||
@ -54,6 +54,7 @@ proxy = [
|
||||
"mcp==1.26.0; python_version >= '3.10'",
|
||||
"litellm-proxy-extras==0.4.65",
|
||||
"litellm-enterprise==0.1.37",
|
||||
"RestrictedPython==8.1",
|
||||
"rich==13.9.4",
|
||||
"polars==1.38.1; python_version >= '3.10'",
|
||||
"soundfile==0.12.1",
|
||||
|
||||
@ -1,85 +1,148 @@
|
||||
import pytest
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.code_validator import (
|
||||
validate_custom_code,
|
||||
CustomCodeValidationError,
|
||||
)
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
|
||||
CustomCodeCompilationError,
|
||||
CustomCodeGuardrail,
|
||||
)
|
||||
|
||||
# Phase 4.1: Test forbidden pattern validation
|
||||
|
||||
# str.mro() + generator gi_code + code.replace(co_names=...) + __setattr__
|
||||
# to swap a function's bytecode and read http_get's real builtins dict.
|
||||
BYTECODE_REWRITE_PAYLOAD = (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
" obj = str.mro()[1]\n"
|
||||
" def g(fn):\n"
|
||||
" yield fn.placeholder\n"
|
||||
" c = g(None).gi_code\n"
|
||||
' gn = "_"+"_gl"+"ob"+"als"+"_"+"_"\n'
|
||||
' cn = "_"+"_co"+"de_"+"_"\n'
|
||||
" obj.__setattr__(g, cn, c.replace(co_names=(gn,)))\n"
|
||||
" for v in g(http_get):\n"
|
||||
" gd = v\n"
|
||||
" break\n"
|
||||
' bn = "_"+"_bu"+"ilt"+"ins"+"_"+"_"\n'
|
||||
' imp = gd[bn]["_"+"_im"+"po"+"rt_"+"_"]\n'
|
||||
' return {"rce": imp("os").popen("id").read()}\n'
|
||||
)
|
||||
|
||||
|
||||
def test_validate_custom_code_import_os():
|
||||
code = "import os\ndef apply_guardrail(inputs, req, ty):\n return allow()"
|
||||
with pytest.raises(CustomCodeValidationError, match="import statements are not"):
|
||||
validate_custom_code(code)
|
||||
def _compile(code: str) -> CustomCodeGuardrail:
|
||||
return CustomCodeGuardrail(custom_code=code, guardrail_name="t")
|
||||
|
||||
|
||||
def test_validate_custom_code_from_subprocess():
|
||||
def test_bytecode_rewrite_rejected_at_compile():
|
||||
with pytest.raises(CustomCodeCompilationError):
|
||||
_compile(BYTECODE_REWRITE_PAYLOAD)
|
||||
|
||||
|
||||
# Call the async http_get primitive without awaiting, then pull f_builtins off
|
||||
# the returned coroutine's cr_frame. INSPECT_ATTRIBUTES covers cr_frame and
|
||||
# f_builtins so this is rejected at compile time.
|
||||
CR_FRAME_PAYLOAD = (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
' co = http_get("http://x")\n'
|
||||
" b = co.cr_frame.f_builtins\n"
|
||||
" co.close()\n"
|
||||
' imp = b["_" + "_imp" + "ort_" + "_"]\n'
|
||||
' return block(imp("os").popen("id").read())\n'
|
||||
)
|
||||
|
||||
|
||||
def test_cr_frame_rejected_at_compile():
|
||||
with pytest.raises(CustomCodeCompilationError):
|
||||
_compile(CR_FRAME_PAYLOAD)
|
||||
|
||||
|
||||
# NFKC homoglyph: U+FF47 'g' normalizes to 'g' at parse time, so "__globals__"
|
||||
# arrives at the AST as "__globals__" and trips the underscore-prefix rule.
|
||||
NFKC_PAYLOAD = (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
' b_key = "buil" + "tins"\n'
|
||||
' i_key = "im" + "port"\n'
|
||||
" b = allow.__\uff47lobals__[b_key]\n"
|
||||
" import_fn = b[i_key]\n"
|
||||
' o = import_fn("o" + "s")\n'
|
||||
' return block(o.popen("id").read())\n'
|
||||
)
|
||||
|
||||
|
||||
def test_nfkc_homoglyph_rejected_at_compile():
|
||||
with pytest.raises(CustomCodeCompilationError):
|
||||
_compile(NFKC_PAYLOAD)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"snippet",
|
||||
[
|
||||
# Literal dunder attribute access.
|
||||
"def apply_guardrail(i, r, t):\n return str.__class__\n",
|
||||
"def apply_guardrail(i, r, t):\n"
|
||||
" return ().__class__.__bases__[0].__subclasses__()\n",
|
||||
# gi_code — on the transformer's restricted-names list.
|
||||
"def apply_guardrail(i, r, t):\n"
|
||||
" def g():\n yield 1\n"
|
||||
" return g().gi_code\n",
|
||||
# Import forms.
|
||||
"import os\ndef apply_guardrail(i, r, t):\n return allow()\n",
|
||||
"from subprocess import call\n"
|
||||
"def apply_guardrail(i, r, t):\n return allow()\n",
|
||||
# __import__ is rejected as an underscore-prefixed name.
|
||||
"def apply_guardrail(i, r, t):\n" ' return __import__("os")\n',
|
||||
],
|
||||
)
|
||||
def test_compile_time_rejections(snippet: str):
|
||||
with pytest.raises(CustomCodeCompilationError):
|
||||
_compile(snippet)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"snippet",
|
||||
[
|
||||
# getattr is not in the sandbox builtins — NameError at call time.
|
||||
"def apply_guardrail(i, r, t):\n"
|
||||
' return getattr(str, "_"+"_class_"+"_")\n',
|
||||
# setattr is guarded_setattr + full_write_guard — setting any attribute
|
||||
# on a user-defined object raises TypeError, whether the name is a
|
||||
# dunder or not.
|
||||
"def apply_guardrail(i, r, t):\n"
|
||||
" def f():\n pass\n"
|
||||
' name = "_" + "_bad_" + "_"\n'
|
||||
" setattr(f, name, None)\n"
|
||||
" return allow()\n",
|
||||
],
|
||||
)
|
||||
def test_runtime_rejections(snippet: str):
|
||||
guardrail = _compile(snippet)
|
||||
fn = guardrail._compiled_function
|
||||
assert fn is not None
|
||||
with pytest.raises((NameError, TypeError, AttributeError, SyntaxError)):
|
||||
fn({"texts": []}, {}, "request")
|
||||
|
||||
|
||||
def test_documented_ssn_example_compiles_and_runs():
|
||||
code = (
|
||||
"from subprocess import call\ndef apply_guardrail(i, r, t):\n return allow()"
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
' for text in inputs["texts"]:\n'
|
||||
' if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"):\n'
|
||||
' return block("SSN detected")\n'
|
||||
" return allow()\n"
|
||||
)
|
||||
with pytest.raises(
|
||||
CustomCodeValidationError, match="import statements are not allowed"
|
||||
):
|
||||
validate_custom_code(code)
|
||||
|
||||
|
||||
def test_validate_custom_code_exec():
|
||||
code = "def apply_guardrail(i, r, t):\n exec('print(1)')\n return allow()"
|
||||
with pytest.raises(CustomCodeValidationError, match=r"exec\(\) is not allowed"):
|
||||
validate_custom_code(code)
|
||||
|
||||
|
||||
def test_validate_custom_code_builtins():
|
||||
code = "def apply_guardrail(i, r, t):\n print(__builtins__)\n return allow()"
|
||||
with pytest.raises(
|
||||
CustomCodeValidationError, match="__builtins__ access is not allowed"
|
||||
):
|
||||
validate_custom_code(code)
|
||||
|
||||
|
||||
def test_validate_custom_code_subclasses():
|
||||
code = "def apply_guardrail(i, r, t):\n print(''.__class__.__mro__[1].__subclasses__())\n return allow()"
|
||||
with pytest.raises(
|
||||
CustomCodeValidationError, match="__subclasses__ access is not allowed"
|
||||
):
|
||||
validate_custom_code(code)
|
||||
|
||||
|
||||
def test_validate_custom_code_clean():
|
||||
code = (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n return allow()\n"
|
||||
)
|
||||
# Should not raise any exception
|
||||
validate_custom_code(code)
|
||||
|
||||
|
||||
# Phase 4.2: Test __builtins__ restriction in execution
|
||||
|
||||
|
||||
def test_custom_code_compile_valid():
|
||||
code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()"
|
||||
guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="test")
|
||||
# if it doesn't fail, we successfully compiled
|
||||
assert guardrail._compiled_function is not None
|
||||
|
||||
|
||||
def test_custom_code_override_builtins():
|
||||
# Verify that even if pattern validation is bypassed, __builtins__ = {} blocks dangerous builtins.
|
||||
# We test this by compiling safe code and verifying builtins are not accessible in the sandbox.
|
||||
code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()"
|
||||
guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="test")
|
||||
# The compiled function's globals should have empty __builtins__
|
||||
fn_globals = guardrail._compiled_function.__globals__
|
||||
assert fn_globals.get("__builtins__") == {}
|
||||
guardrail = _compile(code)
|
||||
fn = guardrail._compiled_function
|
||||
assert fn is not None
|
||||
assert fn({"texts": ["hello"]}, {}, "request") == {"action": "allow"}
|
||||
blocked = fn({"texts": ["my ssn 123-45-6789"]}, {}, "request")
|
||||
assert blocked["action"] == "block"
|
||||
assert blocked["reason"] == "SSN detected"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_code_guardrail_apply():
|
||||
code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()"
|
||||
guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="test")
|
||||
async def test_async_guardrail_compiles_and_runs():
|
||||
code = (
|
||||
"async def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
" return allow()\n"
|
||||
)
|
||||
guardrail = _compile(code)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
result = await guardrail.apply_guardrail(
|
||||
@ -90,5 +153,15 @@ async def test_custom_code_guardrail_apply():
|
||||
assert result["texts"][0] == "test"
|
||||
|
||||
|
||||
# The RBAC endpoint tests are harder to write right here, but the core security
|
||||
# validations are fully covered by the simple tests above.
|
||||
def test_typical_sync_guardrail_still_works():
|
||||
code = (
|
||||
"def apply_guardrail(inputs, request_data, input_type):\n"
|
||||
" return allow()\n"
|
||||
)
|
||||
guardrail = _compile(code)
|
||||
assert guardrail._compiled_function is not None
|
||||
|
||||
|
||||
def test_missing_apply_guardrail_raises():
|
||||
with pytest.raises(CustomCodeCompilationError, match="apply_guardrail"):
|
||||
_compile("x = 1\n")
|
||||
|
||||
15
uv.lock
generated
15
uv.lock
generated
@ -11,7 +11,7 @@ resolution-markers = [
|
||||
]
|
||||
|
||||
[options]
|
||||
exclude-newer = "2026-04-08T16:01:27.663665Z"
|
||||
exclude-newer = "2026-04-12T21:11:57.53582856Z"
|
||||
exclude-newer-span = "P3D"
|
||||
|
||||
[manifest]
|
||||
@ -3602,7 +3602,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "litellm"
|
||||
version = "1.83.6"
|
||||
version = "1.83.8"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@ -3662,6 +3662,7 @@ proxy = [
|
||||
{ name = "pyroscope-io", marker = "sys_platform != 'win32'" },
|
||||
{ name = "python-multipart" },
|
||||
{ name = "pyyaml" },
|
||||
{ name = "restrictedpython" },
|
||||
{ name = "rich" },
|
||||
{ name = "rq" },
|
||||
{ name = "soundfile" },
|
||||
@ -3841,6 +3842,7 @@ requires-dist = [
|
||||
{ name = "pyyaml", marker = "extra == 'proxy'", specifier = "==6.0.3" },
|
||||
{ name = "redisvl", marker = "python_full_version >= '3.9' and python_full_version < '3.14' and extra == 'extra-proxy'", specifier = "==0.4.1" },
|
||||
{ name = "resend", marker = "extra == 'extra-proxy'", specifier = "==2.23.0" },
|
||||
{ name = "restrictedpython", marker = "extra == 'proxy'", specifier = "==8.1" },
|
||||
{ name = "rich", marker = "extra == 'proxy'", specifier = "==13.9.4" },
|
||||
{ name = "rq", marker = "extra == 'proxy'", specifier = "==2.7.0" },
|
||||
{ name = "semantic-router", marker = "python_full_version >= '3.9' and python_full_version < '3.14' and extra == 'semantic-router'", specifier = "==0.1.12" },
|
||||
@ -7460,6 +7462,15 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/8e/67/afbb0978d5399bc9ea200f1d4489a23c9a1dad4eee6376242b8182389c79/respx-0.22.0-py2.py3-none-any.whl", hash = "sha256:631128d4c9aba15e56903fb5f66fb1eff412ce28dd387ca3a81339e52dbd3ad0", size = 25127, upload-time = "2024-12-19T22:33:57.837Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "restrictedpython"
|
||||
version = "8.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/5f/1c/aec08bcb4ab14a1521579fbe21ceff2a634bb1f737f11cf7f9c8bb96e680/restrictedpython-8.1.tar.gz", hash = "sha256:4a69304aceacf6bee74bdf153c728221d4e3109b39acbfe00b3494927080d898", size = 838331, upload-time = "2025-10-19T14:11:32.531Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/1a/c0/3848f4006f7e164ee20833ca984067e4b3fc99fe7f1dfa88b4927e681299/restrictedpython-8.1-py3-none-any.whl", hash = "sha256:4769449c6cdb10f2071649ba386902befff0eff2a8fd6217989fa7b16aeae926", size = 27651, upload-time = "2025-10-19T14:11:30.201Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rfc3339-validator"
|
||||
version = "0.1.4"
|
||||
|
||||
Loading…
Reference in New Issue
Block a user