diff --git a/litellm/proxy/auth/login_throttle.py b/litellm/proxy/auth/login_throttle.py index 0536e293f13..205d7f48601 100644 --- a/litellm/proxy/auth/login_throttle.py +++ b/litellm/proxy/auth/login_throttle.py @@ -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: diff --git a/litellm/redis/stores/login_throttle.py b/litellm/redis/stores/login_throttle.py new file mode 100644 index 00000000000..56394f8fa95 --- /dev/null +++ b/litellm/redis/stores/login_throttle.py @@ -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}") diff --git a/tests/unit/proxy/auth/test_login_utils.py b/tests/unit/proxy/auth/test_login_utils.py index 5148e38aa61..1d77b29a9eb 100644 --- a/tests/unit/proxy/auth/test_login_utils.py +++ b/tests/unit/proxy/auth/test_login_utils.py @@ -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"]