mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(key budgets): restrict end user budget lookups to proxy admins
end_user_id was free-form and the lookup was a global find_unique with nothing tying the end user to the key, its team or the caller. Any valid key could read any end user's alias, budget, live spend and reset times, and a null entity_label told the caller whether an id existed at all. /key/budgets made that reachable for every role, since a caller trivially satisfies the key check by inspecting their own key. End users are a proxy-global namespace, LiteLLM_EndUserTable carries no team or organization column, so there is nothing to scope a non-admin against. Gate the parameter on proxy admin and admin viewer, matching /customer/info and /customer/list. The gate is on the parameter, not the route, so a non-admin reading their own key's budgets is unaffected.
This commit is contained in:
parent
1282a6d232
commit
2c608b05a6
2 changed files with 71 additions and 1 deletions
|
|
@ -3753,6 +3753,12 @@ async def info_key_fn(
|
|||
raise handle_exception_on_proxy(e)
|
||||
|
||||
|
||||
_END_USER_BUDGET_READER_ROLES: Final = (
|
||||
LitellmUserRoles.PROXY_ADMIN,
|
||||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
|
||||
)
|
||||
|
||||
|
||||
def _key_not_found_error() -> ProxyException:
|
||||
return ProxyException(
|
||||
message="Key not found in database",
|
||||
|
|
@ -3800,7 +3806,8 @@ async def key_budgets_fn(
|
|||
hash. Defaults to the key in the Authorization header when omitted (`GET /key/budgets`).
|
||||
- end_user_id: str | None (query parameter) - Also report the budgets that would apply to this
|
||||
end user. Omitted end users produce no `end_user` rows, because nothing binds an end user to
|
||||
a key outside a request.
|
||||
a key outside a request. Proxy admins only, since end users are a proxy-global namespace with
|
||||
no key, team or organization scoping to check a caller against.
|
||||
|
||||
Returns:
|
||||
- key: str - The key that was looked up, echoed back as it was passed in
|
||||
|
|
@ -3839,6 +3846,7 @@ async def key_budgets_fn(
|
|||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
user_custom_auth,
|
||||
)
|
||||
|
||||
try:
|
||||
|
|
@ -3847,6 +3855,15 @@ async def key_budgets_fn(
|
|||
"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
|
||||
)
|
||||
|
||||
if end_user_id is not None and user_api_key_dict.user_role not in _END_USER_BUDGET_READER_ROLES:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=(
|
||||
"Only proxy admins can resolve end user budgets, because end users are not scoped to a key, "
|
||||
f"a team or an organization. Your role={user_api_key_dict.user_role}"
|
||||
),
|
||||
)
|
||||
|
||||
key: Final = key_id or user_api_key_dict.api_key
|
||||
if key is None:
|
||||
raise _key_not_found_error()
|
||||
|
|
@ -3886,6 +3903,12 @@ async def key_budgets_fn(
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
general_settings=general_settings,
|
||||
custom_auth_enabled=user_custom_auth is not None,
|
||||
),
|
||||
token_end_user_max_budget=(
|
||||
user_api_key_dict.end_user_max_budget
|
||||
if end_user_id is not None and end_user_id == user_api_key_dict.end_user_id
|
||||
else None
|
||||
),
|
||||
)
|
||||
return KeyBudgetsResponse(key=key, budgets=budgets)
|
||||
|
|
|
|||
|
|
@ -17195,3 +17195,50 @@ async def test_key_budgets_treat_a_non_positive_cap_as_unset(max_budget):
|
|||
assert by_scope["organization"].status == "unlimited"
|
||||
assert by_scope["key_model"].max_budget is None
|
||||
assert by_scope["key_model"].status == "unlimited"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"role",
|
||||
[
|
||||
LitellmUserRoles.INTERNAL_USER.value,
|
||||
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value,
|
||||
LitellmUserRoles.TEAM.value,
|
||||
LitellmUserRoles.CUSTOMER.value,
|
||||
None,
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_budgets_refuse_to_resolve_an_end_user_for_a_non_admin(role):
|
||||
"""End users are proxy-global, so without this any key holder could read any customer's spend."""
|
||||
caller = UserAPIKeyAuth(api_key="sk-caller", user_role=role)
|
||||
with _budgets_route_world(key_row=_budgets_key_row(), caller=caller) as resolver:
|
||||
response = client.get("/key/budgets?end_user_id=someone-elses-customer")
|
||||
|
||||
assert response.status_code == 403
|
||||
assert "end user" in json.dumps(response.json()).lower()
|
||||
resolver.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"role",
|
||||
[LitellmUserRoles.PROXY_ADMIN.value, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_budgets_resolve_an_end_user_for_an_admin(role):
|
||||
"""The gate has to keep admins working, or it has just deleted the parameter."""
|
||||
caller = UserAPIKeyAuth(api_key="sk-admin", user_role=role)
|
||||
with _budgets_route_world(key_row=_budgets_key_row(), caller=caller) as resolver:
|
||||
response = client.get("/key/budgets?end_user_id=a-customer")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert resolver.await_args.kwargs["end_user_id"] == "a-customer"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_budgets_without_an_end_user_stay_open_to_non_admins():
|
||||
"""The gate is on the parameter, not the route: a caller reading their own key must be unaffected."""
|
||||
caller = UserAPIKeyAuth(api_key="sk-caller", user_role=LitellmUserRoles.INTERNAL_USER.value)
|
||||
with _budgets_route_world(key_row=_budgets_key_row(), caller=caller):
|
||||
response = client.get("/key/budgets")
|
||||
|
||||
assert response.status_code == 200
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue