scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(proxy): improve input validation on management endpoints

Browse source
This commit is contained in:
jayden 2026-04-09 14:14:53 -07:00
parent 3a6db708ce
commit d910a95661
No known key found for this signature in database
4 changed files with 258 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/integrations/dotprompt/prompt_manager.py
View file

@ -7,7 +7,8 @@ from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union
import yaml
from jinja2 import DictLoader, Environment, select_autoescape
from jinja2 import DictLoader, select_autoescape
from jinja2.sandbox import ImmutableSandboxedEnvironment
class PromptTemplate:
@ -59,7 +60,10 @@ class PromptManager:
self.prompt_directory = Path(prompt_directory) if prompt_directory else None
self.prompts: Dict[str, PromptTemplate] = {}
self.prompt_file = prompt_file
self.jinja_env = Environment(
# Sandboxed env: templates can come from user input via /prompts/test,
# so we must block access to unsafe Python attributes and mutation of
# caller-supplied mutables.
self.jinja_env = ImmutableSandboxedEnvironment(
loader=DictLoader({}),
autoescape=select_autoescape(["html", "xml"]),
# Use Handlebars-style delimiters to match Dotprompt spec

77
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -456,6 +456,34 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict:
return data_json
def _check_allowed_routes_caller_permission(
allowed_routes: Optional[list],
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""
Only proxy admins may set `allowed_routes` on a key.
`allowed_routes` bypasses the standard role-based route gate in
RouteChecks.non_proxy_admin_allowed_routes_check, so if a non-admin is
allowed to set it they can grant themselves access to any endpoint.
Non-admins should use `key_type` to pick a preset route bucket instead.
"""
# Empty list is the default on GenerateKeyRequest — treat as "not set".
if not allowed_routes:
return
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
return
raise HTTPException(
status_code=403,
detail={
"error": (
"Only proxy admins can set `allowed_routes` on a key. "
"Use `key_type` to pick a preset route bucket instead."
)
},
)
async def validate_team_id_used_in_service_account_request(
team_id: Optional[str],
prisma_client: Optional[PrismaClient],
@ -740,9 +768,9 @@ async def _common_key_generation_helper( # noqa: PLR0915
request_type="key", **data_json, table_name="key"
)
response[
"soft_budget"
] = data.soft_budget # include the user-input soft budget in the response
response["soft_budget"] = (
data.soft_budget
) # include the user-input soft budget in the response
response = GenerateKeyResponse(**response)
@ -1254,6 +1282,12 @@ async def generate_key_fn(
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN, detail=message
)
_check_allowed_routes_caller_permission(
allowed_routes=data.allowed_routes,
user_api_key_dict=user_api_key_dict,
)
# For non-admin internal users: auto-assign caller's user_id if not provided
# This prevents creating unbound keys with no user association (LIT-1884)
_is_proxy_admin = (
@ -1888,6 +1922,11 @@ async def _validate_update_key_data(
"""Validate permissions and constraints for key update."""
_is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
_check_allowed_routes_caller_permission(
allowed_routes=data.allowed_routes,
user_api_key_dict=user_api_key_dict,
)
# Prevent non-admin from removing user_id (setting to empty string) (LIT-1884)
if data.user_id is not None and data.user_id == "" and not _is_proxy_admin:
raise HTTPException(
@ -3233,10 +3272,10 @@ async def delete_verification_tokens(
try:
if prisma_client:
tokens = [_hash_token_if_needed(token=key) for key in tokens]
_keys_being_deleted: List[
LiteLLM_VerificationToken
] = await prisma_client.db.litellm_verificationtoken.find_many(
where={"token": {"in": tokens}}
_keys_being_deleted: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={"token": {"in": tokens}}
)
)
if len(_keys_being_deleted) == 0:
@ -3436,9 +3475,9 @@ async def _rotate_master_key( # noqa: PLR0915
from litellm.proxy.proxy_server import proxy_config
try:
models: Optional[
List
] = await prisma_client.db.litellm_proxymodeltable.find_many()
models: Optional[List] = (
await prisma_client.db.litellm_proxymodeltable.find_many()
)
except Exception:
models = None
# 2. process model table
@ -4078,11 +4117,11 @@ async def validate_key_list_check(
param="user_id",
code=status.HTTP_403_FORBIDDEN,
)
complete_user_info_db_obj: Optional[
BaseModel
] = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
include={"organization_memberships": True},
complete_user_info_db_obj: Optional[BaseModel] = (
await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_api_key_dict.user_id},
include={"organization_memberships": True},
)
)
if complete_user_info_db_obj is None:
@ -4165,10 +4204,10 @@ async def _fetch_user_team_objects(
if complete_user_info is None or not complete_user_info.teams:
return []
teams: Optional[
List[BaseModel]
] = await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": complete_user_info.teams}}
teams: Optional[List[BaseModel]] = (
await prisma_client.db.litellm_teamtable.find_many(
where={"team_id": {"in": complete_user_info.teams}}
)
)
if teams is None:
return []

27
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -1,6 +1,7 @@
#### CRUD ENDPOINTS for UI Settings #####
import json
from typing import Any, Dict, List, Optional, Union
from urllib.parse import urlparse
from fastapi import APIRouter, Depends, File, HTTPException, UploadFile
@ -817,6 +818,29 @@ async def get_ui_theme_settings():
)
def _validate_public_image_url(value: Optional[str], field_name: str) -> None:
"""
Reject anything that isn't a plain http(s) URL with a host. This value is
later served via the unauthenticated /get_image endpoint, so local paths
like "/etc/passwd" or "file://..." must not be accepted.
"""
if value is None:
return
if not isinstance(value, str) or not value.strip():
return
parsed = urlparse(value.strip())
if parsed.scheme not in ("http", "https") or not parsed.netloc:
raise HTTPException(
status_code=400,
detail={
"error": (
f"Invalid {field_name}: must be an http(s) URL with a host. "
"Local filesystem paths and non-http schemes are not allowed."
)
},
)
@router.patch(
"/update/ui_theme_settings",
tags=["UI Theme Settings"],
@ -831,6 +855,9 @@ async def update_ui_theme_settings(theme_config: UIThemeConfig):
from litellm.proxy.proxy_server import proxy_config, store_model_in_db
_validate_public_image_url(theme_config.logo_url, "logo_url")
_validate_public_image_url(theme_config.favicon_url, "favicon_url")
if store_model_in_db is not True:
raise HTTPException(
status_code=500,

167
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -8579,3 +8579,170 @@ def test_enforce_upperbound_no_config_is_noop():
assert data.tpm_limit == 999999
finally:
litellm.upperbound_key_generate_params = original
class TestAllowedRoutesCallerPermission:
"""
Non-admins must not be able to set `allowed_routes` on a key. The field
bypasses the role-based route gate in
RouteChecks.non_proxy_admin_allowed_routes_check, so allowing a non-admin
to populate it grants them arbitrary endpoint access.
"""
@pytest.mark.asyncio
async def test_non_admin_generate_key_with_allowed_routes_rejected(self):
data = GenerateKeyRequest(
key_alias="escalate",
allowed_routes=["/*"],
)
user_api_key_dict = UserAPIKeyAuth(
user_id="internal-user-123",
user_role=LitellmUserRoles.INTERNAL_USER,
)
mock_prisma_client = AsyncMock()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch(
"litellm.proxy.proxy_server.user_api_key_cache", MagicMock()
), patch("litellm.proxy.proxy_server.user_custom_key_generate", None), patch(
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
new_callable=AsyncMock,
return_value=MagicMock(),
):
with pytest.raises(ProxyException) as exc_info:
await generate_key_fn(
data=data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
assert str(exc_info.value.code) == "403"
assert "allowed_routes" in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_admin_generate_key_with_allowed_routes_allowed(self):
data = GenerateKeyRequest(
key_alias="admin-key",
allowed_routes=["/chat/completions"],
user_id="admin-user",
)
user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
mock_prisma_client = AsyncMock()
stub_response = MagicMock()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch(
"litellm.proxy.proxy_server.user_api_key_cache", MagicMock()
), patch("litellm.proxy.proxy_server.user_custom_key_generate", None), patch(
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
new_callable=AsyncMock,
return_value=stub_response,
):
result = await generate_key_fn(
data=data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
assert result is stub_response
@pytest.mark.asyncio
async def test_non_admin_generate_key_default_empty_allowed_routes_ok(self):
"""
Regression guard: GenerateKeyRequest.allowed_routes defaults to [], so
the helper must treat empty-list as "not set" or every non-admin key
creation breaks.
"""
data = GenerateKeyRequest(key_alias="plain-key")
user_api_key_dict = UserAPIKeyAuth(
user_id="internal-user-123",
user_role=LitellmUserRoles.INTERNAL_USER,
)
mock_prisma_client = AsyncMock()
stub_response = MagicMock()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch(
"litellm.proxy.proxy_server.user_api_key_cache", MagicMock()
), patch("litellm.proxy.proxy_server.user_custom_key_generate", None), patch(
"litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper",
new_callable=AsyncMock,
return_value=stub_response,
):
result = await generate_key_fn(
data=data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
assert result is stub_response
@pytest.mark.asyncio
async def test_non_admin_update_key_with_allowed_routes_rejected(self):
from litellm.proxy.management_endpoints.key_management_endpoints import (
update_key_fn,
)
data = UpdateKeyRequest(key="sk-test", allowed_routes=["/*"])
user_api_key_dict = UserAPIKeyAuth(
user_id="internal-user-123",
user_role=LitellmUserRoles.INTERNAL_USER,
)
mock_prisma_client = AsyncMock()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch(
"litellm.proxy.proxy_server.user_api_key_cache", MagicMock()
), patch("litellm.proxy.proxy_server.user_custom_key_update", None), patch(
"litellm.proxy.proxy_server.llm_router", None
), patch("litellm.proxy.proxy_server.premium_user", True), patch(
"litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()
), patch(
"litellm.proxy.management_endpoints.key_management_endpoints._get_and_validate_existing_key",
new_callable=AsyncMock,
return_value=MagicMock(),
):
with pytest.raises(ProxyException) as exc_info:
await update_key_fn(
request=MagicMock(),
data=data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
assert str(exc_info.value.code) == "403"
assert "allowed_routes" in str(exc_info.value.message)
def test_jinja_prompt_manager_is_sandboxed():
"""
PromptManager renders user-supplied templates via /prompts/test, so its
jinja env must reject access to unsafe Python attributes like
``__class__`` and ``__mro__``.
"""
from jinja2.exceptions import SecurityError
from litellm.integrations.dotprompt.prompt_manager import PromptManager
pm = PromptManager()
template = pm.jinja_env.from_string("{{ ''.__class__.__mro__ }}")
with pytest.raises(SecurityError):
template.render()
def test_validate_public_image_url_rejects_local_paths():
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
_validate_public_image_url,
)
for bad in ("/etc/passwd", "file:///etc/passwd", "../../etc/passwd"):
with pytest.raises(HTTPException) as exc_info:
_validate_public_image_url(bad, "logo_url")
assert exc_info.value.status_code == 400
def test_validate_public_image_url_accepts_http_and_noop_empty():
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
_validate_public_image_url,
)
_validate_public_image_url("https://example.com/logo.png", "logo_url")
_validate_public_image_url("http://cdn.internal/logo.svg", "logo_url")
_validate_public_image_url(None, "logo_url")
_validate_public_image_url("", "logo_url")
_validate_public_image_url(" ", "logo_url")