mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(proxy): unify alias/internal-name tag rate-limit buckets and close caller-forged metadata bypass
Bugbot finding: stamping team_scope on both the team-alias and internal model_name paths didn't unify their Redis keys, since _hash_tag still hashed the caller-visible name by default and that name differs between the two paths. Both index branches now also stamp resolved_group with the team's own public alias, so a team calling its own internal model_name and the same team calling its public alias land on one shared bucket. Veria AI finding: admission resolved the authoritative metadata bucket via get_metadata_variable_name_from_kwargs, which only checks key presence. A caller could forge an empty (or None) litellm_metadata to make admission read no tags/team at all and bypass every configured limit. Admission now reuses the same truthiness-checking resolver already used for success accounting, renamed to reflect both call sites.
This commit is contained in:
Deepanshu
parent
d5dc12fa04
commit
8b3f634fd2
3 changed files with 209 additions and 45 deletions
|
|
@ -27,7 +27,6 @@ from litellm.caching.in_memory_cache import InMemoryCache
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
_get_parent_otel_span_from_kwargs, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching dynamic_rate_limiter_v3's identical import
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
|
||||
|
|
@ -99,7 +98,7 @@ from litellm.proxy.hooks.tag_rate_limits_shared import (
|
|||
policy_fingerprint as _policy_fingerprint,
|
||||
)
|
||||
from litellm.proxy.hooks.tag_rate_limits_shared import (
|
||||
resolve_success_event_metadata_variable_name as _resolve_success_event_metadata_variable_name,
|
||||
resolve_authoritative_metadata_variable_name as _resolve_authoritative_metadata_variable_name,
|
||||
)
|
||||
from litellm.proxy.hooks.tag_rate_limits_shared import (
|
||||
scope_signature as _scope_signature,
|
||||
|
|
@ -150,24 +149,33 @@ class _ConfiguredLimit:
|
|||
# bucket). Otherwise the sorted deployment ids that declared this exact
|
||||
# value -- the bucket is shared among only those deployments.
|
||||
deployment_scope: tuple[str, ...] | None
|
||||
# The team_id this limit was resolved under via `by_team_alias`, or None
|
||||
# when resolved via `by_model_name`. team_public_model_name is only
|
||||
# unique per team, so two teams can publish the identical alias string;
|
||||
# without the team_id folded into the bucket key too, both teams'
|
||||
# identically-named, identically-configured limits would collide on the
|
||||
# same Redis counter despite the index itself correctly scoping the
|
||||
# lookup by (team_id, alias).
|
||||
# The owning team_id, set whenever this limit belongs to a team-owned
|
||||
# deployment -- via `by_team_alias`, or via `by_model_name` for a group
|
||||
# any of whose deployments also declare a `team_public_model_name` (see
|
||||
# `_build_limits_index`). team_public_model_name is only unique per team,
|
||||
# so two teams can publish the identical alias string; without the
|
||||
# team_id folded into the bucket key too, both teams' identically-named,
|
||||
# identically-configured limits would collide on the same Redis counter
|
||||
# despite the index itself correctly scoping the `by_team_alias` lookup
|
||||
# by (team_id, alias).
|
||||
team_scope: str | None = None
|
||||
# The real model_name this limit was found under when `resolve()`'s
|
||||
# direct lookup by the caller-visible model string missed and
|
||||
# `resolve_any()` fell back to resolving via a candidate deployment's
|
||||
# own model_name instead (routing groups, and any other indirection
|
||||
# where Router deliberately keeps the caller-visible name distinct from
|
||||
# every deployment's own model_name). None when resolved directly, in
|
||||
# which case the caller-visible name is already unambiguous and safe to
|
||||
# hash by. Set, this overrides the caller-visible name in the bucket key
|
||||
# so limits from two different underlying model_names sharing one
|
||||
# routing group never collide on one counter.
|
||||
# Overrides the caller-visible name in the bucket key whenever that name
|
||||
# doesn't uniquely identify the bucket. Two independent cases set this:
|
||||
# (1) `resolve()`'s direct lookup missed and `resolve_any()` fell back to
|
||||
# resolving via a candidate deployment's own model_name instead (routing
|
||||
# groups, and any other indirection where Router deliberately keeps the
|
||||
# caller-visible name distinct from every deployment's own model_name) --
|
||||
# here it's stamped with that deployment's own model_name, so limits from
|
||||
# two different underlying model_names sharing one routing group never
|
||||
# collide on one counter; (2) the limit belongs to a team-owned
|
||||
# deployment (`team_scope` is set) -- here it's stamped with the team's
|
||||
# own `team_public_model_name`, the same value regardless of whether this
|
||||
# entry was found via `by_team_alias` (caller used the public alias) or
|
||||
# via `by_model_name` (caller used the deployment's own internal,
|
||||
# auto-generated model_name), so both call shapes land on one shared
|
||||
# bucket instead of splitting a team's usage across two counters. None
|
||||
# when resolved directly and not team-owned, in which case the
|
||||
# caller-visible name is already unambiguous and safe to hash by.
|
||||
resolved_group: str | None = None
|
||||
|
||||
|
||||
|
|
@ -186,8 +194,8 @@ def _model_name_of(deployment: Mapping[str, object]) -> str:
|
|||
|
||||
def _extract_team_id(request_kwargs: Mapping[str, object], metadata_variable_name: str) -> str | None:
|
||||
"""Reads `user_api_key_team_id` from only the one field
|
||||
`get_metadata_variable_name_from_kwargs` names as authoritative for this
|
||||
request -- never falling back to the other field, since
|
||||
`_resolve_authoritative_metadata_variable_name` names as authoritative for
|
||||
this request -- never falling back to the other field, since
|
||||
`litellm_pre_call_utils.py` writes the real, server-authenticated value
|
||||
into that one field alone and leaves the other exactly as the caller
|
||||
sent it. An OR-fallback across both would let a caller's own
|
||||
|
|
@ -462,12 +470,16 @@ def _build_limits_index(model_list: Sequence[Mapping[str, object]]) -> _LimitsIn
|
|||
# model_name, not only its team_public_model_name alias
|
||||
# (litellm auto-generates a name unique per (team_id, uuid),
|
||||
# so every deployment in this group shares one team_id when
|
||||
# any does) -- stamping the identical team_scope here as the
|
||||
# alias entry below gets keeps both paths resolving to the
|
||||
# same bucket, so a caller can't split its usage across two
|
||||
# any does). Stamping team_scope alone is not enough to unify
|
||||
# this with the alias entry below: `_hash_tag` still hashes
|
||||
# the caller-visible name by default, and that name differs
|
||||
# between the two paths (the internal model_name here vs. the
|
||||
# public alias below). Also stamping `resolved_group` with the
|
||||
# team's own alias forces both paths to hash under the
|
||||
# identical name, so a caller can't split its usage across two
|
||||
# independent counters just by alternating which name it calls.
|
||||
tuple(replace(limit, team_scope=team_scope) for limit in configured)
|
||||
if (team_scope := next((key[0] for dep in group if (key := _team_alias_key(dep))), None)) is not None
|
||||
tuple(replace(limit, team_scope=team_key[0], resolved_group=team_key[1]) for limit in configured)
|
||||
if (team_key := next((key for dep in group if (key := _team_alias_key(dep))), None)) is not None
|
||||
else configured
|
||||
)
|
||||
for model_name, deployment_group in groupby(sorted_by_model_name, key=_model_name_of)
|
||||
|
|
@ -487,7 +499,13 @@ def _build_limits_index(model_list: Sequence[Mapping[str, object]]) -> _LimitsIn
|
|||
for aliased_group in (tuple(dep for _key, dep in alias_group),)
|
||||
if (
|
||||
alias_configured := tuple(
|
||||
replace(limit, team_scope=alias_key[0])
|
||||
# resolved_group is already the caller-visible name on
|
||||
# this path (the caller reached this group by dialing the
|
||||
# alias directly), but stamping it explicitly keeps both
|
||||
# index branches symmetric and independent of whatever
|
||||
# value the caller happens to pass as `model_group` into
|
||||
# `_hash_tag`.
|
||||
replace(limit, team_scope=alias_key[0], resolved_group=alias_key[1])
|
||||
for unit in _LIMIT_UNITS
|
||||
for limit in _build_group_limits(aliased_group, unit)
|
||||
)
|
||||
|
|
@ -711,12 +729,15 @@ def _scope_suffix(deployment_scope: tuple[str, ...] | None) -> str:
|
|||
|
||||
|
||||
def _hash_tag(model_group: str, configured: _ConfiguredLimit, tag_value: str, key_hash: str | None) -> str:
|
||||
# resolved_group overrides the caller-visible model_group when this
|
||||
# limit was found via resolve_any()'s per-deployment fallback (routing
|
||||
# groups): the caller-visible name is ambiguous there (shared by every
|
||||
# member model_name), so hashing by it would collide two different
|
||||
# underlying model_names' identically-named limits onto one counter.
|
||||
# See _ConfiguredLimit.resolved_group.
|
||||
# resolved_group overrides the caller-visible model_group in two cases:
|
||||
# resolve_any()'s per-deployment fallback (routing groups), where the
|
||||
# caller-visible name is ambiguous (shared by every member model_name), so
|
||||
# hashing by it would collide two different underlying model_names'
|
||||
# identically-named limits onto one counter; and a team-owned deployment,
|
||||
# where the caller-visible name differs depending on whether the caller
|
||||
# dialed the team's public alias or the deployment's own internal
|
||||
# model_name, so hashing by it would split one team's usage across two
|
||||
# counters. See _ConfiguredLimit.resolved_group.
|
||||
effective_model_group: Final = configured.resolved_group if configured.resolved_group is not None else model_group
|
||||
scope: Final = _scope_suffix(configured.deployment_scope)
|
||||
# team_scope disambiguates two teams that publish the identical
|
||||
|
|
@ -1139,7 +1160,14 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
|
|||
|
||||
resolved_request_kwargs: Final = request_kwargs or _EMPTY_MAPPING
|
||||
stale_request_keys: Final = await self._release_stale_hop_reservations(resolved_request_kwargs)
|
||||
metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(resolved_request_kwargs)
|
||||
# Not `get_metadata_variable_name_from_kwargs` (naive key-presence
|
||||
# check): a caller can forge an empty (or `None`) `litellm_metadata`
|
||||
# on an ordinary request to make that check pick it over the real,
|
||||
# populated `metadata` the proxy wrote authenticated team/tag identity
|
||||
# into, seeing no tags at all and admitting past every configured
|
||||
# limit. See `_resolve_authoritative_metadata_variable_name`'s own
|
||||
# docstring.
|
||||
metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(resolved_request_kwargs)
|
||||
team_id: Final = _extract_team_id(resolved_request_kwargs, metadata_variable_name)
|
||||
# Built from the full routing-group membership, not `healthy_deployments`
|
||||
# (Router's own cooldown-filtered list for this hop): a member that's
|
||||
|
|
@ -1499,12 +1527,12 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
|
|||
# check): at this point `kwargs` is `model_call_details`, which
|
||||
# carries `litellm_metadata` present-but-`None` alongside the
|
||||
# real, populated `metadata` for a standard request -- see
|
||||
# `_resolve_success_event_metadata_variable_name`'s own docstring.
|
||||
# `_resolve_authoritative_metadata_variable_name`'s own docstring.
|
||||
litellm_params_raw: Final = kwargs.get("litellm_params")
|
||||
litellm_params_for_metadata: Final = (
|
||||
litellm_params_raw if isinstance(litellm_params_raw, Mapping) else kwargs
|
||||
)
|
||||
metadata_variable_name: Final = _resolve_success_event_metadata_variable_name(litellm_params_for_metadata)
|
||||
metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(litellm_params_for_metadata)
|
||||
key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name)
|
||||
try:
|
||||
await self.internal_usage_cache.dual_cache.async_delete_cache(
|
||||
|
|
@ -1644,7 +1672,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
|
|||
# top-level here, only nested under kwargs["litellm_params"] (see
|
||||
# Logging.update_environment_variables).
|
||||
litellm_params_for_metadata: Final = kwargs.get("litellm_params") or kwargs
|
||||
metadata_variable_name: Final = _resolve_success_event_metadata_variable_name(litellm_params_for_metadata)
|
||||
metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(litellm_params_for_metadata)
|
||||
team_id: Final = _extract_team_id(litellm_params_for_metadata, metadata_variable_name)
|
||||
key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name)
|
||||
key_alias: Final = _extract_key_alias(litellm_params_for_metadata, metadata_variable_name)
|
||||
|
|
|
|||
|
|
@ -221,15 +221,18 @@ def entry_applies(entry: TagRateLimitEntry, tags: Sequence[str], key_alias: str
|
|||
return key_alias in entry.apply_to_key_alias
|
||||
|
||||
|
||||
def resolve_success_event_metadata_variable_name(
|
||||
litellm_params_for_metadata: Mapping[str, object],
|
||||
def resolve_authoritative_metadata_variable_name(
|
||||
metadata_source: Mapping[str, object],
|
||||
) -> Literal["metadata", "litellm_metadata"]:
|
||||
"""`get_metadata_variable_name_from_kwargs` only checks key presence, which
|
||||
misresolves at `async_log_success_event` time: `kwargs["litellm_params"]`
|
||||
always carries a `litellm_metadata` key (typically `None`) alongside the
|
||||
real, populated `metadata` dict for a standard (non
|
||||
LITELLM_METADATA_ROUTES) request, so the key-presence check always picks
|
||||
`litellm_metadata` there and silently reads no tags/identity at all.
|
||||
misresolves both at admission time and at `async_log_success_event` time:
|
||||
a caller can forge an empty (or merely present-but-`None`) `litellm_metadata`
|
||||
on an ordinary request -- `kwargs["litellm_params"]` also always carries a
|
||||
`litellm_metadata` key (typically `None`) alongside the real, populated
|
||||
`metadata` dict for a standard (non LITELLM_METADATA_ROUTES) request -- and
|
||||
the key-presence check always picks `litellm_metadata` in both cases,
|
||||
silently reading no tags/identity at all and admitting the request against
|
||||
every configured limit.
|
||||
|
||||
A plain truthiness check on `litellm_metadata` isn't enough either: a
|
||||
caller can populate it with unrelated, non-empty content on a route where
|
||||
|
|
@ -241,7 +244,7 @@ def resolve_success_event_metadata_variable_name(
|
|||
there -- so requiring that marker's presence, not mere truthiness, only
|
||||
ever prefers `litellm_metadata` when it is genuinely the field the proxy
|
||||
wrote identity/tags into."""
|
||||
litellm_metadata: Final = litellm_params_for_metadata.get("litellm_metadata")
|
||||
litellm_metadata: Final = metadata_source.get("litellm_metadata")
|
||||
if isinstance(litellm_metadata, Mapping) and "user_api_key_auth" in litellm_metadata:
|
||||
return "litellm_metadata"
|
||||
return "metadata"
|
||||
|
|
|
|||
|
|
@ -214,6 +214,66 @@ def test_extract_team_id_ignores_a_forged_value_in_the_non_authoritative_field()
|
|||
assert _extract_team_id(request_kwargs, "litellm_metadata") == "real-team"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_deployments_ignores_a_forged_empty_litellm_metadata_key(time_controller):
|
||||
"""
|
||||
Veria AI finding: get_metadata_variable_name_from_kwargs picks
|
||||
"litellm_metadata" whenever that key is merely present, regardless of its
|
||||
value. add_litellm_data_to_request writes real, authenticated team/tag
|
||||
identity into "metadata" for an ordinary (non LITELLM_METADATA_ROUTES)
|
||||
request, but leaves any caller-supplied "litellm_metadata" sitting
|
||||
alongside it untouched -- so a caller adding an empty "litellm_metadata"
|
||||
to a chat completions request made admission read no tags/team at all,
|
||||
sailing past every configured limit.
|
||||
"""
|
||||
limiter = _make_limiter(time_controller)
|
||||
deployment = _deployment(
|
||||
"grp",
|
||||
"dep-1",
|
||||
{"request_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]}},
|
||||
)
|
||||
router = litellm.Router(model_list=[deployment])
|
||||
limiter.update_variables(llm_router=router)
|
||||
healthy = router.model_list
|
||||
|
||||
request_kwargs = {"metadata": {"tags": ["end_user_id:u1"]}, "litellm_metadata": {}}
|
||||
await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
with pytest.raises(ProxyRateLimitError):
|
||||
await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_deployments_reads_metadata_when_litellm_metadata_is_present_but_none(time_controller):
|
||||
"""
|
||||
Same misresolution, naturally occurring rather than attacker-forged: a
|
||||
plain chat completion's kwargs carries a "litellm_metadata" key that is
|
||||
always present but set to None, alongside the real, populated "metadata"
|
||||
dict (see resolve_authoritative_metadata_variable_name's own docstring).
|
||||
"""
|
||||
limiter = _make_limiter(time_controller)
|
||||
deployment = _deployment(
|
||||
"grp",
|
||||
"dep-1",
|
||||
{"request_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]}},
|
||||
)
|
||||
router = litellm.Router(model_list=[deployment])
|
||||
limiter.update_variables(llm_router=router)
|
||||
healthy = router.model_list
|
||||
|
||||
request_kwargs = {"metadata": {"tags": ["end_user_id:u1"]}, "litellm_metadata": None}
|
||||
await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
with pytest.raises(ProxyRateLimitError):
|
||||
await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TagRateLimitEntry -- limit validation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -3519,6 +3579,33 @@ def test_build_limits_index_is_also_keyed_by_team_public_model_name():
|
|||
assert by_alias[0].team_scope == "team-1"
|
||||
|
||||
|
||||
def test_build_limits_index_computes_identical_bucket_key_for_alias_and_internal_model_name():
|
||||
"""
|
||||
Bugbot finding: matching team_scope alone does not unify the two paths'
|
||||
buckets, because _hash_tag hashes the caller-visible model_group by
|
||||
default, and that name is `real-model-name` on the by_model_name path but
|
||||
`team-alias-name` on the by_team_alias path. Both entries must also share
|
||||
the identical resolved_group (the team's own alias) so _hash_tag hashes
|
||||
both under one name -- otherwise a team calling its own internal
|
||||
model_name lands on a different Redis counter than the same team calling
|
||||
its public alias, doubling its effective quota.
|
||||
"""
|
||||
deployment = _deployment(
|
||||
"real-model-name",
|
||||
"dep-1",
|
||||
{"token_limits": {"limits": [{"name": "daily", "limit": 500, "period_seconds": 86400}]}},
|
||||
)
|
||||
deployment["model_info"]["team_id"] = "team-1"
|
||||
deployment["model_info"]["team_public_model_name"] = "team-alias-name"
|
||||
index = _build_limits_index([deployment])
|
||||
by_name = index.resolve("real-model-name", team_id=None)[0]
|
||||
by_alias = index.resolve("team-alias-name", team_id="team-1")[0]
|
||||
|
||||
key_via_internal_name = _bucket_key("real-model-name", by_name, tag_value="u1", bucket_id=0)
|
||||
key_via_alias = _bucket_key("team-alias-name", by_alias, tag_value="u1", bucket_id=0)
|
||||
assert key_via_internal_name == key_via_alias
|
||||
|
||||
|
||||
def test_build_limits_index_preserves_key_ttl_seconds_and_max_in_memory_cache_size():
|
||||
"""
|
||||
Regression test: _configured_limit_for_signature used to reconstruct a
|
||||
|
|
@ -3707,6 +3794,52 @@ async def test_filter_deployments_enforces_limit_when_called_with_team_alias(tim
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_deployments_shares_quota_across_alias_and_internal_model_name(time_controller):
|
||||
"""
|
||||
Bugbot finding on a prior fix for this same drift: stamping the identical
|
||||
team_scope on both paths was not enough, since _hash_tag still hashes the
|
||||
caller-visible name (the alias here, the deployment's own internal
|
||||
model_name there) by default. A team calling through its public alias and
|
||||
the identical team calling through the deployment's own internal
|
||||
model_name must draw from the same bucket, or the team gets one quota per
|
||||
name it happens to call with -- doubling its real limit.
|
||||
"""
|
||||
limiter = _make_limiter(time_controller)
|
||||
deployment = _deployment(
|
||||
"real-model-name",
|
||||
"dep-1",
|
||||
{
|
||||
"request_limits": {
|
||||
"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]
|
||||
}
|
||||
},
|
||||
)
|
||||
deployment["model_info"]["team_id"] = "team-1"
|
||||
deployment["model_info"]["team_public_model_name"] = "team-alias-name"
|
||||
router = litellm.Router(model_list=[deployment])
|
||||
limiter.update_variables(llm_router=router)
|
||||
healthy = router.model_list
|
||||
|
||||
# First call exhausts the limit of 1 via the team's public alias.
|
||||
await limiter.async_filter_deployments(
|
||||
model="team-alias-name",
|
||||
healthy_deployments=healthy,
|
||||
messages=None,
|
||||
request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1"}},
|
||||
)
|
||||
# Router.should_include_deployment also lets the same team reach this
|
||||
# deployment by its own internal model_name; that call must be rejected
|
||||
# against the alias call's own bucket, not admitted into a fresh one.
|
||||
with pytest.raises(ProxyRateLimitError):
|
||||
await limiter.async_filter_deployments(
|
||||
model="real-model-name",
|
||||
healthy_deployments=healthy,
|
||||
messages=None,
|
||||
request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key_team_id": "team-1"}},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_deployments_does_not_cross_team_alias_boundary(time_controller):
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue