fix(proxy): refuse /health targets outside the caller's scope instead of probing the rest

This commit is contained in:
mateo-berri 2026-09-04 09:40:00 -07:00
parent 471f51cc4d
commit 69a45e81cc
2 changed files with 85 additions and 20 deletions

View file

@ -1127,6 +1127,16 @@ async def health_endpoint(
if allowed_models is None
or _caller_may_probe_deployment(m, allowed_models, llm_router, user_api_key_dict.team_id)
]
targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id)
if restrict_to_allowed_models and targeted_ids is not None and not targeted_ids:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={
"error": f"key not allowed to health-check model_id {model_id}"
if model_id
else f"key not allowed to health-check model {model}"
},
)
if use_background_health_checks:
# The cached background result covers every model. When the
# caller targets a specific model/model_id we have to narrow the
@ -1134,7 +1144,6 @@ async def health_endpoint(
# healthy_count, otherwise an unhealthy "foo" combined with any
# other healthy model would still report healthy_count > 0 and
# the targeted-503 path would never fire.
targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id)
if restrict_to_allowed_models:
allowed_model_ids: Final = {
(m.get("model_info") or {}).get("id")

View file

@ -1517,7 +1517,7 @@ async def test_health_endpoint_filters_model_list_by_user_access():
):
from fastapi import Response
await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict)
await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None)
assert "model_list" in captured, "health_endpoint did not call _perform_health_check_and_save"
returned_names = {m["model_name"] for m in captured["model_list"]}
@ -1580,7 +1580,7 @@ async def test_health_endpoint_keeps_full_model_list_for_all_proxy_models():
):
from fastapi import Response
await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict)
await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None)
returned_names = {m["model_name"] for m in captured["model_list"]}
assert returned_names == {
@ -1648,7 +1648,7 @@ async def test_health_endpoint_resolves_all_team_models_to_team_allowlist():
):
from fastapi import Response
await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict)
await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None)
returned_names = {m["model_name"] for m in captured["model_list"]}
assert returned_names == {"model-b"}, f"all-team-models key should health-check the team's models: {returned_names}"
@ -1769,6 +1769,8 @@ async def test_health_endpoint_expands_access_group_on_live_path():
await health_endpoint(
response=Response(),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"]),
model=None,
model_id=None,
)
assert [m["model_name"] for m in captured["model_list"]] == ["bedrock-nova"]
@ -2149,7 +2151,7 @@ async def test_health_endpoint_blocks_cross_scope_model_id_under_background_cach
cache filter was driven by an unvalidated ID and the global cache
leaked id-b's entry to the caller.
"""
from fastapi import Response
from fastapi import HTTPException, Response
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
@ -2200,21 +2202,18 @@ async def test_health_endpoint_blocks_cross_scope_model_id_under_background_cach
):
# Calling with model="model-b" rather than model_id="id-b" because
# the model_id branch raises 404 when llm_router is None. The bug
# being verified is the same: targeted resolver must drop entries
# not in the caller's scoped model_list. With the fix, the result
# has no leaked endpoints and the targeted-503 path fires.
result = await health_endpoint(
response=response,
user_api_key_dict=user_api_key_dict,
model="model-b",
model_id=None,
)
# being verified is the same: a target outside the caller's scoped
# model_list is refused before the cache is read.
with pytest.raises(HTTPException) as refused:
await health_endpoint(
response=response,
user_api_key_dict=user_api_key_dict,
model="model-b",
model_id=None,
)
leaked_ids = {ep.get("model_id") for ep in result.get("healthy_endpoints", [])}
leaked_ids |= {ep.get("model_id") for ep in result.get("unhealthy_endpoints", [])}
assert "id-b" not in leaked_ids, "background cache leaked an out-of-scope deployment to a scoped caller"
assert result["healthy_count"] == 0
assert response.status_code == 503
assert refused.value.status_code == 403
assert "leaky-internal.test" not in str(refused.value.detail)
@pytest.mark.asyncio
@ -2810,7 +2809,7 @@ async def _live_probed_model_ids(
side_effect=fake_perform,
),
):
await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict)
await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict, model=None, model_id=None)
return {m["model_info"]["id"] for m in captured["model_list"]}
@ -2883,6 +2882,63 @@ async def test_health_endpoint_hides_another_teams_deployment_on_background_cach
assert result["healthy_count"] == 1
@pytest.mark.asyncio
async def test_health_endpoint_refuses_a_targeted_deployment_outside_the_callers_scope_on_live_path():
"""
A scoped key asking for a deployment it may not see must get a 403 and no
probe at all: probing the rest of its scope instead would report another
deployment's health under the requested id and store it as such.
"""
from fastapi import HTTPException, Response
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
fake_perform = AsyncMock()
with (
_proxy_health_globals(_TEAM_MODEL_LIST, _router_for(_TEAM_MODEL_LIST)),
patch( # test-quality-ok: the probe must never run; no injection seam
"litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save",
fake_perform,
),
pytest.raises(HTTPException) as excinfo,
):
await health_endpoint(
response=Response(),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id=None),
model=None,
model_id="id-team-b",
)
assert excinfo.value.status_code == 403
fake_perform.assert_not_awaited()
@pytest.mark.asyncio
async def test_health_endpoint_refuses_a_targeted_deployment_outside_the_callers_scope_on_background_cache_path():
from fastapi import HTTPException, Response
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
with (
_proxy_health_globals(
_TEAM_MODEL_LIST,
_router_for(_TEAM_MODEL_LIST),
use_background_health_checks=True,
health_check_results=_TEAM_CACHED_RESULTS,
),
pytest.raises(HTTPException) as excinfo,
):
await health_endpoint(
response=Response(),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-group"], team_id="team-a"),
model="bedrock-nova_team-b_9f2c",
model_id=None,
)
assert excinfo.value.status_code == 403
@pytest.mark.asyncio
async def test_health_endpoint_hides_team_deployments_from_a_key_with_no_team_on_background_cache_path():
from fastapi import Response