File size: 5,208 Bytes
b540dd9 f85338f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | from fastapi import Request, HTTPException
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from typing import Dict, Optional
import time
from collections import defaultdict, deque
from datetime import datetime, timedelta
import json
class RateLimitMiddleware(BaseHTTPMiddleware):
"""Rate limiting middleware to protect API endpoints"""
def __init__(self, app, requests_per_window: int = 100, window_seconds: int = 3600):
super().__init__(app)
self.requests_per_window = requests_per_window
self.window_seconds = window_seconds
self.client_requests: Dict[str, deque] = defaultdict(deque)
self.cleanup_interval = 300 # Clean up every 5 minutes
self.last_cleanup = time.time()
def get_client_ip(self, request: Request) -> str:
"""Extract client IP address"""
# Check for forwarded headers first (common in cloud deployments)
forwarded_for = request.headers.get("x-forwarded-for")
if forwarded_for:
return forwarded_for.split(",")[0].strip()
real_ip = request.headers.get("x-real-ip")
if real_ip:
return real_ip
# Fallback to direct client IP
return request.client.host if request.client else "unknown"
def cleanup_old_requests(self):
"""Remove old request timestamps to prevent memory leaks"""
current_time = time.time()
if current_time - self.last_cleanup < self.cleanup_interval:
return
cutoff_time = current_time - self.window_seconds
for client_ip in list(self.client_requests.keys()):
timestamps = self.client_requests[client_ip]
# Remove old timestamps
while timestamps and timestamps[0] < cutoff_time:
timestamps.popleft()
# Remove empty deques
if not timestamps:
del self.client_requests[client_ip]
self.last_cleanup = current_time
async def dispatch(self, request: Request, call_next):
"""Process rate limiting for incoming requests"""
# Skip rate limiting for health checks and docs
if request.url.path in ["/health", "/docs", "/openapi.json", "/"]:
response = await call_next(request)
return response
client_ip = self.get_client_ip(request)
current_time = time.time()
# Cleanup old requests periodically
self.cleanup_old_requests()
# Get request timestamps for this client
timestamps = self.client_requests[client_ip]
# Remove timestamps outside the current window
cutoff_time = current_time - self.window_seconds
while timestamps and timestamps[0] < cutoff_time:
timestamps.popleft()
# Check if client has exceeded rate limit
if len(timestamps) >= self.requests_per_window:
error_response = {
"error": {
"type": "rate_limit_exceeded",
"message": f"Rate limit exceeded. Maximum {self.requests_per_window} requests per {self.window_seconds} seconds.",
"details": {
"requests_made": len(timestamps),
"window_seconds": self.window_seconds,
"reset_time": int(timestamps[0] + self.window_seconds)
}
},
"timestamp": datetime.utcnow().isoformat() + "Z"
}
return JSONResponse(
status_code=429,
content=error_response,
headers={
"X-RateLimit-Limit": str(self.requests_per_window),
"X-RateLimit-Remaining": "0",
"X-RateLimit-Reset": str(int(timestamps[0] + self.window_seconds)),
"Retry-After": str(int(timestamps[0] + self.window_seconds - current_time))
}
)
# Add current request timestamp
timestamps.append(current_time)
# Process the request
response = await call_next(request)
# Add rate limit headers to response
remaining = max(0, self.requests_per_window - len(timestamps))
response.headers["X-RateLimit-Limit"] = str(self.requests_per_window)
response.headers["X-RateLimit-Remaining"] = str(remaining)
if timestamps:
response.headers["X-RateLimit-Reset"] = str(int(timestamps[0] + self.window_seconds))
return response
def create_error_response(error_type: str, message: str, status_code: int = 400, details: Optional[Dict] = None) -> JSONResponse:
"""Create standardized error response"""
error_content = {
"error": {
"type": error_type,
"message": message,
"details": details or {}
},
"timestamp": datetime.utcnow().isoformat() + "Z"
}
return JSONResponse(
status_code=status_code,
content=error_content
) |