"""请求日志和限流中间件""" import logging import re import time from typing import Optional from fastapi import Request from starlette.middleware.base import BaseHTTPMiddleware logger = logging.getLogger(__name__) # 敏感参数名称模式(不区分大小写) SENSITIVE_PARAM_PATTERNS = re.compile( r"^(password|passwd|pwd|token|secret|key|authorization|auth|api_key|" # noqa: E501 r"apikey|access_token|refresh_token|accesstoken|refreshtoken|" # noqa: E501 r"session_id|sessionid|sid|cookie|csrf|xsrf|bearer)$", # noqa: E501 re.IGNORECASE, ) def filter_sensitive_params(query_string: Optional[str]) -> Optional[str]: """ 过滤 query string 中的敏感参数 Args: query_string: 原始 query string,例如 "name=xxx&password=secret&token=abc" Returns: 过滤后的 query string,敏感参数的值被替换为 "***" 如果 query_string 为空或 None,返回原始值 """ if not query_string: return query_string if query_string.startswith("?"): query_string = query_string[1:] if not query_string: return query_string parts = query_string.split("&") filtered_parts = [] for part in parts: if "=" in part: key, value = part.split("=", 1) if SENSITIVE_PARAM_PATTERNS.match(key): filtered_parts.append(f"{key}=***") else: filtered_parts.append(part) else: # 没有 = 的参数,保留原样 filtered_parts.append(part) return "&".join(filtered_parts) class RequestLoggingMiddleware(BaseHTTPMiddleware): """请求日志中间件""" async def dispatch(self, request: Request, call_next): # 记录请求开始时间 start_time = time.time() # 过滤 query string 中的敏感参数 raw_query = str(request.url.query) if request.url.query else "" safe_query = filter_sensitive_params(raw_query) # 记录请求信息(不包含敏感参数) if safe_query: logger.info(f"Request: {request.method} {request.url.path}?{safe_query}") # noqa: E501 else: 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): """基于 IP 的简单限流中间件 Args: app: ASGI 应用 max_requests: 窗口期内最大请求数 window_seconds: 时间窗口(秒) paths: 限流的路径列表,None 表示所有路径 """ def __init__(self, app, max_requests: int = 100, window_seconds: int = 60, paths: list[str] | None = None): super().__init__(app) self.max_requests = max_requests self.window_seconds = window_seconds self.paths = set(paths) if paths else None self.requests: dict[str, list[float]] = {} async def dispatch(self, request: Request, call_next): # 如果配置了路径过滤,只对指定路径限流 if self.paths is not None and request.url.path not in self.paths: return await call_next(request) # 获取客户端 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": ( # noqa: E501 f"Too many requests. Limit: " f"{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