diff --git a/backend/open_webui/routers/auths.py b/backend/open_webui/routers/auths.py index 31b0e97e7b..75a3e0961d 100644 --- a/backend/open_webui/routers/auths.py +++ b/backend/open_webui/routers/auths.py @@ -77,7 +77,6 @@ from open_webui.utils.auth import ( from open_webui.utils.groups import apply_default_group_assignment from open_webui.utils.misc import parse_duration, validate_email_format from open_webui.utils.rate_limit import RateLimiter -from open_webui.utils.redis import get_redis_client from pydantic import BaseModel, StrictStr, field_validator from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession @@ -88,12 +87,11 @@ log = logging.getLogger(__name__) # Forgive us our failed attempts, as we forgive those # who exceed their allotted rate against this gate. -signin_rate_limiter = RateLimiter(redis_client=get_redis_client(), limit=5 * 3, window=60 * 3) +signin_rate_limiter = RateLimiter(limit=5 * 3, window=60 * 3) # Best-effort throttle only: there is no caller identity before the provider answers, # and deployments may derive request.client from proxy headers. token_exchange_rate_limiter = ( RateLimiter( - redis_client=get_redis_client(), limit=OAUTH_TOKEN_EXCHANGE_RATE_LIMIT, window=OAUTH_TOKEN_EXCHANGE_RATE_LIMIT_WINDOW, ) @@ -816,7 +814,7 @@ async def signin( db=db, ) else: - if signin_rate_limiter.is_limited(form_data.email.lower()): + if await signin_rate_limiter.is_limited(request.app.state.redis, form_data.email.lower()): raise HTTPException( status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail=ERROR_MESSAGES.RATE_LIMIT_EXCEEDED, @@ -1637,8 +1635,8 @@ async def token_exchange( detail='Token exchange is disabled', ) - if token_exchange_rate_limiter and token_exchange_rate_limiter.is_limited( - request.client.host if request.client else 'unknown' + if token_exchange_rate_limiter and await token_exchange_rate_limiter.is_limited( + request.app.state.redis, request.client.host if request.client else 'unknown' ): raise HTTPException( status_code=status.HTTP_429_TOO_MANY_REQUESTS, diff --git a/backend/open_webui/utils/rate_limit.py b/backend/open_webui/utils/rate_limit.py index 9602c04a14..c19f5d11d0 100644 --- a/backend/open_webui/utils/rate_limit.py +++ b/backend/open_webui/utils/rate_limit.py @@ -1,7 +1,8 @@ import time -from typing import Dict, Optional +from typing import Optional from open_webui.env import REDIS_KEY_PREFIX +from redis.asyncio import Redis class RateLimiter: @@ -10,30 +11,26 @@ class RateLimiter: Falls back to in-memory storage if Redis is not available. """ - # In-memory fallback storage - _memory_store: Dict[str, Dict[int, int]] = {} - def __init__( self, - redis_client, limit: int, window: int, bucket_size: int = 60, enabled: bool = True, ): """ - :param redis_client: Redis client instance or None :param limit: Max allowed events in the window :param window: Time window in seconds :param bucket_size: Bucket resolution :param enabled: Turn on/off rate limiting globally """ - self.r = redis_client self.limit = limit self.window = window self.bucket_size = bucket_size self.num_buckets = window // bucket_size self.enabled = enabled + # bucket index -> rate-limit key -> hits + self._memory_store: dict[int, dict[str, int]] = {} def _bucket_key(self, key: str, bucket_index: int) -> str: return f'{REDIS_KEY_PREFIX}:ratelimit:{key.lower()}:{bucket_index}' @@ -41,10 +38,13 @@ class RateLimiter: def _current_bucket(self) -> int: return int(time.time()) // self.bucket_size - def _redis_available(self) -> bool: - return self.r is not None + def _prune_memory_store(self, now_bucket: int) -> None: + min_bucket = now_bucket - self.num_buckets + expired = [bucket_index for bucket_index in self._memory_store if bucket_index < min_bucket] + for bucket_index in expired: + del self._memory_store[bucket_index] - def is_limited(self, key: str) -> bool: + async def is_limited(self, redis: Redis | None, key: str) -> bool: """ Main rate-limit check. Gracefully handles missing or failing Redis. @@ -52,85 +52,63 @@ class RateLimiter: if not self.enabled: return False - if self._redis_available(): + if redis is not None: try: - return self._is_limited_redis(key) + return await self._is_limited_redis(redis, key) except Exception: return self._is_limited_memory(key) else: return self._is_limited_memory(key) - def get_count(self, key: str) -> int: + async def get_count(self, redis: Redis | None, key: str) -> int: if not self.enabled: return 0 - if self._redis_available(): + if redis is not None: try: - return self._get_count_redis(key) + return await self._get_count_redis(redis, key) except Exception: return self._get_count_memory(key) else: return self._get_count_memory(key) - def remaining(self, key: str) -> int: - used = self.get_count(key) + async def remaining(self, redis: Redis | None, key: str) -> int: + used = await self.get_count(redis, key) return max(0, self.limit - used) - def _is_limited_redis(self, key: str) -> bool: + async def _is_limited_redis(self, redis: Redis, key: str) -> bool: now_bucket = self._current_bucket() bucket_key = self._bucket_key(key, now_bucket) - attempts = self.r.incr(bucket_key) + attempts = await redis.incr(bucket_key) if attempts == 1: - self.r.expire(bucket_key, self.window + self.bucket_size) + await redis.expire(bucket_key, self.window + self.bucket_size) # Collect buckets buckets = [self._bucket_key(key, now_bucket - i) for i in range(self.num_buckets + 1)] - counts = self.r.mget(buckets) + counts = await redis.mget(buckets) total = sum(int(c) for c in counts if c) return total > self.limit - def _get_count_redis(self, key: str) -> int: + async def _get_count_redis(self, redis: Redis, key: str) -> int: now_bucket = self._current_bucket() buckets = [self._bucket_key(key, now_bucket - i) for i in range(self.num_buckets + 1)] - counts = self.r.mget(buckets) + counts = await redis.mget(buckets) return sum(int(c) for c in counts if c) def _is_limited_memory(self, key: str) -> bool: now_bucket = self._current_bucket() + self._prune_memory_store(now_bucket) - # Init storage - if key not in self._memory_store: - self._memory_store[key] = {} + current_bucket_counts = self._memory_store.setdefault(now_bucket, {}) + current_bucket_counts[key] = current_bucket_counts.get(key, 0) + 1 - store = self._memory_store[key] - - # Increment bucket - store[now_bucket] = store.get(now_bucket, 0) + 1 - - # Drop expired buckets - min_bucket = now_bucket - self.num_buckets - expired = [b for b in store if b < min_bucket] - for b in expired: - del store[b] - - # Count totals - total = sum(store.values()) + total = sum(bucket_counts.get(key, 0) for bucket_counts in self._memory_store.values()) return total > self.limit def _get_count_memory(self, key: str) -> int: now_bucket = self._current_bucket() - if key not in self._memory_store: - return 0 - - store = self._memory_store[key] - min_bucket = now_bucket - self.num_buckets - - # Remove expired - expired = [b for b in store if b < min_bucket] - for b in expired: - del store[b] - - return sum(store.values()) + self._prune_memory_store(now_bucket) + return sum(bucket_counts.get(key, 0) for bucket_counts in self._memory_store.values())