|
3 | 3 | Prevents abuse and manages request quotas |
4 | 4 | """ |
5 | 5 | import time |
6 | | -from typing import Optional, Dict |
| 6 | +from typing import Optional, Dict, Callable |
7 | 7 | from datetime import datetime, timedelta |
| 8 | +from functools import wraps |
8 | 9 | import hashlib |
9 | 10 | import secrets |
10 | 11 | from dataclasses import dataclass |
| 12 | +from fastapi import HTTPException, Request |
| 13 | + |
| 14 | + |
| 15 | +# In-memory rate limit storage (per-process, resets on restart) |
| 16 | +_rate_limit_store: Dict[str, list] = {} |
| 17 | + |
| 18 | + |
| 19 | +def rate_limit(requests_per_minute: int = 60): |
| 20 | + """ |
| 21 | + Simple rate limit decorator for FastAPI routes. |
| 22 | + Uses in-memory storage - suitable for single-instance deployments. |
| 23 | + For production, use Redis-backed RateLimiter class instead. |
| 24 | + """ |
| 25 | + def decorator(func: Callable): |
| 26 | + @wraps(func) |
| 27 | + async def wrapper(*args, **kwargs): |
| 28 | + # Try to get request from kwargs or args |
| 29 | + request = kwargs.get('request') |
| 30 | + if not request: |
| 31 | + for arg in args: |
| 32 | + if isinstance(arg, Request): |
| 33 | + request = arg |
| 34 | + break |
| 35 | + |
| 36 | + # Get client identifier (IP or fallback) |
| 37 | + if request: |
| 38 | + client_id = request.client.host if request.client else "unknown" |
| 39 | + else: |
| 40 | + client_id = "unknown" |
| 41 | + |
| 42 | + key = f"{func.__name__}:{client_id}" |
| 43 | + now = time.time() |
| 44 | + window_start = now - 60 |
| 45 | + |
| 46 | + # Clean old entries and get current count |
| 47 | + if key not in _rate_limit_store: |
| 48 | + _rate_limit_store[key] = [] |
| 49 | + |
| 50 | + _rate_limit_store[key] = [t for t in _rate_limit_store[key] if t > window_start] |
| 51 | + |
| 52 | + if len(_rate_limit_store[key]) >= requests_per_minute: |
| 53 | + raise HTTPException( |
| 54 | + status_code=429, |
| 55 | + detail=f"Rate limit exceeded. Max {requests_per_minute} requests per minute." |
| 56 | + ) |
| 57 | + |
| 58 | + _rate_limit_store[key].append(now) |
| 59 | + return await func(*args, **kwargs) |
| 60 | + |
| 61 | + return wrapper |
| 62 | + return decorator |
11 | 63 |
|
12 | 64 |
|
13 | 65 | @dataclass |
|
0 commit comments