Xiaoxia AI
90 lines
2.6 KiB
Python
90 lines
2.6 KiB
Python
"""
|
|
请求日志中间件
|
|
"""
|
|
|
|
import logging
|
|
import time
|
|
|
|
from fastapi import Request
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RequestLoggingMiddleware(BaseHTTPMiddleware):
|
|
"""请求日志中间件"""
|
|
|
|
async def dispatch(self, request: Request, call_next):
|
|
# 记录请求开始时间
|
|
start_time = time.time()
|
|
|
|
# 记录请求信息
|
|
logger.info(f"Request: {request.method} {request.url.path}")
|
|
|
|
# 处理请求
|
|
response = await call_next(request)
|
|
|
|
# 计算处理时间
|
|
process_time = time.time() - start_time
|
|
|
|
# 记录响应信息
|
|
logger.info(
|
|
f"Response: {request.method} {request.url.path} " f"status={response.status_code} time={process_time:.3f}s"
|
|
)
|
|
|
|
# 添加响应头
|
|
response.headers["X-Process-Time"] = str(process_time)
|
|
|
|
return response
|
|
|
|
|
|
class RateLimitMiddleware(BaseHTTPMiddleware):
|
|
"""简单的速率限制中间件（基于内存）"""
|
|
|
|
def __init__(self, app, max_requests: int = 100, window_seconds: int = 60):
|
|
super().__init__(app)
|
|
self.max_requests = max_requests
|
|
self.window_seconds = window_seconds
|
|
self.requests = {} # {ip: [(timestamp, ...)]}
|
|
|
|
async def dispatch(self, request: Request, call_next):
|
|
# 获取客户端 IP
|
|
client_ip = request.client.host
|
|
current_time = time.time()
|
|
|
|
# 清理过期记录
|
|
if client_ip in self.requests:
|
|
self.requests[client_ip] = [
|
|
ts for ts in self.requests[client_ip] if current_time - ts < self.window_seconds
|
|
]
|
|
|
|
# 检查速率限制
|
|
request_count = len(self.requests.get(client_ip, []))
|
|
|
|
if request_count >= self.max_requests:
|
|
from fastapi.responses import JSONResponse
|
|
|
|
return JSONResponse(
|
|
status_code=429,
|
|
content={
|
|
"error": {
|
|
"code": "RATE_LIMIT_EXCEEDED",
|
|
"message": f"Too many requests. Limit: {self.max_requests} per {self.window_seconds}s",
|
|
}
|
|
},
|
|
)
|
|
|
|
# 记录请求
|
|
if client_ip not in self.requests:
|
|
self.requests[client_ip] = []
|
|
self.requests[client_ip].append(current_time)
|
|
|
|
# 处理请求
|
|
response = await call_next(request)
|
|
|
|
# 添加速率限制信息到响应头
|
|
response.headers["X-RateLimit-Limit"] = str(self.max_requests)
|
|
response.headers["X-RateLimit-Remaining"] = str(self.max_requests - len(self.requests[client_ip]))
|
|
|
|
return response
|