scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-04 07:27:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 299

refactor(redis-store): move login throttling behind a store

Browse source 
Co-authored-by: Yujong Lee <yujong@berri.ai>
This commit is contained in:
Devin AI 2026-10-30 16:15:27 +00:00
parent b8e43df44b
commit e0c2ec6522
3 changed files with 178 additions and 169 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

126
litellm/proxy/auth/login_throttle.py
View file

@ -9,27 +9,24 @@ Recovery is the master key over the API, which never passes through here, or wai
from __future__ import annotations
import asyncio
import hashlib
import ipaddress
import math
import time
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import timedelta
from functools import cache
from typing import Final, Literal, NamedTuple, Protocol, TypeAlias
from typing import Final, Literal, Protocol, TypeAlias
from fastapi import Request, status
from pydantic import TypeAdapter, ValidationError
from redis.exceptions import RedisError
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError
from litellm.caching.redis_cache import RedisCache
from litellm.constants import (
EMPTY_MAPPING,
LOGIN_THROTTLE_CACHE_KEY_PREFIX,
LOGIN_THROTTLE_MAX_TRACKED_BLOCKS,
LOGIN_THROTTLE_MAX_TRACKED_COUNTERS,
LOGIN_THROTTLE_NOT_BLOCKED,
@ -37,6 +34,14 @@ from litellm.constants import (
)
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy.auth.network import TrustedProxyConfig, resolve_client_ip
from litellm.redis.stores.login_throttle import (
BlockTtls,
LoginKeys,
LoginThrottleStore,
TransportLoginThrottleStore,
login_keys,
)
from litellm.redis.transport.contract import RedisFailure
from litellm.secret_managers.main import get_secret_bool
_LOGIN_THROTTLE_TARGET: Final = "login_throttle"
@ -54,15 +59,12 @@ WINDOW_KEY: Final = "failed_login_window_seconds"
BLOCK_KEY: Final = "failed_login_block_seconds"
TRUSTED_PROXY_RANGES_KEY: Final = "trusted_proxy_ranges"
_REDIS_FAILURES: Final = (RedisError, RedisCircuitBreakerOpenError, OSError, asyncio.TimeoutError)
_LOCAL_BLOCK_EXPIRY: Final = TypeAdapter[float | None](float | None)
_SOURCE_LIMIT_OVERRIDES: Final = TypeAdapter[Mapping[str, object]](Mapping[str, object])
_RANGE_ENTRIES: Final = TypeAdapter[tuple[object, ...]](tuple[object, ...])
Scope: TypeAlias = Literal["user", "source"]
_BlockTtls: TypeAlias = tuple[int, int]
_LUA_BLOCK_TTLS: Final = TypeAdapter[_BlockTtls](_BlockTtls)
_Network: TypeAlias = ipaddress.IPv4Network | ipaddress.IPv6Network
@ -78,25 +80,6 @@ class LocalStore(Protocol):
def delete_cache(self, key: str) -> None: ...
# KEYS: pair counter, pair block, source counter, source block (one cluster slot via the source hash tag)
# ARGV: pair limit, source limit (0 = source scope off), window seconds, block seconds
# Both scripts return {pair block TTL, source block TTL}; 0 or below means not blocked
_BLOCK_TTLS_LUA: Final = "return {redis.call('TTL', KEYS[2]), redis.call('TTL', KEYS[4])}"
_RECORD_FAILURE_LUA: Final = (
"local function bump(count_key, block_key, limit) "
"local blocked = redis.call('TTL', block_key) "
"if blocked > 0 then return blocked end "
"local count = redis.call('INCR', count_key) "
"if redis.call('TTL', count_key) < 0 then redis.call('EXPIRE', count_key, ARGV[3]) end "
"if count > limit then redis.call('SET', block_key, '1', 'EX', ARGV[4]) return tonumber(ARGV[4]) end "
"return 0 end "
"local user_block = bump(KEYS[1], KEYS[2], tonumber(ARGV[1])) "
"local source_block = 0 "
"if tonumber(ARGV[2]) > 0 and user_block == 0 then "
"source_block = bump(KEYS[3], KEYS[4], tonumber(ARGV[2])) end "
"return {user_block, source_block}"
)
_COUNTERS: Final = InMemoryCache(
max_size_in_memory=LOGIN_THROTTLE_MAX_TRACKED_COUNTERS, default_ttl=DEFAULT_FAILED_LOGIN_WINDOW_SECONDS
)
@ -253,13 +236,6 @@ def source_group(client_ip: str) -> str:
return str(address)
class _Keys(NamedTuple):
pair_counter: str
pair_block: str
source_counter: str
source_block: str
@dataclass(frozen=True, slots=True)
class Block:
scope: Scope
@ -312,15 +288,11 @@ class LoginThrottle:
enabled=not exempt and not _rate_limit_disabled(),
)
def _keys(self, username: str) -> _Keys:
group: Final = source_group(self.client_ip)
user: Final = hashlib.sha256(username.casefold().encode("utf-8")).hexdigest()
return _Keys(
pair_counter=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:user:{user}",
pair_block=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:block:user:{user}",
source_counter=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:source",
source_block=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:block:source",
)
def _keys(self, username: str) -> LoginKeys:
return login_keys(source_group(self.client_ip), username)
def _store(self) -> LoginThrottleStore | None:
return None if self.redis_cache is None else TransportLoginThrottleStore(self.redis_cache.transport())
async def attempt(self, username: str) -> LoginAttempt:
"""Refuses a blocked key before any credential is looked at; otherwise hands back the attempt to settle."""
@ -338,7 +310,7 @@ class LoginThrottle:
)
raise self.refused(block.retry_after)
async def _active_block(self, keys: _Keys) -> Block | None:
async def _active_block(self, keys: LoginKeys) -> Block | None:
local: Final = self._local_block_ttls(keys)
shared: Final = await self._shared_block_ttls(keys)
user_ttl: Final = max(local[0], shared[0])
@ -349,19 +321,15 @@ class LoginThrottle:
return Block(scope="user", retry_after=user_ttl)
return None
@with_service_target("login_throttle")
async def _shared_block_ttls(self, keys: _Keys) -> _BlockTtls:
if self.redis_cache is None:
return LOGIN_THROTTLE_NOT_BLOCKED
try:
return _LUA_BLOCK_TTLS.validate_python(
await self.redis_cache.async_register_script(_BLOCK_TTLS_LUA)(keys, ())
)
except _REDIS_FAILURES as err:
self._warn_redis(err)
@with_service_target(_LOGIN_THROTTLE_TARGET)
async def _shared_block_ttls(self, keys: LoginKeys) -> BlockTtls:
store: Final = self._store()
if store is None:
return LOGIN_THROTTLE_NOT_BLOCKED
blocks: Final = await store.blocks(keys)
return LOGIN_THROTTLE_NOT_BLOCKED if isinstance(blocks, RedisFailure) else blocks
def _local_block_ttls(self, keys: _Keys) -> _BlockTtls:
def _local_block_ttls(self, keys: LoginKeys) -> BlockTtls:
return self._local_block_ttl(keys.pair_block), self._local_block_ttl(keys.source_block)
def _local_block_ttl(self, block_key: str) -> int:
@ -370,19 +338,24 @@ class LoginThrottle:
return 0
return max(math.ceil(expires_at - time.time()), 0)
@with_service_target("login_throttle")
async def record_failure(self, username: str) -> _BlockTtls:
@with_service_target(_LOGIN_THROTTLE_TARGET)
async def record_failure(self, username: str) -> BlockTtls:
keys: Final = self._keys(username)
source_limit: Final = self.source_limit or 0
if self.redis_cache is not None:
try:
return _LUA_BLOCK_TTLS.validate_python(
await self.redis_cache.async_register_script(_RECORD_FAILURE_LUA)(
keys, (self.user_limit, source_limit, self.window_seconds, self.block_seconds)
)
)
except _REDIS_FAILURES as err:
self._warn_redis(err)
store: Final = self._store()
shared: Final = (
None
if store is None
else await store.record_failure(
keys,
user_limit=self.user_limit,
source_limit=source_limit,
window=timedelta(seconds=self.window_seconds),
block=timedelta(seconds=self.block_seconds),
)
)
if shared is not None and not isinstance(shared, RedisFailure):
return shared
user_block: Final = self._local_bump(keys.pair_counter, keys.pair_block, self.user_limit)
if source_limit == 0 or user_block > 0:
return user_block, 0
@ -400,20 +373,11 @@ class LoginThrottle:
@with_service_target(_LOGIN_THROTTLE_TARGET)
async def clear_pair(self, username: str) -> None:
pair_counter: Final = self._keys(username).pair_counter
if self.redis_cache is not None:
try:
await self.redis_cache.async_delete_cache(pair_counter)
except _REDIS_FAILURES as err:
self._warn_redis(err)
self.counters.delete_cache(pair_counter)
def _warn_redis(self, err: Exception) -> None:
verbose_proxy_logger.warning(
"Redis failed while counting Admin UI sign-in attempts; using this worker's own counters "
"until it recovers: %s",
err,
)
keys: Final = self._keys(username)
store: Final = self._store()
if store is not None:
_ = await store.clear(keys)
self.counters.delete_cache(keys.pair_counter)
@staticmethod
def refused(retry_after: int) -> ProxyException:

90
litellm/redis/stores/login_throttle.py Normal file
View file

@ -0,0 +1,90 @@
import hashlib
from dataclasses import dataclass
from datetime import timedelta
from typing import Final, NamedTuple, Protocol, TypeAlias
from pydantic import TypeAdapter, ValidationError
from litellm.constants import LOGIN_THROTTLE_CACHE_KEY_PREFIX
from litellm.redis.transport.contract import RedisFailure, RedisTransport
BlockTtls: TypeAlias = tuple[int, int]
_BLOCK_TTLS: Final = TypeAdapter[BlockTtls](BlockTtls)
# KEYS: pair counter, pair block, source counter, source block (one cluster slot via the source hash tag)
# ARGV: pair limit, source limit (0 = source scope off), window seconds, block seconds
# Both scripts return {pair block TTL, source block TTL}; 0 or below means not blocked
_BLOCK_TTLS_LUA: Final = "return {redis.call('TTL', KEYS[2]), redis.call('TTL', KEYS[4])}"
_RECORD_FAILURE_LUA: Final = (
"local function bump(count_key, block_key, limit) "
"local blocked = redis.call('TTL', block_key) "
"if blocked > 0 then return blocked end "
"local count = redis.call('INCR', count_key) "
"if redis.call('TTL', count_key) < 0 then redis.call('EXPIRE', count_key, ARGV[3]) end "
"if count > limit then redis.call('SET', block_key, '1', 'EX', ARGV[4]) return tonumber(ARGV[4]) end "
"return 0 end "
"local user_block = bump(KEYS[1], KEYS[2], tonumber(ARGV[1])) "
"local source_block = 0 "
"if tonumber(ARGV[2]) > 0 and user_block == 0 then "
"source_block = bump(KEYS[3], KEYS[4], tonumber(ARGV[2])) end "
"return {user_block, source_block}"
)
class LoginKeys(NamedTuple):
pair_counter: str
pair_block: str
source_counter: str
source_block: str
def login_keys(source_group: str, username: str) -> LoginKeys:
user: Final = hashlib.sha256(username.casefold().encode("utf-8")).hexdigest()
return LoginKeys(
pair_counter=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{source_group}}}:user:{user}",
pair_block=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{source_group}}}:block:user:{user}",
source_counter=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{source_group}}}:source",
source_block=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{source_group}}}:block:source",
)
class LoginThrottleStore(Protocol):
async def blocks(self, keys: LoginKeys) -> BlockTtls | RedisFailure: ...
async def record_failure(
self, keys: LoginKeys, *, user_limit: int, source_limit: int, window: timedelta, block: timedelta
) -> BlockTtls | RedisFailure: ...
async def clear(self, keys: LoginKeys) -> None | RedisFailure: ...
@dataclass(frozen=True, slots=True)
class TransportLoginThrottleStore:
transport: RedisTransport
async def blocks(self, keys: LoginKeys) -> BlockTtls | RedisFailure:
return _ttls(await self.transport.eval(_BLOCK_TTLS_LUA, keys, ()))
async def record_failure(
self, keys: LoginKeys, *, user_limit: int, source_limit: int, window: timedelta, block: timedelta
) -> BlockTtls | RedisFailure:
args: Final = (user_limit, source_limit, _whole_seconds(window), _whole_seconds(block))
return _ttls(await self.transport.eval(_RECORD_FAILURE_LUA, keys, args))
async def clear(self, keys: LoginKeys) -> None | RedisFailure:
deleted: Final = await self.transport.delete(keys.pair_counter)
return deleted if isinstance(deleted, RedisFailure) else None
def _whole_seconds(ttl: timedelta) -> int:
return ttl // timedelta(seconds=1)
def _ttls(reply: object) -> BlockTtls | RedisFailure:
if isinstance(reply, RedisFailure):
return reply
try:
return _BLOCK_TTLS.validate_python(reply, strict=True)
except ValidationError as e:
return RedisFailure(kind="decode", message=f"login throttle reply is not two block TTLs: {e}")

131
tests/unit/proxy/auth/test_login_utils.py
View file

@ -14,10 +14,13 @@ from types import SimpleNamespace
from typing import TYPE_CHECKING, Final
from unittest.mock import AsyncMock, MagicMock, patch
import fakeredis
import httpx
import pytest
import respx
from tests._support.fake_redis import fake_redis_cache
if TYPE_CHECKING:
from litellm.proxy.auth.login_throttle import LoginThrottle
@ -1367,66 +1370,8 @@ async def test_disabling_the_control_lets_every_attempt_through(monkeypatch):
assert _local_count(throttle, throttle._keys("admin").pair_counter) == 0
class _FakeRedis:
"""Redis whose only writes are the throttle's two scripts, run atomically as one call each.
Mirrors the Lua: a blocked key returns its remaining block time and is not counted; a counter
is expired on first write; one over the limit installs the block; a blocked pair stops the
source from being counted. The real scripts are exercised against a live Redis in the PR's
proof, this fake only has to be faithful enough for the worker-sharing tests.
"""
def __init__(self):
self.values: dict = {}
self.ttls: dict = {}
self.scripts: list[str] = []
def async_register_script(self, script: str):
from litellm.proxy.auth import login_throttle as lt
async def _run(keys, args):
self.scripts.append(script)
if script == lt._BLOCK_TTLS_LUA:
return [self._ttl(keys[1]), self._ttl(keys[3])]
assert script == lt._RECORD_FAILURE_LUA
user_limit, source_limit, window, block = (int(a) for a in args)
user_block = self._bump(keys[0], keys[1], user_limit, window, block)
if source_limit > 0 and user_block == 0:
return [user_block, self._bump(keys[2], keys[3], source_limit, window, block)]
return [user_block, 0]
return _run
def _ttl(self, key: str) -> int:
return self.ttls.get(key, -2) if key in self.values else -2
def _bump(self, count_key: str, block_key: str, limit: int, window: int, block: int) -> int:
if self._ttl(block_key) > 0:
return self._ttl(block_key)
self.values[count_key] = self.values.get(count_key, 0) + 1
self.ttls.setdefault(count_key, window)
if self.values[count_key] > limit:
self.values[block_key] = 1
self.ttls[block_key] = block
return block
return 0
async def async_delete_cache(self, key):
self.values.pop(key, None)
self.ttls.pop(key, None)
class _DownRedis(_FakeRedis):
"""Redis whose every call fails, as during an outage or an open circuit breaker."""
def async_register_script(self, script: str):
async def _run(keys, args):
raise ConnectionError("redis is down")
return _run
async def async_delete_cache(self, key):
raise ConnectionError("redis is down")
async def _redis_keys(server: fakeredis.FakeServer) -> list[str]:
return [key.decode() for key in await fakeredis.FakeAsyncRedis(server=server).keys("*")]
@pytest.mark.asyncio
@ -1437,57 +1382,65 @@ async def test_redis_is_the_only_counter_while_it_answers(monkeypatch):
monkeypatch.setenv("UI_USERNAME", "admin")
monkeypatch.setenv("UI_PASSWORD", "right")
monkeypatch.setenv("DATABASE_URL", "postgresql://stub")
redis = _FakeRedis()
first_worker = _throttle(user_limit=2, stores=_stores(), redis_cache=redis)
second_worker = _throttle(user_limit=2, stores=_stores(), redis_cache=redis)
server = fakeredis.FakeServer()
raw = fakeredis.FakeAsyncRedis(server=server)
redis = fake_redis_cache(server)
first_worker = _throttle(user_limit=2, window_seconds=60, block_seconds=300, stores=_stores(), redis_cache=redis)
second_worker = _throttle(user_limit=2, window_seconds=60, block_seconds=300, stores=_stores(), redis_cache=redis)
keys = first_worker._keys("user@corp.com")
assert [await _fail(first_worker, username="user@corp.com") for _ in range(3)] == ["401"] * 3
assert not [k for k in first_worker.counters.cache_dict if str(k).startswith(LOGIN_THROTTLE_CACHE_KEY_PREFIX)], (
"with Redis answering, no worker may keep a counter of its own"
)
assert not first_worker.blocks.cache_dict
assert await raw.get(keys.pair_counter) == b"3"
assert 0 < await raw.ttl(keys.pair_counter) <= 60
assert 0 < await raw.ttl(keys.pair_block) <= 300, "the block lives in Redis, where every worker reads it"
assert await _fail(second_worker, username="user@corp.com") == "429", "the second worker sees the block"
block_keys = [k for k in redis.values if ":block:user:" in k]
assert block_keys, "the block lives in Redis, where every worker reads it"
for key in block_keys:
await redis.async_delete_cache(key)
await raw.delete(keys.pair_block)
await _db_login(second_worker, "user@corp.com", "right", correct=True)
assert not [k for k in redis.values if ":user:" in k and ":block:" not in k], (
"success clears the shared pair counter"
)
assert keys.pair_counter not in await _redis_keys(server), "success clears the shared pair counter"
@pytest.mark.asyncio
async def test_the_shared_block_counts_the_source_once_the_pair_is_not_blocked(monkeypatch):
monkeypatch.setenv("UI_USERNAME", "admin")
monkeypatch.setenv("UI_PASSWORD", "right")
server = fakeredis.FakeServer()
raw = fakeredis.FakeAsyncRedis(server=server)
throttle = _throttle(user_limit=5, source_limit=2, block_seconds=120, redis_cache=fake_redis_cache(server))
keys = [throttle._keys(f"user-{i}@corp.com") for i in range(3)]
assert [await _fail(throttle, username=f"user-{i}@corp.com") for i in range(3)] == ["401"] * 3
assert await raw.get(keys[0].source_counter) == b"3"
assert 0 < await raw.ttl(keys[0].source_block) <= 120
assert await _fail(throttle, username="someone-else@corp.com") == "429"
@pytest.mark.asyncio
async def test_a_redis_outage_falls_back_to_this_workers_own_counter(monkeypatch):
"""With Redis raising, guesses are still counted and blocked per worker, with a warning, instead of unbounded."""
import logging
from litellm._logging import verbose_proxy_logger
"""With Redis down, guesses are still counted and blocked per worker instead of unbounded."""
from litellm.proxy._types import ProxyException
monkeypatch.setenv("UI_USERNAME", "admin")
monkeypatch.setenv("UI_PASSWORD", "right")
_set_throttle_clock(monkeypatch, 1_000.0)
throttle = _throttle(user_limit=2, block_seconds=300, redis_cache=_DownRedis())
server = fakeredis.FakeServer()
throttle = _throttle(user_limit=2, block_seconds=300, redis_cache=fake_redis_cache(server))
server.connected = False
records: list[logging.LogRecord] = []
handler = logging.Handler()
handler.emit = records.append
verbose_proxy_logger.addHandler(handler)
try:
assert [await _fail(throttle) for _ in range(3)] == ["401"] * 3
_set_throttle_clock(monkeypatch, 1_000.6)
with pytest.raises(ProxyException) as blocked:
await _guess(throttle)
finally:
verbose_proxy_logger.removeHandler(handler)
assert [await _fail(throttle) for _ in range(3)] == ["401"] * 3
_set_throttle_clock(monkeypatch, 1_000.6)
with pytest.raises(ProxyException) as blocked:
await _guess(throttle)
assert blocked.value.code == "429"
assert blocked.value.headers.get("Retry-After") == "300"
assert any("Redis failed while counting Admin UI sign-in attempts" in r.getMessage() for r in records)
@pytest.mark.asyncio
@ -1496,7 +1449,9 @@ async def test_a_failed_redis_delete_still_clears_this_workers_counter(monkeypat
monkeypatch.setenv("UI_USERNAME", "admin")
monkeypatch.setenv("UI_PASSWORD", "right")
monkeypatch.setenv("DATABASE_URL", "postgresql://stub")
throttle = _throttle(user_limit=5, redis_cache=_DownRedis())
server = fakeredis.FakeServer()
throttle = _throttle(user_limit=5, redis_cache=fake_redis_cache(server))
server.connected = False
key = throttle._keys("user@corp.com").pair_counter
assert [await _fail(throttle, username="user@corp.com") for _ in range(2)] == ["401", "401"]