""" Rate Limiting Middleware Provides token bucket and sliding window rate limiting """ import time import asyncio import logging from typing import Optional, Dict, Callable, Any from collections import defaultdict from dataclasses import dataclass, field from datetime import datetime from fastapi import Request, Response, HTTPException, status from fastapi.responses import JSONResponse from starlette.middleware.base import BaseHTTPMiddleware from app.config import settings logger = logging.getLogger(__name__) @dataclass class RateLimitInfo: """Information about a rate limit bucket""" tokens: float last_update: float requests_per_minute: int = 0 requests_per_hour: int = 0 minute_window_start: float = 0 hour_window_start: float = 0 class RateLimitStore: """ In-memory rate limit storage For production, use Redis with RedisRateLimitStore """ def __init__(self): self._store: Dict[str, RateLimitInfo] = {} self._lock = asyncio.Lock() async def get(self, key: str) -> Optional[RateLimitInfo]: async with self._lock: return self._store.get(key) async def set(self, key: str, info: RateLimitInfo): async with self._lock: self._store[key] = info async def cleanup_expired(self, max_age: float = 3600): """Remove entries older than max_age seconds""" current_time = time.time() async with self._lock: expired_keys = [ key for key, info in self._store.items() if current_time - info.last_update > max_age ] for key in expired_keys: del self._store[key] class RedisRateLimitStore: """ Redis-based rate limit storage for distributed systems """ def __init__(self, redis_url: str): self.redis_url = redis_url self._redis = None async def _get_redis(self): if self._redis is None: try: import redis.asyncio as redis self._redis = redis.from_url(self.redis_url) except ImportError: logger.warning("Redis not available, falling back to in-memory storage") return None return self._redis async def get(self, key: str) -> Optional[RateLimitInfo]: redis = await self._get_redis() if redis is None: return None try: import json data = await redis.get(f"ratelimit:{key}") if data: info_dict = json.loads(data) return RateLimitInfo(**info_dict) except Exception as e: logger.error(f"Redis get error: {e}") return None async def set(self, key: str, info: RateLimitInfo, ttl: int = 3600): redis = await self._get_redis() if redis is None: return try: import json data = json.dumps({ "tokens": info.tokens, "last_update": info.last_update, "requests_per_minute": info.requests_per_minute, "requests_per_hour": info.requests_per_hour, "minute_window_start": info.minute_window_start, "hour_window_start": info.hour_window_start }) await redis.setex(f"ratelimit:{key}", ttl, data) except Exception as e: logger.error(f"Redis set error: {e}") class RateLimiter: """ Token bucket rate limiter with sliding window counters """ def __init__( self, requests_per_minute: int = 60, requests_per_hour: int = 1000, burst_size: int = 10, store: Optional[RateLimitStore] = None ): self.requests_per_minute = requests_per_minute self.requests_per_hour = requests_per_hour self.burst_size = burst_size self.store = store or RateLimitStore() # Token bucket refill rate (tokens per second) self.refill_rate = requests_per_minute / 60.0 def _get_client_id(self, request: Request) -> str: """Extract client identifier from request""" # Check for API key api_key = request.headers.get("X-API-Key") if api_key: return f"api:{api_key[:16]}" # Check for authenticated user auth_header = request.headers.get("Authorization") if auth_header and auth_header.startswith("Bearer "): # Use token hash for identification token = auth_header[7:] return f"token:{hash(token) % 1000000}" # Fall back to IP address forwarded = request.headers.get("X-Forwarded-For") if forwarded: ip = forwarded.split(",")[0].strip() else: ip = request.client.host if request.client else "unknown" return f"ip:{ip}" async def check_rate_limit(self, request: Request) -> tuple[bool, Dict[str, Any]]: """ Check if request is within rate limits Returns (allowed, info) tuple """ client_id = self._get_client_id(request) current_time = time.time() # Get or create rate limit info info = await self.store.get(client_id) if info is None: info = RateLimitInfo( tokens=float(self.burst_size), last_update=current_time, requests_per_minute=0, requests_per_hour=0, minute_window_start=current_time, hour_window_start=current_time ) # Reset sliding windows if needed if current_time - info.minute_window_start >= 60: info.requests_per_minute = 0 info.minute_window_start = current_time if current_time - info.hour_window_start >= 3600: info.requests_per_hour = 0 info.hour_window_start = current_time # Refill tokens (token bucket algorithm) time_passed = current_time - info.last_update info.tokens = min( self.burst_size, info.tokens + time_passed * self.refill_rate ) info.last_update = current_time # Check limits rate_limit_info = { "limit_minute": self.requests_per_minute, "limit_hour": self.requests_per_hour, "remaining_minute": max(0, self.requests_per_minute - info.requests_per_minute - 1), "remaining_hour": max(0, self.requests_per_hour - info.requests_per_hour - 1), "reset_minute": int(info.minute_window_start + 60), "reset_hour": int(info.hour_window_start + 3600), "tokens_remaining": int(info.tokens) } # Check token bucket (for burst protection) if info.tokens < 1: rate_limit_info["retry_after"] = int(1 / self.refill_rate) return False, rate_limit_info # Check per-minute limit if info.requests_per_minute >= self.requests_per_minute: rate_limit_info["retry_after"] = int(60 - (current_time - info.minute_window_start)) return False, rate_limit_info # Check per-hour limit if info.requests_per_hour >= self.requests_per_hour: rate_limit_info["retry_after"] = int(3600 - (current_time - info.hour_window_start)) return False, rate_limit_info # Consume token and increment counters info.tokens -= 1 info.requests_per_minute += 1 info.requests_per_hour += 1 await self.store.set(client_id, info) return True, rate_limit_info class RateLimitMiddleware(BaseHTTPMiddleware): """ FastAPI middleware for rate limiting """ def __init__( self, app, requests_per_minute: int = 60, requests_per_hour: int = 1000, burst_size: int = 10, exclude_paths: Optional[list] = None, use_redis: bool = False, redis_url: Optional[str] = None ): super().__init__(app) # Set up store store = None if use_redis and redis_url: store = RedisRateLimitStore(redis_url) self.limiter = RateLimiter( requests_per_minute=requests_per_minute, requests_per_hour=requests_per_hour, burst_size=burst_size, store=store ) self.exclude_paths = exclude_paths or [ "/health", "/docs", "/redoc", "/openapi.json", "/api/versions" ] async def dispatch(self, request: Request, call_next) -> Response: # Skip rate limiting for excluded paths path = request.url.path if any(path.startswith(excluded) for excluded in self.exclude_paths): return await call_next(request) # Check rate limit allowed, info = await self.limiter.check_rate_limit(request) if not allowed: return JSONResponse( status_code=status.HTTP_429_TOO_MANY_REQUESTS, content={ "error": "rate_limit_exceeded", "message": "Too many requests. Please slow down.", "retry_after": info.get("retry_after", 60) }, headers={ "X-RateLimit-Limit-Minute": str(info["limit_minute"]), "X-RateLimit-Limit-Hour": str(info["limit_hour"]), "X-RateLimit-Remaining-Minute": str(info["remaining_minute"]), "X-RateLimit-Remaining-Hour": str(info["remaining_hour"]), "X-RateLimit-Reset-Minute": str(info["reset_minute"]), "X-RateLimit-Reset-Hour": str(info["reset_hour"]), "Retry-After": str(info.get("retry_after", 60)) } ) # Process request response = await call_next(request) # Add rate limit headers to successful responses response.headers["X-RateLimit-Limit-Minute"] = str(info["limit_minute"]) response.headers["X-RateLimit-Limit-Hour"] = str(info["limit_hour"]) response.headers["X-RateLimit-Remaining-Minute"] = str(info["remaining_minute"]) response.headers["X-RateLimit-Remaining-Hour"] = str(info["remaining_hour"]) response.headers["X-RateLimit-Reset-Minute"] = str(info["reset_minute"]) response.headers["X-RateLimit-Reset-Hour"] = str(info["reset_hour"]) return response def rate_limit( requests_per_minute: int = 60, requests_per_hour: int = 1000 ): """ Decorator for per-route rate limiting Usage: @router.get("/expensive-operation") @rate_limit(requests_per_minute=10, requests_per_hour=100) async def expensive_operation(): ... """ _limiter = RateLimiter( requests_per_minute=requests_per_minute, requests_per_hour=requests_per_hour ) def decorator(func: Callable): async def wrapper(request: Request, *args, **kwargs): allowed, info = await _limiter.check_rate_limit(request) if not allowed: raise HTTPException( status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail={ "error": "rate_limit_exceeded", "message": "Too many requests for this endpoint", "retry_after": info.get("retry_after", 60) } ) return await func(request, *args, **kwargs) return wrapper return decorator # Convenience function to create middleware def create_rate_limit_middleware( requests_per_minute: int = None, requests_per_hour: int = None, burst_size: int = None, use_redis: bool = False ): """Create rate limit middleware with default settings from config""" return lambda app: RateLimitMiddleware( app, requests_per_minute=requests_per_minute or settings.RATE_LIMIT_PER_MINUTE, requests_per_hour=requests_per_hour or settings.RATE_LIMIT_PER_HOUR, burst_size=burst_size or settings.RATE_LIMIT_BURST, use_redis=use_redis, redis_url=settings.REDIS_URL if use_redis else None )