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:
yucheng 2026-09-11 23:25:33 +00:00
parent fcca3c239e
commit 383edfe953
2 changed files with 72 additions and 10 deletions

View file

@ -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."""

View file

@ -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.