scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

fix Unauthenticated RCE and Sandbox Escape in Custom Code Guardrail

Browse source
This commit is contained in:
Harshit28j 2026-02-24 13:23:55 +05:30 • committed by Sameer Kankute
parent 893eeef1ad
commit d7adecd996
6 changed files with 339 additions and 114 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm/proxy/auth/user_api_key_auth.py
View file

@ -593,9 +593,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
user_id=user_id,
team_id=team_id,
team_alias=(
team_object.team_alias
if team_object is not None
else None
team_object.team_alias if team_object is not None else None
),
team_metadata=team_object.metadata
if team_object is not None
@ -709,12 +707,12 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if isinstance(api_key, str):
return UserAPIKeyAuth(
api_key=api_key,
user_role=LitellmUserRoles.PROXY_ADMIN,
user_role=LitellmUserRoles.INTERNAL_USER,
parent_otel_span=parent_otel_span,
)
else:
return UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_role=LitellmUserRoles.INTERNAL_USER,
parent_otel_span=parent_otel_span,
)
elif api_key is None: # only require api key if master key is set
@ -846,7 +844,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
)
valid_token.parent_otel_span = parent_otel_span
if _end_user_object is not None:
valid_token.end_user_object_permission = _end_user_object.object_permission
valid_token.end_user_object_permission = (
_end_user_object.object_permission
)
return valid_token
@ -954,7 +954,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if isinstance(
api_key, str
): # if generated token, make sure it starts with sk-.
_masked_key = "{}****{}".format(api_key[:4], api_key[-4:]) if len(api_key) > 8 else "****"
_masked_key = (
"{}****{}".format(api_key[:4], api_key[-4:])
if len(api_key) > 8
else "****"
)
assert api_key.startswith(
"sk-"
), "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format(
@ -1304,9 +1308,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if _end_user_object is not None:
valid_token_dict.update(end_user_params)
valid_token_dict["end_user_object_permission"] = (
_end_user_object.object_permission
)
valid_token_dict[
"end_user_object_permission"
] = _end_user_object.object_permission
# check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions
# sso/login, ui/login, /key functions and /user functions

134
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -245,7 +245,10 @@ class CreateGuardrailRequest(BaseModel):
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def create_guardrail(request: CreateGuardrailRequest):
async def create_guardrail(
request: CreateGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Create a new guardrail
@ -295,6 +298,13 @@ async def create_guardrail(request: CreateGuardrailRequest):
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy._types import LitellmUserRoles
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -334,7 +344,11 @@ class UpdateGuardrailRequest(BaseModel):
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest):
async def update_guardrail(
guardrail_id: str,
request: UpdateGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Update an existing guardrail
@ -384,6 +398,13 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest):
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy._types import LitellmUserRoles
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -431,7 +452,10 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest):
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def delete_guardrail(guardrail_id: str):
async def delete_guardrail(
guardrail_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Delete a guardrail
@ -452,6 +476,13 @@ async def delete_guardrail(guardrail_id: str):
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy._types import LitellmUserRoles
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -497,7 +528,11 @@ async def delete_guardrail(guardrail_id: str):
tags=["Guardrails"],
dependencies=[Depends(user_api_key_auth)],
)
async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest):
async def patch_guardrail(
guardrail_id: str,
request: PatchGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Partially update an existing guardrail
@ -545,6 +580,13 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest):
"""
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
from litellm.proxy.proxy_server import prisma_client
from litellm.proxy._types import LitellmUserRoles
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to manage guardrails",
)
if prisma_client is None:
raise HTTPException(status_code=500, detail="Prisma client not initialized")
@ -1302,9 +1344,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 = {
@ -1367,7 +1409,10 @@ class TestCustomCodeGuardrailResponse(BaseModel):
dependencies=[Depends(user_api_key_auth)],
response_model=TestCustomCodeGuardrailResponse,
)
async def test_custom_code_guardrail(request: TestCustomCodeGuardrailRequest):
async def test_custom_code_guardrail(
request: TestCustomCodeGuardrailRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Test custom code guardrail logic without creating a guardrail.
@ -1441,62 +1486,35 @@ async def test_custom_code_guardrail(request: TestCustomCodeGuardrailRequest):
```
"""
import concurrent.futures
import re
from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import (
get_custom_code_primitives,
)
from litellm.proxy._types import LitellmUserRoles
# Security validation patterns
FORBIDDEN_PATTERNS = [
# 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"),
]
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Admin access required to test custom code guardrails",
)
EXECUTION_TIMEOUT_SECONDS = 5
try:
# Step 0: Security validation - check for forbidden patterns
code = request.custom_code
for pattern, error_msg in FORBIDDEN_PATTERNS:
if re.search(pattern, code):
return TestCustomCodeGuardrailResponse(
success=False,
error=f"Security violation: {error_msg}",
error_type="compilation",
)
from litellm.proxy.guardrails.guardrail_hooks.custom_code.code_validator import (
validate_custom_code,
CustomCodeValidationError,
)
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()
@ -1612,10 +1630,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(

59
litellm/proxy/guardrails/guardrail_hooks/custom_code/code_validator.py Normal file
View file

@ -0,0 +1,59 @@
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"),
]
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}")

60
litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py
View file

@ -41,18 +41,19 @@ from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast
from fastapi import HTTPException
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (CustomGuardrail,
log_guardrail_information)
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.base import \
GuardrailConfigModel
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
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import \
Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
class CustomCodeGuardrailError(Exception):
@ -154,10 +155,19 @@ 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
# Create a restricted execution environment
# Only include our safe primitives
exec_globals = get_custom_code_primitives().copy()
# 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)
@ -390,6 +400,12 @@ 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
@ -399,12 +415,42 @@ class CustomCodeGuardrail(CustomGuardrail):
try:
self.custom_code = new_code
self._compile_custom_code()
# Inline compilation instead of calling _compile_custom_code()
# to avoid deadlock (re-acquiring self._compile_lock)
exec_globals = get_custom_code_primitives().copy()
exec_globals["__builtins__"] = {}
exec(compile(self.custom_code, "<guardrail>", "exec"), exec_globals)
if "apply_guardrail" not in exec_globals:
raise CustomCodeCompilationError(
"Custom code must define an 'apply_guardrail' function. "
"Expected signature: apply_guardrail(inputs, request_data, input_type)"
)
apply_fn = exec_globals["apply_guardrail"]
if not callable(apply_fn):
raise CustomCodeCompilationError(
"'apply_guardrail' must be a callable function"
)
self._compiled_function = apply_fn
verbose_proxy_logger.info(
f"Custom code guardrail '{self.guardrail_name}': Code updated successfully"
)
except SyntaxError as e:
# Rollback on failure
self.custom_code = old_code
self._compiled_function = old_function
self._compile_error = f"Syntax error in custom code: {e}"
raise CustomCodeCompilationError(self._compile_error) from e
except CustomCodeCompilationError:
# Rollback on failure
self.custom_code = old_code
self._compiled_function = old_function
raise
except Exception as e:
# Rollback on failure
self.custom_code = old_code
self._compiled_function = old_function
self._compile_error = f"Failed to compile custom code: {e}"
raise CustomCodeCompilationError(self._compile_error) from e

82
litellm/proxy/proxy_server.py
View file

@ -1216,9 +1216,7 @@ try:
# Case 2: Runtime UI exists and is ready
if has_content and is_pre_restructured:
verbose_proxy_logger.info(
f"Using pre-restructured UI at {runtime_ui_path}"
)
verbose_proxy_logger.info(f"Using pre-restructured UI at {runtime_ui_path}")
ui_path = runtime_ui_path
# Case 3: Runtime UI exists but needs restructuring
@ -2994,6 +2992,10 @@ class ProxyConfig:
if master_key is not None and isinstance(master_key, str):
litellm_master_key_hash = hash_token(master_key)
else:
verbose_proxy_logger.critical(
"LITELLM_MASTER_KEY is not set! All unauthenticated requests will be treated as INTERNAL_USER. This is insecure for production."
)
### USER API KEY CACHE IN-MEMORY TTL ###
user_api_key_cache_ttl = general_settings.get(
"user_api_key_cache_ttl", None
@ -3796,6 +3798,7 @@ class ProxyConfig:
parsed = value
elif isinstance(value, str):
import json
try:
parsed = yaml.safe_load(value)
except (yaml.YAMLError, json.JSONDecodeError):
@ -4381,10 +4384,12 @@ class ProxyConfig:
if self._should_load_db_object(object_type="model_cost_map"):
await self._check_and_reload_model_cost_map(prisma_client=prisma_client)
if self._should_load_db_object(object_type="anthropic_beta_headers"):
await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client)
await self._check_and_reload_anthropic_beta_headers(
prisma_client=prisma_client
)
if self._should_load_db_object(object_type="sso_settings"):
await self._init_sso_settings_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="cache_settings"):
@ -4614,7 +4619,9 @@ class ProxyConfig:
f"Error in _check_and_reload_model_cost_map: {str(e)}"
)
async def _check_and_reload_anthropic_beta_headers(self, prisma_client: PrismaClient):
async def _check_and_reload_anthropic_beta_headers(
self, prisma_client: PrismaClient
):
"""
Check if anthropic beta headers config needs to be reloaded based on database configuration.
This function runs every 10 seconds as part of _init_non_llm_objects_in_db.
@ -4705,7 +4712,11 @@ class ProxyConfig:
)
# Count providers in config
provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description")
provider_count = sum(
1
for k in new_config.keys()
if k != "provider_aliases" and k != "description"
)
verbose_proxy_logger.info(
f"Anthropic beta headers config reloaded successfully. Providers: {provider_count}"
)
@ -5687,8 +5698,7 @@ class ProxyStartupEvent:
):
_db_val = _db_gs_record.param_value.get("store_model_in_db")
if _db_val is True or (
isinstance(_db_val, str)
and _db_val.lower() == "true"
isinstance(_db_val, str) and _db_val.lower() == "true"
):
store_model_in_db = True
verbose_proxy_logger.info(
@ -6155,6 +6165,7 @@ class ProxyStartupEvent:
"Pyroscope profiling will not run. Install with: pip install pyroscope-io"
)
#### API ENDPOINTS ####
@router.get(
"/v1/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"]
@ -10993,18 +11004,14 @@ async def get_favicon():
from fastapi.responses import Response
current_dir = os.path.dirname(os.path.abspath(__file__))
default_favicon = os.path.join(
current_dir, "_experimental", "out", "favicon.ico"
)
default_favicon = os.path.join(current_dir, "_experimental", "out", "favicon.ico")
favicon_url = os.getenv("LITELLM_FAVICON_URL", "")
if not favicon_url:
if os.path.exists(default_favicon):
return FileResponse(default_favicon, media_type="image/x-icon")
raise HTTPException(
status_code=404, detail="Default favicon not found"
)
raise HTTPException(status_code=404, detail="Default favicon not found")
if favicon_url.startswith(("http://", "https://")):
try:
@ -11019,9 +11026,7 @@ async def get_favicon():
)
response = await async_client.get(favicon_url)
if response.status_code == 200:
content_type = response.headers.get(
"content-type", "image/x-icon"
)
content_type = response.headers.get("content-type", "image/x-icon")
return Response(
content=response.content,
media_type=content_type,
@ -11033,12 +11038,8 @@ async def get_favicon():
response.status_code,
)
if os.path.exists(default_favicon):
return FileResponse(
default_favicon, media_type="image/x-icon"
)
raise HTTPException(
status_code=404, detail="Favicon not found"
)
return FileResponse(default_favicon, media_type="image/x-icon")
raise HTTPException(status_code=404, detail="Favicon not found")
except HTTPException:
raise
except Exception as e:
@ -11046,20 +11047,14 @@ async def get_favicon():
"Error downloading favicon from %s: %s", favicon_url, e
)
if os.path.exists(default_favicon):
return FileResponse(
default_favicon, media_type="image/x-icon"
)
raise HTTPException(
status_code=404, detail="Favicon not found"
)
return FileResponse(default_favicon, media_type="image/x-icon")
raise HTTPException(status_code=404, detail="Favicon not found")
else:
if os.path.exists(favicon_url):
return FileResponse(favicon_url, media_type="image/x-icon")
if os.path.exists(default_favicon):
return FileResponse(default_favicon, media_type="image/x-icon")
raise HTTPException(
status_code=404, detail="Favicon not found"
)
raise HTTPException(status_code=404, detail="Favicon not found")
#### INVITATION MANAGEMENT ####
@ -12545,7 +12540,9 @@ async def reload_anthropic_beta_headers(
},
)
provider_count = sum(1 for k in new_config.keys() if k not in ["provider_aliases", "description"])
provider_count = sum(
1 for k in new_config.keys() if k not in ["provider_aliases", "description"]
)
verbose_proxy_logger.info(
f"Anthropic beta headers config reloaded successfully in current pod. Providers: {provider_count}"
)
@ -12557,7 +12554,9 @@ async def reload_anthropic_beta_headers(
"timestamp": current_time.isoformat(),
}
except Exception as e:
verbose_proxy_logger.exception(f"Failed to reload anthropic beta headers: {str(e)}")
verbose_proxy_logger.exception(
f"Failed to reload anthropic beta headers: {str(e)}"
)
raise HTTPException(
status_code=500, detail=f"Failed to reload anthropic beta headers: {str(e)}"
)
@ -12679,7 +12678,8 @@ async def cancel_anthropic_beta_headers_reload(
f"Failed to cancel anthropic beta headers reload: {str(e)}"
)
raise HTTPException(
status_code=500, detail=f"Failed to cancel anthropic beta headers reload: {str(e)}"
status_code=500,
detail=f"Failed to cancel anthropic beta headers reload: {str(e)}",
)
@ -12726,7 +12726,9 @@ async def get_anthropic_beta_headers_reload_status(
)
if config_record is None or config_record.param_value is None:
verbose_proxy_logger.info("No anthropic beta headers reload configuration found")
verbose_proxy_logger.info(
"No anthropic beta headers reload configuration found"
)
return {
"scheduled": False,
"interval_hours": None,
@ -12752,7 +12754,9 @@ async def get_anthropic_beta_headers_reload_status(
# Use pod's in-memory last reload time
if last_anthropic_beta_headers_reload is not None:
try:
last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload)
last_reload_time = datetime.fromisoformat(
last_anthropic_beta_headers_reload
)
time_since_last_reload = current_time - last_reload_time
hours_since_last_reload = time_since_last_reload.total_seconds() / 3600

94
tests/litellm/proxy/guardrails/test_custom_code_security.py Normal file
View file

@ -0,0 +1,94 @@
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 (
CustomCodeGuardrail,
)
# Phase 4.1: Test forbidden pattern validation
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 test_validate_custom_code_from_subprocess():
code = (
"from subprocess import call\ndef apply_guardrail(i, r, t):\n return allow()"
)
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__") == {}
@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")
from litellm.types.utils import GenericGuardrailAPIInputs
result = await guardrail.apply_guardrail(
inputs=GenericGuardrailAPIInputs(texts=["test"]),
request_data={},
input_type="request",
)
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.