mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(rate-limiting): isolate this hook's in-memory cache from the shared proxy cache
internal_usage_cache is the same DualCache instance the proxy's key/team parallel-request limiter uses for its own authentication- bound counters, and its default InMemoryCache evicts at 200 items. Without isolation, a caller flooding this hook's own caller-controlled tag buckets past that ceiling could evict an unrelated, authentication- bound counter and let some other caller exceed a limit nothing here configured. Give this hook a dedicated in-memory layer while still sharing the real Redis connection when one is configured, so cross- instance correctness is unaffected.
This commit is contained in:
parent
c62d79f7df
commit
e6c504645b
2 changed files with 60 additions and 1 deletions
|
|
@ -616,7 +616,16 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer
|
||||||
internal_usage_cache: DualCache,
|
internal_usage_cache: DualCache,
|
||||||
time_provider: Callable[[], datetime] | None = None,
|
time_provider: Callable[[], datetime] | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.internal_usage_cache = InternalUsageCache(dual_cache=internal_usage_cache)
|
# A dedicated in-memory layer, not the proxy-wide `internal_usage_cache`
|
||||||
|
# passed in: that instance is shared with the key/team parallel-request
|
||||||
|
# limiter's own authentication-bound counters, and its default
|
||||||
|
# InMemoryCache evicts at 200 items. Without this isolation, a caller
|
||||||
|
# flooding this hook's own caller-controlled tag buckets past that
|
||||||
|
# ceiling could evict an unrelated, authentication-bound counter and
|
||||||
|
# exceed a limit nothing here configured. The real Redis connection
|
||||||
|
# (if any) is still shared, so cross-instance correctness is unaffected.
|
||||||
|
isolated_dual_cache: Final = DualCache(redis_cache=internal_usage_cache.redis_cache)
|
||||||
|
self.internal_usage_cache = InternalUsageCache(dual_cache=isolated_dual_cache)
|
||||||
self._v3 = _PROXY_MaxParallelRequestsHandler_v3(self.internal_usage_cache, time_provider=time_provider)
|
self._v3 = _PROXY_MaxParallelRequestsHandler_v3(self.internal_usage_cache, time_provider=time_provider)
|
||||||
self._time_provider = time_provider or datetime.now
|
self._time_provider = time_provider or datetime.now
|
||||||
self._index = _TagRateLimitIndex(time_provider=self._time_provider)
|
self._index = _TagRateLimitIndex(time_provider=self._time_provider)
|
||||||
|
|
|
||||||
|
|
@ -2131,3 +2131,53 @@ async def test_concurrency_scope_by_key_hash_gives_independent_reservations_per_
|
||||||
messages=None,
|
messages=None,
|
||||||
request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyB"}},
|
request_kwargs={"metadata": {"tags": ["end_user_id:u1"], "user_api_key": "keyB"}},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# in-memory cache isolation -- caller-controlled tag buckets must never evict
|
||||||
|
# the shared cache's other, authentication-bound counters
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_flooding_tag_buckets_does_not_evict_the_shared_cache_authentication_bound_counter(
|
||||||
|
time_controller,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
The proxy-wide internal_usage_cache passed into this limiter is also
|
||||||
|
used by the key/team parallel-request limiter for its own,
|
||||||
|
authentication-bound counters, and its default InMemoryCache evicts at
|
||||||
|
200 items. Without a dedicated in-memory layer for this hook's own
|
||||||
|
caller-controlled tag buckets, an attacker sending 200+ distinct tag
|
||||||
|
values could evict an unrelated authentication-bound counter and let
|
||||||
|
some other caller exceed a limit nothing here configured.
|
||||||
|
"""
|
||||||
|
shared_cache = DualCache()
|
||||||
|
await shared_cache.async_set_cache(key="authentication_bound_counter", value="do-not-evict")
|
||||||
|
|
||||||
|
limiter = _PROXY_TagRateLimiter(internal_usage_cache=shared_cache, time_provider=time_controller.now)
|
||||||
|
router = litellm.Router(
|
||||||
|
model_list=[
|
||||||
|
_deployment(
|
||||||
|
"grp",
|
||||||
|
"dep-1",
|
||||||
|
{
|
||||||
|
"request_limits": {
|
||||||
|
"limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 1000, "period_seconds": 60}]
|
||||||
|
}
|
||||||
|
},
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
limiter.update_variables(llm_router=router)
|
||||||
|
healthy = router.model_list
|
||||||
|
|
||||||
|
for i in range(250):
|
||||||
|
await limiter.async_filter_deployments(
|
||||||
|
model="grp",
|
||||||
|
healthy_deployments=healthy,
|
||||||
|
messages=None,
|
||||||
|
request_kwargs={"metadata": {"tags": [f"end_user_id:flood-{i}"]}},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert await shared_cache.async_get_cache(key="authentication_bound_counter") == "do-not-evict"
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue