import json import time import uuid from dataclasses import replace from types import MappingProxyType from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.middleware.cors import CORSMiddleware from starlette.middleware.gzip import GZipMiddleware from starlette.requests import Request from starlette.responses import Response from starlette.types import ASGIApp from app.api.v1.module_system.params.service import ParamsService from app.common.response import ErrorResponse from app.config.setting import settings from app.core.exceptions import CustomException from app.core.logger import logger from app.core.request_context import RequestContext, clear_current_tenant, reset_correlation_id, set_correlation_id, set_current_tenant from app.core.security import decode_access_token from app.utils.ip_local_util import get_client_ip def _strip_bearer(authorization: str) -> str | None: """从 Authorization header 提取 token，非 Bearer 返回 None。""" v = authorization.strip() if v[:7].lower() == "bearer ": v = v[7:].strip() elif v[:6].lower() == "bearer": v = v[6:].strip() else: return None return v or None # 中间件配置的「安全默认值」：Redis 不可用 / 解析异常时启用，确保中间件行为可预测。 # 使用 MappingProxyType 防止任何地方误改导致跨请求污染。 _DEFAULT_CONFIG: MappingProxyType = MappingProxyType({ "demo_enable": False, "ip_white_list": (), "ip_black_list": (), "white_api_list_path": (), }) class CustomCORSMiddleware(CORSMiddleware): def __init__(self, app: ASGIApp) -> None: super().__init__( app, allow_origins=settings.ALLOW_ORIGINS, allow_methods=settings.ALLOW_METHODS, allow_headers=settings.ALLOW_HEADERS, allow_credentials=settings.ALLOW_CREDENTIALS, expose_headers=settings.CORS_EXPOSE_HEADERS, ) class RequestLogMiddleware(BaseHTTPMiddleware): """请求日志 & 演示模式拦截""" def __init__(self, app: ASGIApp) -> None: super().__init__(app) @staticmethod def _hydrate_session_id(request: Request) -> None: """从 ctx / JWT 中提取 session_id 并写入 ctx。""" ctx = getattr(request.state, "ctx", None) sid = ctx.session_id if ctx else None if not sid and ctx and ctx.jwt_user_info: sid = ctx.jwt_user_info.get("session_id") if not sid: token = _strip_bearer(request.headers.get("Authorization", "")) if token: try: payload = decode_access_token(token) sid = getattr(payload, "sub", None) if payload else None except Exception: sid = None if sid: request.state.ctx = replace(ctx or RequestContext(), session_id=sid) async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: start_time = time.time() self._hydrate_session_id(request) client_ip = get_client_ip(request) logger.info("请求: {} {} | client={}", request.method, request.url.path, client_ip or "unknown") try: path = request.url.path config = await self._load_config(request) is_blacklisted = bool(client_ip and client_ip in config["ip_black_list"]) in_demo = ( config["demo_enable"] and request.method != "GET" and (client_ip is None or client_ip not in config["ip_white_list"]) and not _is_path_whitelisted(path, config["white_api_list_path"]) ) if is_blacklisted or in_demo: logger.warning( "请求被拦截: {} {} | ip={} | 原因={}", request.method, path, client_ip, "IP黑名单" if is_blacklisted else "演示模式", ) return ErrorResponse( msg="IP已被黑名单" if is_blacklisted else "演示环境，禁止操作" ) response = await call_next(request) process_time = round(time.time() - start_time, 5) response.headers["X-Process-Time"] = str(process_time) logger.info("响应: {} | {:.1f}ms", response.status_code, process_time * 1000) return response except CustomException as e: logger.exception(f"中间件异常: {e!s}") return ErrorResponse(msg="系统异常，请联系管理员", data=str(e)) @staticmethod async def _load_config(request: Request) -> dict: """加载中间件配置（带 60 秒内存缓存），失败时返回全部默认值。""" redis = getattr(request.app.state, "redis", None) if not redis: return _DEFAULT_CONFIG try: tenant_id = await _extract_tenant_from_token(request) or 1 return await ParamsService.get_system_config_for_middleware(redis, tenant_id) except Exception: return _DEFAULT_CONFIG class CustomGZipMiddleware(GZipMiddleware): def __init__(self, app: ASGIApp) -> None: super().__init__(app, minimum_size=settings.GZIP_MIN_SIZE, compresslevel=settings.GZIP_COMPRESS_LEVEL) class CorrelationIdMiddleware(BaseHTTPMiddleware): def __init__(self, app: ASGIApp) -> None: self._header = "X-Correlation-ID" super().__init__(app) async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: cid = request.headers.get(self._header) or str(uuid.uuid4()) token = set_correlation_id(cid) try: response = await call_next(request) response.headers[self._header] = cid return response finally: reset_correlation_id(token) _TENANT_WHITELIST_PREFIXES = ("/docs", "/redoc", "/openapi.json", "/metrics", "/static/") _WHITELIST_ALL = ( "/api/v1/system/auth/login", "/api/v1/system/auth/captcha", "/api/v1/system/auth/refresh", "/api/v1/health", "/api/v1/common/health", ) + tuple(settings.TENANT_WHITELIST_PATHS) def _tenant_is_whitelisted(path: str) -> bool: """白名单路径：精确匹配公共接口，前缀匹配文档 / 静态资源。""" for prefix in (*_WHITELIST_ALL, *_TENANT_WHITELIST_PREFIXES): if path == prefix or path.startswith(prefix): return True return False async def _extract_tenant_from_token(request: Request) -> int | None: """从 JWT + Redis 会话解析租户 ID；结果挂到 request.state 上以便本请求内复用。 返回 None 表示未登录 / 会话过期，调用方应避免回退到平台租户（1）。 """ if hasattr(request.state, "tenant_id_resolved"): return request.state.tenant_id request.state.tenant_id_resolved = True request.state.tenant_id = None token = _strip_bearer(request.headers.get("Authorization", "")) if not token: return None try: payload = decode_access_token(token) if not payload or not hasattr(payload, "sub"): return None session_id = payload.sub redis = getattr(request.app.state, "redis", None) raw = await await_redis_get(redis, session_id) if redis else None user_info = json.loads(raw) if raw else None base = getattr(request.state, "ctx", None) or RequestContext() request.state.ctx = replace(base, jwt_payload=payload, jwt_user_info=user_info) if user_info and user_info.get("tenant_id"): request.state.tenant_id = int(user_info["tenant_id"]) except Exception: pass return request.state.tenant_id async def await_redis_get(redis, key: str) -> str | None: """获取用户会话；Redis 不可用时返回 None。""" from app.common.enums import RedisInitKeyConfig from app.core.redis_crud import RedisCURD try: return await RedisCURD(redis).get(f"{RedisInitKeyConfig.USER_SESSION.key}:{key}") except Exception: return None def _is_path_whitelisted(path: str, whitelist: list) -> bool: """精确匹配；``*`` 结尾表示前缀通配。""" for item in whitelist: if not isinstance(item, str) or not item: continue if item.endswith("*"): if path.startswith(item.rstrip("*")): return True elif path == item: return True return False class TenantMiddleware(BaseHTTPMiddleware): def __init__(self, app: ASGIApp) -> None: super().__init__(app) async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: if request.method == "OPTIONS" or _tenant_is_whitelisted(request.url.path): return await call_next(request) try: set_current_tenant(await _extract_tenant_from_token(request)) except Exception: logger.exception("租户中间件异常: path={}", request.url.path) try: return await call_next(request) finally: clear_current_tenant()