scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

fix(proxy): authorize /health/test_connection against loaded deployment's team_id (VERIA-441) (#31767)

Browse source 
* fix(proxy): authorize /health/test_connection against loaded deployment's team_id (VERIA-441)

POST /health/test_connection looked up a deployment by request-supplied model_info.id, dumped its
litellm_params (including api_key) into the outbound probe, merged request params over it, and then
authorized the call against the caller-supplied model_info.team_id. A team admin could pass another
team's deployment id together with their own team_id and an attacker-controlled api_base, sending
the victim team's provider key to that URL.

Capture the loaded deployment's model_info alongside its litellm_params in both the id-lookup and
the model_name fallback paths, and pass that captured value to can_user_make_model_call. When no
deployment is loaded (caller is probing fresh, request-supplied credentials), keep using the
request body's model_info; no foreign deployment is in scope and the existing role check still
requires admin or team-admin.

Add two regression tests that wrap (not mock) ModelManagementAuthChecks.can_user_make_model_call,
one per resolution path, asserting HTTP 403 and that the auth check was reached with the loaded
deployment's team_id. Both fail on the pre-fix code.

* test(health): add positive-path regression through real auth (VERIA-441)

The two deny tests already exercise the real (wrapped) ModelManagementAuthChecks.
Add a matching positive-path test so a mutation that swaps the auth team_id for
a deny-all value on the legit path also fails: loaded deployment owned by team-X,
caller admin of team-X -> asserts HTTP 200 and that the auth check ran with the
LOADED deployment's team_id.

* refactor(test): rename health endpoint tests for clarity (VERIA-441)

Rename test functions and variables from attacker/victim/owner framing to
neutral team-a/team-b terminology. Update docstrings to remove exploit-specific
language. Tests remain functionally identical, covering deny paths (cross-team
deployments) and the positive path (same-team deployments).
This commit is contained in:
yucheng-berri 2026-07-01 17:05:49 -07:00 • committed by GitHub
parent 700afbb6b2
commit a2f5bb1868
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 308 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -1807,6 +1807,7 @@ async def test_model_connection(
# Look up model configuration from router if model name is provided
# This gets the litellm_params from proxy config (with resolved env vars)
config_litellm_params: dict = {}
loaded_model_info: Optional[dict] = None
if llm_router is not None:
# Prefer disambiguation by deployment id (`model_info.id`) when
# the caller supplies it. This is required when multiple
@ -1825,6 +1826,7 @@ async def test_model_connection(
if deployment_by_id is not None:
config_litellm_params = deployment_by_id.litellm_params.model_dump(exclude_none=True)
loaded_model_info = deployment_by_id.model_info.model_dump(exclude_none=True)
elif model_name:
# Fall back to model_name lookup for callers (e.g. the
# "Add Model" wizard, or curl) that don't supply an id.
@ -1846,6 +1848,7 @@ async def test_model_connection(
# config. These already have resolved environment
# variables from proxy config.
config_litellm_params = dict(deployments[0].get("litellm_params", {}))
loaded_model_info = dict(deployments[0].get("model_info") or {})
except Exception as e:
verbose_proxy_logger.debug(
f"Could not find model {model_name} in router: {e}. Proceeding with request params only."
@ -1856,11 +1859,12 @@ async def test_model_connection(
litellm_params = {**config_litellm_params, **request_litellm_params}
## Auth check
auth_model_info = loaded_model_info if loaded_model_info is not None else model_info
await ModelManagementAuthChecks.can_user_make_model_call(
model_params=Deployment(
model_name="test_model",
litellm_params=LiteLLM_Params(**litellm_params),
model_info=model_info,
model_info=auth_model_info,
),
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,

303
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
View file

@ -696,6 +696,309 @@ async def test_test_model_connection_falls_back_to_deployments_zero_without_id()
assert model_params.get("api_key") == "fake-key-A"
@pytest.mark.asyncio
async def test_test_model_connection_uses_loaded_deployment_team_id():
"""
/health/test_connection must authorize using the team_id of the
deployment it actually loaded (by model_info.id), not the team_id
supplied in the request body. Requesting team A's deployment while
authenticated as an admin of team B must be denied.
"""
from fastapi import HTTPException
from litellm.proxy._types import LiteLLM_TeamTable
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelManagementAuthChecks,
)
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
mock_request = MagicMock()
requester_team_id = "team-b"
deployment_owner_team_id = "team-a"
deployment_id = "team-a-deployment-id"
requester_user_api_key_dict = UserAPIKeyAuth(
token="requester-token",
user_id="team-b-admin-user",
team_id=requester_team_id,
user_role=LitellmUserRoles.INTERNAL_USER,
)
mock_prisma_client = MagicMock()
other_team_deployment = Deployment(
model_name="team-a-model",
litellm_params=LiteLLM_Params(
model="openai/gpt-4o",
api_base="https://team-a-api.invalid/v1",
api_key="TEAM-A-API-KEY",
),
model_info=ModelInfo(id=deployment_id, team_id=deployment_owner_team_id),
)
mock_router = MagicMock()
mock_router.get_deployment.return_value = other_team_deployment
async def fake_find_unique(*, where):
team_id = where["team_id"]
if team_id == requester_team_id:
return SimpleNamespace(
model_dump=lambda: LiteLLM_TeamTable(
team_id=requester_team_id,
members_with_roles=[
{
"user_id": "team-b-admin-user",
"role": "admin",
}
],
).model_dump()
)
if team_id == deployment_owner_team_id:
return SimpleNamespace(
model_dump=lambda: LiteLLM_TeamTable(
team_id=deployment_owner_team_id,
members_with_roles=[],
).model_dump()
)
return None
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
patch("litellm.proxy.proxy_server.llm_router", mock_router),
patch("litellm.proxy.proxy_server.premium_user", True),
patch.object(
ModelManagementAuthChecks,
"can_user_make_model_call",
wraps=ModelManagementAuthChecks.can_user_make_model_call,
) as spy_auth_check,
patch(
"litellm.proxy.management_endpoints.model_management_endpoints.TeamRepository"
) as MockTeamRepo,
):
mock_team_repo_instance = MagicMock()
mock_team_repo_instance.table.find_unique = AsyncMock(
side_effect=fake_find_unique
)
MockTeamRepo.return_value = mock_team_repo_instance
with pytest.raises(HTTPException) as exc_info:
await health_test_model_connection(
request=mock_request,
mode="chat",
litellm_params={
"model": "openai/gpt-4o",
"api_base": "https://swapped-base.invalid/v1",
},
model_info={
"id": deployment_id,
"team_id": requester_team_id,
},
user_api_key_dict=requester_user_api_key_dict,
)
assert exc_info.value.status_code == 403
assert spy_auth_check.called
passed_model_params = spy_auth_check.call_args.kwargs["model_params"]
assert passed_model_params.model_info.team_id == deployment_owner_team_id, (
"Auth check must run against the loaded deployment's team_id "
f"({deployment_owner_team_id!r}); got "
f"{passed_model_params.model_info.team_id!r}."
)
@pytest.mark.asyncio
async def test_test_model_connection_uses_loaded_deployment_team_id_via_model_name_fallback():
"""
Companion to the id-lookup case: when the caller provides only a model
name (no `model_info.id`) and that name resolves via the router's
`model_name` fallback to a deployment owned by a different team, the
auth check must still run against the loaded deployment's `team_id`,
not the caller-supplied one in the request body.
"""
from fastapi import HTTPException
from litellm.proxy._types import LiteLLM_TeamTable
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelManagementAuthChecks,
)
mock_request = MagicMock()
requester_team_id = "team-b-2"
deployment_owner_team_id = "team-a-2"
requester_user_api_key_dict = UserAPIKeyAuth(
token="requester-token-2",
user_id="team-b-admin-user-2",
team_id=requester_team_id,
user_role=LitellmUserRoles.INTERNAL_USER,
)
mock_prisma_client = MagicMock()
other_team_deployment_dict = {
"model_name": "shared-model-name",
"litellm_params": {
"model": "openai/gpt-4o",
"api_base": "https://team-a-api-2.invalid/v1",
"api_key": "TEAM-A-API-KEY-2",
},
"model_info": {
"id": "team-a-deployment-id-2",
"team_id": deployment_owner_team_id,
},
}
mock_router = MagicMock()
mock_router.get_model_list.return_value = [other_team_deployment_dict]
async def fake_find_unique(*, where):
return SimpleNamespace(
model_dump=lambda: LiteLLM_TeamTable(
team_id=where["team_id"],
members_with_roles=(
[{"user_id": "team-b-admin-user-2", "role": "admin"}]
if where["team_id"] == requester_team_id
else []
),
).model_dump()
)
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
patch("litellm.proxy.proxy_server.llm_router", mock_router),
patch("litellm.proxy.proxy_server.premium_user", True),
patch.object(
ModelManagementAuthChecks,
"can_user_make_model_call",
wraps=ModelManagementAuthChecks.can_user_make_model_call,
) as spy_auth_check,
patch(
"litellm.proxy.management_endpoints.model_management_endpoints.TeamRepository"
) as MockTeamRepo,
):
mock_team_repo_instance = MagicMock()
mock_team_repo_instance.table.find_unique = AsyncMock(
side_effect=fake_find_unique
)
MockTeamRepo.return_value = mock_team_repo_instance
with pytest.raises(HTTPException) as exc_info:
await health_test_model_connection(
request=mock_request,
mode="chat",
litellm_params={
"model": "shared-model-name",
"api_base": "https://swapped-base-2.invalid/v1",
},
model_info={"team_id": requester_team_id},
user_api_key_dict=requester_user_api_key_dict,
)
assert exc_info.value.status_code == 403
passed_model_params = spy_auth_check.call_args.kwargs["model_params"]
assert passed_model_params.model_info.team_id == deployment_owner_team_id
@pytest.mark.asyncio
async def test_test_model_connection_authorized_team_admin_passes_real_auth():
"""
Positive-path companion to the deny tests above. When the caller is a
genuine admin of the team that owns the loaded deployment, the real
(unmocked) auth check must pass and the endpoint must reach the outbound
health probe. Guards against a regression that swaps the auth `team_id`
for something deny-all on the legit path.
"""
from litellm.proxy._types import LiteLLM_TeamTable
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelManagementAuthChecks,
)
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
mock_request = MagicMock()
owner_team_id = "team-owner"
owner_admin_user_id = "team-owner-admin"
owned_deployment_id = "owned-deployment-id"
owner_admin_api_key_dict = UserAPIKeyAuth(
token="owner-admin-token",
user_id=owner_admin_user_id,
team_id=owner_team_id,
user_role=LitellmUserRoles.INTERNAL_USER,
)
mock_prisma_client = MagicMock()
owned_deployment = Deployment(
model_name="owner-model",
litellm_params=LiteLLM_Params(
model="openai/gpt-4o-mini",
api_base="https://owner-real-api.invalid/v1",
api_key="owner-team-api-key",
),
model_info=ModelInfo(id=owned_deployment_id, team_id=owner_team_id),
)
mock_router = MagicMock()
mock_router.get_deployment.return_value = owned_deployment
async def fake_find_unique(*, where):
if where["team_id"] == owner_team_id:
return SimpleNamespace(
model_dump=lambda: LiteLLM_TeamTable(
team_id=owner_team_id,
members_with_roles=[
{"user_id": owner_admin_user_id, "role": "admin"}
],
).model_dump()
)
return None
health_result = {"status": "healthy", "response_time_ms": 50}
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
patch("litellm.proxy.proxy_server.llm_router", mock_router),
patch("litellm.proxy.proxy_server.premium_user", True),
patch.object(
ModelManagementAuthChecks,
"can_user_make_model_call",
wraps=ModelManagementAuthChecks.can_user_make_model_call,
) as spy_auth_check,
patch(
"litellm.proxy.management_endpoints.model_management_endpoints.TeamRepository"
) as MockTeamRepo,
patch(
"litellm.proxy.health_endpoints._health_endpoints.litellm.ahealth_check",
AsyncMock(return_value=health_result),
),
patch(
"litellm.proxy.health_endpoints._health_endpoints.run_with_timeout",
AsyncMock(return_value=health_result),
),
):
mock_team_repo_instance = MagicMock()
mock_team_repo_instance.table.find_unique = AsyncMock(
side_effect=fake_find_unique
)
MockTeamRepo.return_value = mock_team_repo_instance
result = await health_test_model_connection(
request=mock_request,
mode="chat",
litellm_params={"model": "openai/gpt-4o-mini"},
model_info={"id": owned_deployment_id, "team_id": owner_team_id},
user_api_key_dict=owner_admin_api_key_dict,
)
assert result["status"] == "success"
passed_model_params = spy_auth_check.call_args.kwargs["model_params"]
assert passed_model_params.model_info.team_id == owner_team_id
@pytest.mark.asyncio
@pytest.mark.parametrize(
"status,error_message",