import asyncio import contextlib import time from collections import defaultdict from collections.abc import Callable from dataclasses import dataclass, field from enum import StrEnum from functools import wraps class RateLimitScope(StrEnum): LOGIN = "login" API = "api" SIGNUP = "signup" PASSWORD_RESET = "password_reset" TOKEN_REFRESH = "token_refresh" @dataclass class ScopeConfig: max_attempts: int = 10 window_seconds: int = 60 cooldown_seconds: int = 300 loopback_exempt: bool = True DEFAULT_SCOPE_CONFIGS = { RateLimitScope.LOGIN: ScopeConfig(max_attempts=10, window_seconds=60, cooldown_seconds=300), RateLimitScope.API: ScopeConfig(max_attempts=100, window_seconds=60, cooldown_seconds=60), RateLimitScope.SIGNUP: ScopeConfig(max_attempts=3, window_seconds=3600, cooldown_seconds=7200), RateLimitScope.PASSWORD_RESET: ScopeConfig(max_attempts=5, window_seconds=300, cooldown_seconds=600), RateLimitScope.TOKEN_REFRESH: ScopeConfig(max_attempts=30, window_seconds=60, cooldown_seconds=120), } _LOCAL_IPS = frozenset({"127.0.0.1", "::1", "localhost"}) @dataclass class _BucketState: timestamps: list[float] = field(default_factory=list) cooldown_until: float = 0.0 lock: asyncio.Lock = field(default_factory=asyncio.Lock) class AuthRateLimiter: """认证级限流器：内存滑动窗口 + 多 Scope + 本地回环豁免 + 锁定冷却。 特性： - 多 Scope 独立配置（login/api/signup/password_reset/token_refresh） - 滑动窗口计数，窗口外记录自动淘汰 - 本地回环 (127.0.0.1, ::1) 可豁免 - 锁定冷却期 5min（超限后拒绝所有尝试） - asyncio.Lock 串行化防竞态 - 后台定期清理失效 bucket，防止 DDoS 下内存持续增长 """ _CLEANUP_INTERVAL = 300 # seconds between stale bucket cleanups def __init__(self, scope_configs: dict[RateLimitScope, ScopeConfig] | None = None): self._configs = scope_configs or DEFAULT_SCOPE_CONFIGS self._buckets: dict[str, dict[RateLimitScope, _BucketState]] = defaultdict(dict) self._lock = asyncio.Lock() self._last_cleanup: float = 0.0 self._cleanup_task: asyncio.Task | None = None self._start_background_cleanup() def _is_loopback(self, identifier: str) -> bool: if identifier in _LOCAL_IPS: return True host = identifier.split(":")[0] if host.startswith("127."): return True try: import ipaddress ip = ipaddress.ip_address(host) return ip.is_loopback except ValueError: return False async def _get_bucket(self, identifier: str, scope: RateLimitScope) -> _BucketState: async with self._lock: scope_buckets = self._buckets[identifier] if scope not in scope_buckets: scope_buckets[scope] = _BucketState() return scope_buckets[scope] async def _cleanup_stale_buckets(self, now: float) -> None: async with self._lock: stale_identifiers: list[str] = [] for identifier, scopes in self._buckets.items(): active_scopes: dict[RateLimitScope, _BucketState] = {} for scope, bucket in scopes.items(): config = self._configs.get(scope) window = config.window_seconds if config else 60 bucket.timestamps = [t for t in bucket.timestamps if t > now - window] has_activity = bucket.timestamps or bucket.cooldown_until > now if has_activity: active_scopes[scope] = bucket if active_scopes: self._buckets[identifier] = active_scopes else: stale_identifiers.append(identifier) for identifier in stale_identifiers: del self._buckets[identifier] def _start_background_cleanup(self) -> None: async def _loop() -> None: while True: await asyncio.sleep(self._CLEANUP_INTERVAL) await self._cleanup_stale_buckets(time.monotonic()) try: loop = asyncio.get_running_loop() self._cleanup_task = loop.create_task(_loop()) except RuntimeError: pass async def shutdown(self) -> None: if self._cleanup_task and not self._cleanup_task.done(): self._cleanup_task.cancel() with contextlib.suppress(asyncio.CancelledError): await self._cleanup_task async def check_and_record( self, identifier: str, scope: RateLimitScope, ) -> bool: """检查是否允许请求，并记录本次尝试。 返回 True 表示允许，False 表示被限流。 """ config = self._configs.get(scope) if config is None: return True if config.loopback_exempt and self._is_loopback(identifier): return True now = time.monotonic() if now - self._last_cleanup > self._CLEANUP_INTERVAL: self._last_cleanup = now await self._cleanup_stale_buckets(now) bucket = await self._get_bucket(identifier, scope) now = time.monotonic() async with bucket.lock: if bucket.cooldown_until > now: return False window_start = now - config.window_seconds bucket.timestamps = [t for t in bucket.timestamps if t > window_start] if len(bucket.timestamps) >= config.max_attempts: bucket.cooldown_until = now + config.cooldown_seconds return False bucket.timestamps.append(now) return True async def check_and_record_or_raise( self, identifier: str, scope: RateLimitScope, ) -> None: if not await self.check_and_record(identifier, scope): config = self._configs.get(scope) retry_after = config.cooldown_seconds if config else 60 raise RateLimitExceeded(scope=scope, retry_after=retry_after) async def clear_cooldown(self, identifier: str, scope: RateLimitScope): bucket = await self._get_bucket(identifier, scope) async with bucket.lock: bucket.cooldown_until = 0.0 async def reset(self, identifier: str, scope: RateLimitScope | None = None): async with self._lock: if scope: if identifier in self._buckets and scope in self._buckets[identifier]: async with self._buckets[identifier][scope].lock: self._buckets[identifier].pop(scope, None) else: self._buckets.pop(identifier, None) def get_stats(self) -> dict: now = time.monotonic() stats: dict[str, dict] = {} for identifier, scopes in self._buckets.items(): stats[identifier] = {} for scope, bucket in scopes.items(): config = self._configs.get(scope) window = config.window_seconds if config else 60 active = len([t for t in bucket.timestamps if t > now - window]) stats[identifier][scope.value] = { "attempts_in_window": active, "cooldown_active": bucket.cooldown_until > now, "cooldown_remaining": max(0, bucket.cooldown_until - now), } return stats def wrap( self, scope: RateLimitScope, identifier_fn: Callable[..., str] | None = None, ): """装饰器：将限流包装到异步函数上。 使用方式： @auth_rate_limiter.wrap(RateLimitScope.LOGIN) async def login(request: Request): ... """ def decorator(fn): @wraps(fn) async def wrapper(*args, **kwargs): ident = identifier_fn(*args, **kwargs) if identifier_fn else "default" await self.check_and_record_or_raise(ident, scope) return await fn(*args, **kwargs) return wrapper return decorator class RateLimitExceeded(Exception): def __init__(self, scope: RateLimitScope, retry_after: int = 60): self.scope = scope self.retry_after = retry_after super().__init__(f"Rate limit exceeded for scope '{scope.value}', retry after {retry_after}s") auth_rate_limiter = AuthRateLimiter()