mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
fix(proxy): keep counting the failed sign-ins Redis missed once it answers again
A guess is recorded in exactly one place, Redis or this worker's own store when Redis refused it, so the count is the sum of the two. Redis is read through async_batch_get_counts, which raises on failure, instead of async_get_cache, which swallows it into None and read as an empty counter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
fcca3c239e
commit
383edfe953
2 changed files with 72 additions and 10 deletions
|
|
@ -199,12 +199,18 @@ class LoginThrottle:
|
|||
return await self._outcome(work(redis_cache))
|
||||
|
||||
async def _failures(self, store: DualCache, key: str) -> int:
|
||||
"""The shared count while Redis answers, so every worker sees one budget; this worker's
|
||||
own count only while it does not, so an outage degrades to per-worker accounting."""
|
||||
shared: Final = await self._shared(lambda redis_cache: redis_cache.async_get_cache(key))
|
||||
if shared is not _UNAVAILABLE:
|
||||
return _as_count(shared)
|
||||
return _as_count(await self._outcome(store.async_get_cache(key=key)))
|
||||
"""The shared count plus this worker's own.
|
||||
|
||||
A failure is written to exactly one of the two: Redis, or this worker's store when Redis
|
||||
refused it. So the local store is empty while Redis is healthy, and once Redis answers
|
||||
again the guesses it missed still count. Read through ``async_batch_get_counts`` because
|
||||
``async_get_cache`` turns a failed GET into ``None``, which would pass as an empty counter.
|
||||
"""
|
||||
local: Final = _as_count(await self._outcome(store.async_get_cache(key=key)))
|
||||
shared: Final = await self._shared(lambda redis_cache: redis_cache.async_batch_get_counts([key]))
|
||||
if not isinstance(shared, tuple):
|
||||
return local
|
||||
return _as_count(shared[0]) + local
|
||||
|
||||
async def _remaining_window(self, key: str) -> int:
|
||||
"""Seconds until this counter expires.
|
||||
|
|
@ -269,9 +275,11 @@ class LoginThrottle:
|
|||
shared: Final = await self._shared(
|
||||
lambda redis_cache: redis_cache.async_increment_with_floor(key, 1, self.window_seconds)
|
||||
)
|
||||
if shared is not _UNAVAILABLE:
|
||||
return _as_count(shared)
|
||||
return _as_count(await self._outcome(store.async_increment_cache(key=key, value=1, ttl=self.window_seconds)))
|
||||
if shared is _UNAVAILABLE:
|
||||
return _as_count(
|
||||
await self._outcome(store.async_increment_cache(key=key, value=1, ttl=self.window_seconds))
|
||||
)
|
||||
return _as_count(shared) + _as_count(await self._outcome(store.async_get_cache(key=key)))
|
||||
|
||||
async def record_failure(self, username: str) -> FailureCounts:
|
||||
"""Count one rejected credential guess against this username and against this source."""
|
||||
|
|
|
|||
|
|
@ -1156,6 +1156,9 @@ class _FakeRedis:
|
|||
async def async_get_cache(self, key, **kwargs):
|
||||
return self.values.get(key)
|
||||
|
||||
async def async_batch_get_counts(self, key_list):
|
||||
return tuple(self.values.get(key) for key in key_list)
|
||||
|
||||
async def async_increment_with_floor(self, key, value, ttl):
|
||||
self.values[key] = self.values.get(key, 0) + value
|
||||
self.ttls.setdefault(key, ttl)
|
||||
|
|
@ -1204,9 +1207,16 @@ async def test_counters_are_written_with_their_expiry_and_re_armed_if_stripped(m
|
|||
|
||||
|
||||
class _DownRedis(_FakeRedis):
|
||||
"""Redis whose every call raises, as during an outage or an open circuit breaker."""
|
||||
"""Redis whose every call fails, as during an outage or an open circuit breaker.
|
||||
|
||||
`async_get_cache` returns None rather than raising, as the real one does: it swallows the
|
||||
error, so a failed GET is indistinguishable from an empty key to anyone reading through it.
|
||||
"""
|
||||
|
||||
async def async_get_cache(self, key, **kwargs):
|
||||
return None
|
||||
|
||||
async def async_batch_get_counts(self, key_list):
|
||||
raise ConnectionError("redis is down")
|
||||
|
||||
async def async_increment_with_floor(self, key, value, ttl):
|
||||
|
|
@ -1280,6 +1290,50 @@ async def test_a_redis_outage_falls_back_to_this_workers_own_counter(monkeypatch
|
|||
assert blocked.value.headers.get("Retry-After") == "900"
|
||||
|
||||
|
||||
class _WriteRefusingRedis(_FakeRedis):
|
||||
"""Redis that answers reads but raises on writes until `recover()` is called."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.writable = False
|
||||
|
||||
def recover(self):
|
||||
self.writable = True
|
||||
|
||||
async def async_increment_with_floor(self, key, value, ttl):
|
||||
if not self.writable:
|
||||
raise ConnectionError("redis write failed")
|
||||
return await super().async_increment_with_floor(key, value, ttl)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failures_redis_refused_still_count_once_redis_recovers(monkeypatch):
|
||||
"""Regression: a guess Redis could not record must not be forgotten when Redis comes back.
|
||||
|
||||
Such a guess lands in this worker's own store. Reading only Redis afterwards handed the
|
||||
attacker that guess again, so the budget was the limit plus however many writes failed.
|
||||
"""
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
monkeypatch.setenv("UI_USERNAME", "admin")
|
||||
monkeypatch.setenv("UI_PASSWORD", "right")
|
||||
redis = _WriteRefusingRedis()
|
||||
throttle = _throttle(max_attempts=2, redis_cache=redis)
|
||||
|
||||
with pytest.raises(ProxyException, match="Invalid credentials"):
|
||||
await _guess(throttle)
|
||||
assert not redis.values, "the refused write must not have reached Redis"
|
||||
|
||||
redis.recover()
|
||||
with pytest.raises(ProxyException, match="Invalid credentials"):
|
||||
await _guess(throttle)
|
||||
assert [v for k, v in redis.values.items() if ":user:" in k] == [1], "only the recorded guess is in Redis"
|
||||
|
||||
with pytest.raises(ProxyException) as blocked:
|
||||
await _guess(throttle)
|
||||
assert blocked.value.code == "429", "the guess Redis missed and the one it took must add up to the limit"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_counters_do_not_share_the_key_authentication_cache(monkeypatch):
|
||||
"""Regression: throttle entries must not evict cached credentials.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue