diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index 4bfe9d318747..75229bacc8fe 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -588,24 +588,21 @@ async def update_project( # noqa: PLR0915 param="project_id", ) - # Validate team exists and get team object for limit + permission checks - team_id_to_check = data.team_id or existing_project.team_id - team_obj_for_checks = None - if team_id_to_check is not None: - team_obj_for_checks = await _validate_team_exists( - team_id=team_id_to_check, prisma_client=prisma_client + # Permission to *edit* the project must be evaluated against the + # project's CURRENT team. Sourcing the team from `data.team_id` + # would let an admin of any team pass the check by supplying their + # own team_id, hijacking the project (VERIA-55). + target_team_id = data.team_id or existing_project.team_id + target_team_obj = None + if target_team_id is not None: + target_team_obj = await _validate_team_exists( + team_id=target_team_id, prisma_client=prisma_client ) - # Check if user has permission to update this project has_permission = await _check_user_permission_for_project( user_api_key_dict=user_api_key_dict, team_id=existing_project.team_id, prisma_client=prisma_client, - team_object=( - LiteLLM_TeamTable(**team_obj_for_checks.model_dump()) - if team_obj_for_checks - else None - ), ) if not has_permission: @@ -614,10 +611,32 @@ async def update_project( # noqa: PLR0915 detail={"error": "Only admins or team admins can update projects"}, ) + # Reassigning to a different team also requires admin rights on the + # destination team — otherwise a team admin could shed projects into + # an unsuspecting team's namespace. + if data.team_id is not None and data.team_id != existing_project.team_id: + can_assign_to_target = await _check_user_permission_for_project( + user_api_key_dict=user_api_key_dict, + team_id=data.team_id, + prisma_client=prisma_client, + team_object=( + LiteLLM_TeamTable(**target_team_obj.model_dump()) + if target_team_obj + else None + ), + ) + if not can_assign_to_target: + raise HTTPException( + status_code=403, + detail={ + "error": "Cannot reassign project to a team you are not an admin of" + }, + ) + # Validate project limits against team limits - if team_obj_for_checks is not None: + if target_team_obj is not None: _check_team_project_limits( - team_object=LiteLLM_TeamTable(**team_obj_for_checks.model_dump()), + team_object=LiteLLM_TeamTable(**target_team_obj.model_dump()), data=data, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2e5140e0e343..9159a8ff9daa 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -11,7 +11,7 @@ import re import secrets from datetime import datetime, timezone -from typing import Any, List, Optional, Tuple, Union, cast +from typing import Any, Iterator, List, Optional, Tuple, Union, cast import fastapi from fastapi import HTTPException, Request, WebSocket, status @@ -2271,10 +2271,6 @@ async def _enforce_key_and_fallback_model_access( route=route, request=request, ) - fallback_models = cast( - Optional[List[ALL_FALLBACK_MODEL_VALUES]], - request_data.get("fallbacks", None), - ) if model is not None: await can_key_call_model( @@ -2284,20 +2280,69 @@ async def _enforce_key_and_fallback_model_access( llm_router=llm_router, ) - if fallback_models is not None: - for m in fallback_models: - await can_key_call_model( - model=m["model"] if isinstance(m, dict) else m, - llm_model_list=llm_model_list, - valid_token=valid_token, - llm_router=llm_router, - ) - await is_valid_fallback_model( - model=m["model"] if isinstance(m, dict) else m, - llm_router=llm_router, - user_model=None, + # Validate every fallback model name reachable by this request. + # All three fields (``fallbacks``, ``context_window_fallbacks``, + # ``content_policy_fallbacks``) are forwarded to the router as + # per-request kwargs whether they appear at the top level of + # ``request_data`` or nested under ``router_settings_override``. + # Both surfaces must be validated against the API key's model + # allowlist or a caller can smuggle a restricted model. VERIA-44. + fallback_names: List[str] = [] + override_settings = request_data.get("router_settings_override") + for _fb_key in ROUTER_FALLBACK_FIELDS: + fallback_names.extend( + iter_router_fallback_model_names(request_data.get(_fb_key)) + ) + if isinstance(override_settings, dict): + fallback_names.extend( + iter_router_fallback_model_names(override_settings.get(_fb_key)) ) + for _name in dict.fromkeys(fallback_names): # dedupe, preserve order + await can_key_call_model( + model=_name, + llm_model_list=llm_model_list, + valid_token=valid_token, + llm_router=llm_router, + ) + await is_valid_fallback_model( + model=_name, + llm_router=llm_router, + user_model=None, + ) + + +ROUTER_FALLBACK_FIELDS: Tuple[str, ...] = ( + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", +) + + +def iter_router_fallback_model_names(fallbacks: Any) -> Iterator[str]: + """Yield leaf model names from any of the supported fallbacks shapes. + + Handles the simple top-level shape (``str`` or ``{"model": str}``) and + the nested router-config shape (``[{primary: [fallback_list]}]``). + """ + if not isinstance(fallbacks, list): + return + for entry in fallbacks: + if isinstance(entry, str): + yield entry + elif isinstance(entry, dict): + if isinstance(entry.get("model"), str): + yield entry["model"] + continue + for fallback_list in entry.values(): + if not isinstance(fallback_list, list): + continue + for m in fallback_list: + if isinstance(m, str): + yield m + elif isinstance(m, dict) and isinstance(m.get("model"), str): + yield m["model"] + async def _run_post_custom_auth_checks( valid_token: UserAPIKeyAuth, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a01f5e63211d..1f66d9fec859 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -809,6 +809,20 @@ async def _common_key_generation_helper( # noqa: PLR0915 from litellm.proxy.proxy_server import prisma_client, user_api_key_cache if prisma_client: + # Mirror the membership rule applied to /key/update: when the + # caller specifies an organization_id, require that they are a + # member of (or proxy admin over) the target organization. + _is_proxy_admin = ( + user_api_key_dict.user_role is not None + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ) + if not _is_proxy_admin: + await _validate_caller_can_assign_key_org( + user_api_key_dict=user_api_key_dict, + organization_id=data.organization_id, + prisma_client=prisma_client, + ) + org_table = await get_org_object( org_id=data.organization_id, user_api_key_cache=user_api_key_cache, @@ -1168,6 +1182,42 @@ def check_org_key_rpm_tpm_limits( ) +async def _validate_caller_can_assign_key_org( + user_api_key_dict: UserAPIKeyAuth, + organization_id: str, + prisma_client: PrismaClient, +) -> None: + """Reject ``/key/update`` requests that point a key at an organization + the caller does not belong to. + + Mirrors the org-membership rule already enforced on ``/key/list`` in + ``validate_key_list_check``. Proxy admins are checked at the call site. + """ + if user_api_key_dict.user_id is None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Cannot assign a key to an organization without a user_id on the caller's token", + ) + + user_row = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_api_key_dict.user_id}, + include={"organization_memberships": True}, + ) + memberships = ( + getattr(user_row, "organization_memberships", None) if user_row else None + ) + member_org_ids = { + membership.organization_id + for membership in (memberships or []) + if membership.organization_id is not None + } + if organization_id not in member_org_ids: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"Caller is not a member of organization_id={organization_id}", + ) + + async def _check_org_key_limits( org_table: LiteLLM_OrganizationTable, data: Union[GenerateKeyRequest, UpdateKeyRequest], @@ -2168,10 +2218,26 @@ async def _validate_update_key_data( user_api_key_cache=user_api_key_cache, ) + # When the caller asks to change the key's organization_id, require that + # they are a member of (or a proxy admin over) the target organization. + # Without this gate, any caller could assign their key to an arbitrary + # organization_id by passing it in the request body — VERIA-55 secondary + # IDOR. The check mirrors the membership rule already used on the + # `/key/list` filter path in `validate_key_list_check`. + _existing_org_id = getattr(existing_key_row, "organization_id", None) + if ( + data.organization_id is not None + and data.organization_id != _existing_org_id + and not _is_proxy_admin + ): + await _validate_caller_can_assign_key_org( + user_api_key_dict=user_api_key_dict, + organization_id=data.organization_id, + prisma_client=prisma_client, + ) + # Check org key limits only when throughput-related fields or organization_id change - _org_id_to_check = data.organization_id or getattr( - existing_key_row, "organization_id", None - ) + _org_id_to_check = data.organization_id or _existing_org_id _throughput_fields_changed = ( data.organization_id is not None or data.tpm_limit is not None @@ -3868,6 +3934,22 @@ async def _execute_virtual_key_regeneration( """Generate new token, update DB, invalidate cache, and return response.""" from litellm.proxy.proxy_server import hash_token + # Apply the same membership rule used on /key/update: when the caller + # asks to point the regenerated key at a different organization_id, + # require they are a member of (or proxy admin over) the target org. + if data is not None and data.organization_id is not None: + _existing_org_id = getattr(key_in_db, "organization_id", None) + _is_proxy_admin = ( + user_api_key_dict.user_role is not None + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ) + if data.organization_id != _existing_org_id and not _is_proxy_admin: + await _validate_caller_can_assign_key_org( + user_api_key_dict=user_api_key_dict, + organization_id=data.organization_id, + prisma_client=prisma_client, + ) + new_token = await get_new_token(data=data) new_token_hash = hash_token(new_token) new_token_key_name = f"sk-...{new_token[-4:]}" diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 17cc43745602..bfe6b8484faf 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -6,6 +6,18 @@ import litellm from litellm.proxy._types import UserAPIKeyAuth +# Router-internal mock_testing_* flag names — kept in sync with +# ``litellm.types.router.MockRouterTestingParams`` by the test +# ``test_mock_testing_kwarg_names_matches_dataclass``. Hardcoding (rather +# than deriving via ``dataclasses.fields(MockRouterTestingParams)`` at +# import time) avoids a cyclic import: ``litellm.types.router`` imports +# back into proxy modules before this module finishes loading. +_MOCK_TESTING_KWARG_NAMES: tuple = ( + "mock_testing_fallbacks", + "mock_testing_context_fallbacks", + "mock_testing_content_policy_fallbacks", +) + if TYPE_CHECKING: from litellm.router import Router as _Router @@ -322,6 +334,13 @@ async def route_request( # noqa: PLR0915 - Complex routing function, refactorin """ await add_shared_session_to_data(data) + # Strip router-internal mock_testing_* flags. Combined with an + # unauthorized fallback in ``router_settings_override`` they let a + # caller deterministically execute requests against restricted + # models. VERIA-44. + for _key in _MOCK_TESTING_KWARG_NAMES: + data.pop(_key, None) + team_id = get_team_id_from_data(data) router_model_names = llm_router.model_names if llm_router is not None else [] diff --git a/tests/test_litellm/proxy/auth/test_router_override_fallback_auth.py b/tests/test_litellm/proxy/auth/test_router_override_fallback_auth.py new file mode 100644 index 000000000000..fc0e9aec501a --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_router_override_fallback_auth.py @@ -0,0 +1,237 @@ +""" +VERIA-44: ``router_settings_override.fallbacks`` must be validated +against the API key's model allowlist at auth time. Without this, the +override is promoted to per-request kwargs after auth and lets a caller +execute requests against models their API key cannot call. +""" + +from typing import List +from unittest.mock import AsyncMock, patch + +import pytest + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import ( + _enforce_key_and_fallback_model_access, + iter_router_fallback_model_names, +) + + +def _key_with_models(models: List[str]) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="hashed", + user_id="caller", + user_role=LitellmUserRoles.INTERNAL_USER, + models=models, + ) + + +# ── iter_router_fallback_model_names ───────────────────────────────────────── + + +def testiter_router_fallback_model_names_router_config_shape(): + """Router-config shape: ``[{primary: [fallback_list]}]``.""" + assert list( + iter_router_fallback_model_names( + [{"gpt-3.5-turbo": ["gpt-4", "claude-3"]}, {"gpt-4o": ["o1"]}] + ) + ) == ["gpt-4", "claude-3", "o1"] + + +def testiter_router_fallback_model_names_simple_string_shape(): + """Simple top-level shape: list of strings.""" + assert list(iter_router_fallback_model_names(["gpt-4", "claude-3"])) == [ + "gpt-4", + "claude-3", + ] + + +def testiter_router_fallback_model_names_client_side_shape(): + """ClientSideFallbackModel shape: ``[{"model": "..."}]``.""" + assert list( + iter_router_fallback_model_names([{"model": "gpt-4"}, {"model": "claude-3"}]) + ) == ["gpt-4", "claude-3"] + + +def testiter_router_fallback_model_names_empty_or_none(): + assert list(iter_router_fallback_model_names(None)) == [] + assert list(iter_router_fallback_model_names([])) == [] + assert list(iter_router_fallback_model_names("not a list")) == [] + + +# ── _enforce_key_and_fallback_model_access ──────────────────────────────────── + + +@pytest.mark.asyncio +async def test_router_override_fallbacks_validated_against_key_allowlist(): + """A fallback nested inside ``router_settings_override`` is validated + against the API key's allowed models — not just the top-level + ``fallbacks`` field.""" + valid_token = _key_with_models(["gpt-3.5-turbo"]) + request_data = { + "model": "gpt-3.5-turbo", + "router_settings_override": { + "fallbacks": [{"gpt-3.5-turbo": ["unauthorized-model"]}], + }, + } + + seen_models: List[str] = [] + + async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): + seen_models.append(model) + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.can_key_call_model", + side_effect=fake_can_key_call_model, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", + new=AsyncMock(), + ), + ): + await _enforce_key_and_fallback_model_access( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=None, + llm_model_list=None, + llm_router=None, + ) + + # Both the primary model and the override-nested fallback must be + # checked against the API key's allowlist. + assert "gpt-3.5-turbo" in seen_models + assert "unauthorized-model" in seen_models + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallback_field", + [ + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + ], +) +async def test_router_override_all_fallback_fields_validated(fallback_field): + """All three fallback fields the router accepts as per-request kwargs + are validated — context_window_fallbacks and content_policy_fallbacks + are promoted in route_llm_request.py too.""" + valid_token = _key_with_models(["gpt-3.5-turbo"]) + request_data = { + "model": "gpt-3.5-turbo", + "router_settings_override": { + fallback_field: [{"gpt-3.5-turbo": ["smuggled-model"]}], + }, + } + + seen: List[str] = [] + + async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): + seen.append(model) + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.can_key_call_model", + side_effect=fake_can_key_call_model, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", + new=AsyncMock(), + ), + ): + await _enforce_key_and_fallback_model_access( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=None, + llm_model_list=None, + llm_router=None, + ) + + assert "smuggled-model" in seen + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallback_field", + [ + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + ], +) +async def test_top_level_fallback_fields_validated(fallback_field): + """All three top-level fallback fields are forwarded to the router as + per-request kwargs, so all three must be validated against the API + key's allowlist. Greptile P1 follow-up: previously only the + ``fallbacks`` field was walked at the top level.""" + valid_token = _key_with_models(["gpt-3.5-turbo"]) + request_data = { + "model": "gpt-3.5-turbo", + fallback_field: [{"gpt-3.5-turbo": ["top-level-smuggled"]}], + } + + seen: List[str] = [] + + async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): + seen.append(model) + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.can_key_call_model", + side_effect=fake_can_key_call_model, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", + new=AsyncMock(), + ), + ): + await _enforce_key_and_fallback_model_access( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=None, + llm_model_list=None, + llm_router=None, + ) + + assert "top-level-smuggled" in seen + + +@pytest.mark.asyncio +async def test_router_override_without_fallbacks_does_not_break_auth(): + """``router_settings_override`` set without any fallback fields is a + no-op for the auth check — only the primary model is validated.""" + valid_token = _key_with_models(["gpt-3.5-turbo"]) + request_data = { + "model": "gpt-3.5-turbo", + "router_settings_override": {"num_retries": 3, "timeout": 30}, + } + + seen: List[str] = [] + + async def fake_can_key_call_model(model, llm_model_list, valid_token, llm_router): + seen.append(model) + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.can_key_call_model", + side_effect=fake_can_key_call_model, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.is_valid_fallback_model", + new=AsyncMock(), + ), + ): + await _enforce_key_and_fallback_model_access( + valid_token=valid_token, + request_data=request_data, + route="/v1/chat/completions", + request=None, + llm_model_list=None, + llm_router=None, + ) + + assert seen == ["gpt-3.5-turbo"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_project_org_authz.py b/tests/test_litellm/proxy/management_endpoints/test_project_org_authz.py new file mode 100644 index 000000000000..bd982480d60a --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_project_org_authz.py @@ -0,0 +1,195 @@ +""" +Unit tests for the VERIA-55 fixes: + +- Project update permission must be evaluated against the project's *current* + team, not a team supplied in the request body. +- Key update may not assign a key to an organization the caller is not a + member of. +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +# --------------------------------------------------------------------------- +# /project/update — _check_user_permission_for_project +# --------------------------------------------------------------------------- + + +def _make_prisma_with_team(team_id: str, admins: list): + prisma = MagicMock() + team_row = MagicMock() + team_row.team_id = team_id + team_row.admins = admins + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + return prisma + + +@pytest.mark.asyncio +async def test_project_perm_check_uses_current_team_not_caller_supplied(): + """The permission check must look at the project's existing team. Even + if the caller is admin of an unrelated team, they must not pass when no + explicit team_object is forced through.""" + from enterprise.litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _check_user_permission_for_project, + ) + + # Project lives on team-A, caller is admin only of team-B. + prisma = _make_prisma_with_team(team_id="team-A", admins=["alice"]) + caller = UserAPIKeyAuth( + user_id="bob", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + has_perm = await _check_user_permission_for_project( + user_api_key_dict=caller, + team_id="team-A", + prisma_client=prisma, + ) + assert has_perm is False + prisma.db.litellm_teamtable.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_project_perm_check_allows_team_admin_of_existing_team(): + from enterprise.litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _check_user_permission_for_project, + ) + + prisma = _make_prisma_with_team(team_id="team-A", admins=["alice"]) + alice = UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + has_perm = await _check_user_permission_for_project( + user_api_key_dict=alice, + team_id="team-A", + prisma_client=prisma, + ) + assert has_perm is True + + +@pytest.mark.asyncio +async def test_project_perm_check_proxy_admin_always_allowed(): + from enterprise.litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _check_user_permission_for_project, + ) + + prisma = MagicMock() + admin = UserAPIKeyAuth( + user_id="root", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ) + + has_perm = await _check_user_permission_for_project( + user_api_key_dict=admin, + team_id="team-A", + prisma_client=prisma, + ) + assert has_perm is True + # Admin shortcut should not even hit the DB. + prisma.db.litellm_teamtable.find_unique.assert_not_called() + + +# --------------------------------------------------------------------------- +# /key/update — _validate_caller_can_assign_key_org +# --------------------------------------------------------------------------- + + +def _make_prisma_with_user_orgs(user_id: str, org_ids: list): + prisma = MagicMock() + user_row = MagicMock() + user_row.organization_memberships = [ + MagicMock(organization_id=org_id) for org_id in org_ids + ] + prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + return prisma + + +@pytest.mark.asyncio +async def test_assign_key_org_allows_member(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_caller_can_assign_key_org, + ) + + prisma = _make_prisma_with_user_orgs("alice", ["org-1", "org-2"]) + caller = UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + # Should not raise. + await _validate_caller_can_assign_key_org( + user_api_key_dict=caller, + organization_id="org-2", + prisma_client=prisma, + ) + + +@pytest.mark.asyncio +async def test_assign_key_org_blocks_non_member(): + """The IDOR: caller asks to point a key at an org they don't belong to.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_caller_can_assign_key_org, + ) + + prisma = _make_prisma_with_user_orgs("alice", ["org-1"]) + caller = UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + with pytest.raises(HTTPException) as exc_info: + await _validate_caller_can_assign_key_org( + user_api_key_dict=caller, + organization_id="someone-elses-org", + prisma_client=prisma, + ) + assert exc_info.value.status_code == 403 + assert "someone-elses-org" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_assign_key_org_blocks_caller_without_user_id(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_caller_can_assign_key_org, + ) + + prisma = MagicMock() + caller = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + with pytest.raises(HTTPException) as exc_info: + await _validate_caller_can_assign_key_org( + user_api_key_dict=caller, + organization_id="org-1", + prisma_client=prisma, + ) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_assign_key_org_blocks_caller_with_no_memberships(): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_caller_can_assign_key_org, + ) + + prisma = MagicMock() + user_row = MagicMock() + user_row.organization_memberships = None + prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + + caller = UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + with pytest.raises(HTTPException) as exc_info: + await _validate_caller_can_assign_key_org( + user_api_key_dict=caller, + organization_id="org-1", + prisma_client=prisma, + ) + assert exc_info.value.status_code == 403 diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index bfea21e705ea..98b0b6be0255 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -241,6 +241,53 @@ async def test_route_request_with_router_settings_override_preserves_existing(): assert call_kwargs["timeout"] == 30 +def test_mock_testing_kwarg_names_matches_dataclass(): + """``_MOCK_TESTING_KWARG_NAMES`` is hardcoded to avoid a cyclic import + against ``litellm.types.router``. This test guards against drift — + if a new ``mock_testing_*`` field is added to ``MockRouterTestingParams`` + the strip list must be updated to keep covering it.""" + from dataclasses import fields + + from litellm.proxy.route_llm_request import _MOCK_TESTING_KWARG_NAMES + from litellm.types.router import MockRouterTestingParams + + assert set(_MOCK_TESTING_KWARG_NAMES) == { + f.name for f in fields(MockRouterTestingParams) + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "mock_flag", + [ + "mock_testing_fallbacks", + "mock_testing_context_fallbacks", + "mock_testing_content_policy_fallbacks", + ], +) +async def test_route_request_strips_mock_testing_flags(mock_flag): + """VERIA-44: router-internal testing flags must not survive a + user-supplied request body. Without this strip, an attacker can + combine ``mock_testing_fallbacks=true`` with an unauthorized fallback + in ``router_settings_override`` to deterministically execute requests + against restricted models.""" + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}], + mock_flag: True, + } + llm_router = MagicMock() + llm_router.acompletion.return_value = "ok" + + await route_request(data, llm_router, None, "acompletion") + + call_kwargs = llm_router.acompletion.call_args[1] + assert mock_flag not in call_kwargs + # The flag is also gone from the original data dict so any subsequent + # processing (e.g. logging) doesn't see it either. + assert mock_flag not in data + + @pytest.mark.parametrize( "route_type", ["agenerate_content", "agenerate_content_stream"] )