diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index 79005bf71b6..45ae5943598 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -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) diff --git a/litellm/proxy/hooks/tag_rate_limits_shared.py b/litellm/proxy/hooks/tag_rate_limits_shared.py index 30735f1a9c5..14e8866b142 100644 --- a/litellm/proxy/hooks/tag_rate_limits_shared.py +++ b/litellm/proxy/hooks/tag_rate_limits_shared.py @@ -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" diff --git a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py index 6fcaf865c5d..e9288566da0 100644 --- a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py @@ -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): """