mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 01:17:46 +08:00
fawney19
1d9c77522a
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
487 lines
18 KiB
Python
487 lines
18 KiB
Python
"""
|
||
统一的插件中间件（纯 ASGI 实现）
|
||
负责协调所有插件的调用
|
||
|
||
注意：使用纯 ASGI middleware 而非 BaseHTTPMiddleware，
|
||
以避免 Starlette 已知的流式响应兼容性问题。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
from typing import TYPE_CHECKING
|
||
|
||
from starlette.requests import Request
|
||
|
||
if TYPE_CHECKING:
|
||
from sqlalchemy.orm import Session
|
||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||
|
||
from src.config import config
|
||
from src.core.logger import logger
|
||
from src.plugins.manager import get_plugin_manager
|
||
from src.plugins.rate_limit.base import RateLimitResult
|
||
|
||
|
||
class PluginMiddleware:
|
||
"""
|
||
统一的插件调用中间件（纯 ASGI 实现）
|
||
|
||
职责:
|
||
- 性能监控
|
||
- 限流控制 (可选)
|
||
- 数据库会话生命周期管理
|
||
|
||
注意: 认证由各路由通过 Depends() 显式声明，不在中间件层处理
|
||
|
||
为什么使用纯 ASGI 而非 BaseHTTPMiddleware:
|
||
- BaseHTTPMiddleware 会缓冲整个响应体，对流式响应不友好
|
||
- BaseHTTPMiddleware 与 StreamingResponse 存在已知兼容性问题
|
||
- 纯 ASGI 可以直接透传流式响应，无额外开销
|
||
"""
|
||
|
||
_NOTIFICATION_CACHE_TTL = 30.0 # 通知模块开关缓存 TTL (秒)
|
||
|
||
def __init__(self, app: ASGIApp) -> None:
|
||
self.app = app
|
||
self.plugin_manager = get_plugin_manager()
|
||
self._notification_enabled_cache: bool | None = None
|
||
self._notification_cache_expires: float = 0.0
|
||
|
||
# 从配置读取速率限制值
|
||
self.public_api_rate_limit = config.public_api_rate_limit
|
||
|
||
# 完全跳过限流的路径（静态资源、文档等）
|
||
self.skip_rate_limit_paths = [
|
||
"/health",
|
||
"/readyz",
|
||
"/docs",
|
||
"/redoc",
|
||
"/openapi.json",
|
||
"/favicon.ico",
|
||
"/static/",
|
||
"/assets/",
|
||
"/api/admin/", # 管理后台已有JWT认证，不需要额外限流
|
||
"/api/auth/", # 认证端点（由路由层的 IPRateLimiter 处理）
|
||
"/api/users/", # 用户端点
|
||
"/api/monitoring/", # 监控端点
|
||
]
|
||
|
||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||
"""ASGI 入口点"""
|
||
if scope["type"] != "http":
|
||
# 非 HTTP 请求（如 WebSocket）直接透传
|
||
await self.app(scope, receive, send)
|
||
return
|
||
|
||
# 构建 Request 对象以便复用现有逻辑
|
||
request = Request(scope, receive, send)
|
||
|
||
# 记录请求开始时间
|
||
start_time = time.time()
|
||
|
||
# 设置 request.state 属性
|
||
# 注意：Starlette 的 Request 对象总是有 state 属性（State 实例）
|
||
request.state.request_id = request.headers.get("x-request-id", "")
|
||
request.state.start_time = start_time
|
||
# 标记：若请求过程中通过 Depends(get_db) 创建了会话，则由本中间件统一管理其生命周期
|
||
request.state.db_managed_by_middleware = True
|
||
|
||
# 1. 限流检查（在调用下游之前）
|
||
rate_limit_result = await self._call_rate_limit_plugins(request)
|
||
if rate_limit_result and not rate_limit_result.allowed:
|
||
# 限流触发，返回429
|
||
await self._send_rate_limit_response(send, rate_limit_result)
|
||
return
|
||
|
||
# 2. 预处理插件调用
|
||
await self._call_pre_request_plugins(request)
|
||
|
||
# 用于捕获响应状态码
|
||
response_status_code: int = 0
|
||
|
||
async def send_wrapper(message: Message) -> None:
|
||
nonlocal response_status_code
|
||
|
||
if message["type"] == "http.response.start":
|
||
response_status_code = message.get("status", 0)
|
||
await self._maybe_release_streaming_db_session(request, message)
|
||
|
||
await send(message)
|
||
|
||
# 3. 调用下游应用
|
||
exception_occurred: Exception | None = None
|
||
try:
|
||
await self.app(scope, receive, send_wrapper)
|
||
except Exception as e:
|
||
exception_occurred = e
|
||
# 错误处理插件调用
|
||
await self._call_error_plugins(request, e, start_time)
|
||
raise
|
||
finally:
|
||
# 4. 数据库会话清理（无论成功与否）
|
||
await self._cleanup_db_session(request, exception_occurred)
|
||
|
||
# 5. 后处理插件调用（仅在成功时）
|
||
if not exception_occurred and response_status_code > 0:
|
||
await self._call_post_request_plugins(request, response_status_code, start_time)
|
||
|
||
async def _send_rate_limit_response(self, send: Send, result: RateLimitResult) -> None:
|
||
"""发送 429 限流响应"""
|
||
import json
|
||
|
||
body = json.dumps(
|
||
{
|
||
"type": "error",
|
||
"error": {
|
||
"type": "rate_limit_error",
|
||
"message": result.message or "Rate limit exceeded",
|
||
},
|
||
}
|
||
).encode("utf-8")
|
||
|
||
headers = [(b"content-type", b"application/json")]
|
||
if result.headers:
|
||
for key, value in result.headers.items():
|
||
headers.append((key.lower().encode(), str(value).encode()))
|
||
|
||
await send(
|
||
{
|
||
"type": "http.response.start",
|
||
"status": 429,
|
||
"headers": headers,
|
||
}
|
||
)
|
||
await send(
|
||
{
|
||
"type": "http.response.body",
|
||
"body": body,
|
||
}
|
||
)
|
||
|
||
def _finalize_db_session(
|
||
self,
|
||
db: Session,
|
||
*,
|
||
should_commit: bool,
|
||
should_rollback: bool,
|
||
log_prefix: str = "",
|
||
) -> None:
|
||
"""统一的数据库会话清理逻辑
|
||
|
||
Args:
|
||
db: SQLAlchemy 会话
|
||
should_commit: 是否需要提交（仅当 should_rollback=False 时生效）
|
||
should_rollback: 是否需要回滚（优先于 commit）
|
||
log_prefix: 日志前缀，用于区分调用场景
|
||
"""
|
||
try:
|
||
if should_rollback:
|
||
try:
|
||
db.rollback()
|
||
except Exception as rollback_error:
|
||
logger.debug(f"{log_prefix}回滚事务时出错（可忽略）: {rollback_error}")
|
||
elif should_commit:
|
||
try:
|
||
db.commit()
|
||
except Exception as commit_error:
|
||
logger.error(f"{log_prefix}提交失败: {commit_error}")
|
||
try:
|
||
db.rollback()
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
try:
|
||
db.close()
|
||
except Exception as close_error:
|
||
logger.debug(f"{log_prefix}关闭数据库连接时出错（可忽略）: {close_error}")
|
||
|
||
async def _cleanup_db_session(self, request: Request, exception: Exception | None) -> None:
|
||
"""清理数据库会话
|
||
|
||
事务策略：
|
||
- 如果 request.state.tx_committed_by_route 为 True，说明路由已自行提交，中间件只负责 close
|
||
- 否则由中间件统一 commit/rollback
|
||
|
||
这避免了双重提交的问题，同时保持向后兼容。
|
||
"""
|
||
from sqlalchemy.orm import Session
|
||
|
||
if getattr(request.state, "db_released_early", False):
|
||
return
|
||
if not getattr(request.state, "db_managed_by_middleware", False):
|
||
return
|
||
|
||
db = getattr(request.state, "db", None)
|
||
if not isinstance(db, Session):
|
||
return
|
||
|
||
tx_committed_by_route = getattr(request.state, "tx_committed_by_route", False)
|
||
self._finalize_db_session(
|
||
db,
|
||
should_commit=not tx_committed_by_route and exception is None,
|
||
should_rollback=exception is not None,
|
||
)
|
||
|
||
async def _maybe_release_streaming_db_session(self, request: Request, message: Message) -> None:
|
||
"""在 SSE 响应开始时提前释放请求级 DB session。"""
|
||
if getattr(request.state, "db_released_early", False):
|
||
return
|
||
if not getattr(request.state, "db_managed_by_middleware", False):
|
||
return
|
||
|
||
headers = message.get("headers") or []
|
||
content_type = None
|
||
for key, value in headers:
|
||
if key.lower() == b"content-type":
|
||
content_type = value.decode("utf-8", errors="ignore").lower()
|
||
break
|
||
|
||
if not content_type or "text/event-stream" not in content_type:
|
||
return
|
||
|
||
from sqlalchemy.orm import Session
|
||
|
||
db = getattr(request.state, "db", None)
|
||
if not isinstance(db, Session):
|
||
return
|
||
|
||
tx_committed_by_route = getattr(request.state, "tx_committed_by_route", False)
|
||
self._finalize_db_session(
|
||
db,
|
||
should_commit=not tx_committed_by_route,
|
||
should_rollback=False,
|
||
log_prefix="流式响应提前",
|
||
)
|
||
request.state.db = None
|
||
request.state.db_released_early = True
|
||
|
||
def _get_client_ip(self, request: Request) -> str:
|
||
"""
|
||
获取客户端 IP 地址，支持代理头
|
||
|
||
优先级：X-Real-IP > X-Forwarded-For > 直连 IP
|
||
X-Real-IP 由最外层 Nginx 设置，最可靠
|
||
"""
|
||
# 优先检查 X-Real-IP（由最外层 Nginx 设置，最可靠）
|
||
real_ip = request.headers.get("x-real-ip")
|
||
if real_ip:
|
||
return real_ip.strip()
|
||
|
||
# 检查 X-Forwarded-For，取第一个 IP（原始客户端）
|
||
forwarded_for = request.headers.get("x-forwarded-for")
|
||
if forwarded_for:
|
||
ips = [ip.strip() for ip in forwarded_for.split(",") if ip.strip()]
|
||
if ips:
|
||
return ips[0]
|
||
|
||
# 回退到直连 IP
|
||
if request.client:
|
||
return request.client.host
|
||
|
||
return "unknown"
|
||
|
||
async def _get_rate_limit_key_and_config(
|
||
self, request: Request
|
||
) -> tuple[str | None, int | None]:
|
||
"""
|
||
获取速率限制的key和配置
|
||
|
||
策略说明:
|
||
- /api/public/* 端点: 使用服务器级别 IP 限制
|
||
- /api/admin/* 端点: 跳过（在 skip_rate_limit_paths 中跳过）
|
||
- /api/auth/* 端点: 跳过（由路由层的 IPRateLimiter 处理）
|
||
|
||
Returns:
|
||
(key, rate_limit_value) - key用于标识限制对象，rate_limit_value是限制值
|
||
"""
|
||
path = request.url.path
|
||
|
||
# /api/public/* 端点: 使用服务器级别 IP 地址作为限制 key
|
||
if path.startswith("/api/public/"):
|
||
client_ip = self._get_client_ip(request)
|
||
key = f"public_ip:{client_ip}"
|
||
rate_limit = self.public_api_rate_limit
|
||
request.state.rate_limit_key_type = "public_ip"
|
||
request.state.rate_limit_value = rate_limit
|
||
return key, rate_limit
|
||
|
||
# 其他端点不应用速率限制（或已在 skip_rate_limit_paths 中跳过）
|
||
return None, None
|
||
|
||
async def _call_rate_limit_plugins(self, request: Request) -> RateLimitResult | None:
|
||
"""调用限流插件"""
|
||
|
||
# 跳过不需要限流的路径（支持前缀匹配）
|
||
for skip_path in self.skip_rate_limit_paths:
|
||
if request.url.path == skip_path or request.url.path.startswith(skip_path):
|
||
return None
|
||
|
||
# 获取限流插件
|
||
rate_limit_plugin = self.plugin_manager.get_plugin("rate_limit")
|
||
if not rate_limit_plugin or not rate_limit_plugin.enabled:
|
||
# 如果没有限流插件，允许通过
|
||
return None
|
||
|
||
# 获取速率限制的 key 和配置
|
||
key, rate_limit_value = await self._get_rate_limit_key_and_config(request)
|
||
if not key:
|
||
# 不需要限流的端点（如未分类路径），静默跳过
|
||
return None
|
||
|
||
try:
|
||
# 检查速率限制，传入数据库配置的限制值
|
||
result = await rate_limit_plugin.check_limit(
|
||
key=key,
|
||
endpoint=request.url.path,
|
||
method=request.method,
|
||
rate_limit=rate_limit_value, # 传入配置的限制值
|
||
)
|
||
# 类型检查：确保返回的是RateLimitResult类型
|
||
if isinstance(result, RateLimitResult):
|
||
# 如果检查通过，实际消耗令牌
|
||
if result.allowed:
|
||
await rate_limit_plugin.consume(
|
||
key=key,
|
||
amount=1,
|
||
rate_limit=rate_limit_value,
|
||
)
|
||
else:
|
||
# 限流触发，记录日志
|
||
logger.warning(
|
||
"速率限制触发: {}",
|
||
getattr(request.state, "rate_limit_key_type", "unknown"),
|
||
)
|
||
return result
|
||
return None
|
||
except ConnectionError as e:
|
||
# Redis 连接错误：根据配置决定
|
||
logger.warning(f"Rate limit connection error: {e}")
|
||
if config.rate_limit_fail_open:
|
||
return None
|
||
else:
|
||
return RateLimitResult(
|
||
allowed=False,
|
||
remaining=0,
|
||
retry_after=30,
|
||
message="Rate limit service unavailable",
|
||
)
|
||
except TimeoutError as e:
|
||
# 超时错误：可能是负载过高，根据配置决定
|
||
logger.warning(f"Rate limit timeout: {e}")
|
||
if config.rate_limit_fail_open:
|
||
return None
|
||
else:
|
||
return RateLimitResult(
|
||
allowed=False, remaining=0, retry_after=30, message="Rate limit service timeout"
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Rate limit error: {type(e).__name__}: {e}")
|
||
# 其他异常：根据配置决定
|
||
if config.rate_limit_fail_open:
|
||
# fail-open: 异常时放行请求（优先可用性）
|
||
return None
|
||
else:
|
||
# fail-close: 异常时拒绝请求（优先安全性）
|
||
return RateLimitResult(
|
||
allowed=False, remaining=0, retry_after=60, message="Rate limit service error"
|
||
)
|
||
|
||
async def _call_pre_request_plugins(self, request: Request) -> None:
|
||
"""调用请求前的插件（当前保留扩展点）"""
|
||
pass
|
||
|
||
async def _call_post_request_plugins(
|
||
self, request: Request, status_code: int, start_time: float
|
||
) -> None:
|
||
"""调用请求后的插件"""
|
||
|
||
duration = time.time() - start_time
|
||
|
||
# 监控插件 - 记录指标
|
||
monitor_plugin = self.plugin_manager.get_plugin("monitor")
|
||
if monitor_plugin and monitor_plugin.enabled:
|
||
try:
|
||
# 使用路由模板而非实际路径，避免动态段导致 Prometheus 标签基数爆炸
|
||
route = request.scope.get("route")
|
||
endpoint_label = route.path if route and hasattr(route, "path") else "unknown"
|
||
|
||
monitor_labels = {
|
||
"method": request.method,
|
||
"endpoint": endpoint_label,
|
||
"status": str(status_code),
|
||
"status_class": f"{status_code // 100}xx",
|
||
}
|
||
|
||
# 记录请求计数
|
||
await monitor_plugin.increment(
|
||
"http_requests_total",
|
||
labels=monitor_labels,
|
||
)
|
||
|
||
# 记录请求时长
|
||
await monitor_plugin.timing(
|
||
"http_request_duration",
|
||
duration,
|
||
labels=monitor_labels,
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Monitor plugin failed: {e}")
|
||
|
||
async def _is_notification_email_module_enabled(self, request: Request) -> bool:
|
||
"""检查通知邮件模块是否启用（带内存缓存，避免 5xx 雪崩时放大 DB 压力）。"""
|
||
now = time.time()
|
||
if self._notification_enabled_cache is not None and now < self._notification_cache_expires:
|
||
return self._notification_enabled_cache
|
||
|
||
from sqlalchemy.orm import Session
|
||
|
||
from src.database import create_session
|
||
from src.services.system.config import SystemConfigService
|
||
|
||
config_key = "module.notification_email.enabled"
|
||
try:
|
||
request_db = getattr(request.state, "db", None)
|
||
if isinstance(request_db, Session):
|
||
result = bool(SystemConfigService.get_config(request_db, config_key, default=False))
|
||
else:
|
||
db = create_session()
|
||
try:
|
||
result = bool(SystemConfigService.get_config(db, config_key, default=False))
|
||
finally:
|
||
db.close()
|
||
except Exception as e:
|
||
logger.warning("读取通知邮件模块开关失败: {}", e)
|
||
return False
|
||
|
||
self._notification_enabled_cache = result
|
||
self._notification_cache_expires = now + self._NOTIFICATION_CACHE_TTL
|
||
return result
|
||
|
||
async def _call_error_plugins(
|
||
self, request: Request, error: Exception, start_time: float
|
||
) -> None:
|
||
"""调用错误处理插件"""
|
||
from fastapi import HTTPException
|
||
|
||
duration = time.time() - start_time
|
||
|
||
# 通知插件 - 发送严重错误通知
|
||
if not isinstance(error, HTTPException) or error.status_code >= 500:
|
||
if not await self._is_notification_email_module_enabled(request):
|
||
return
|
||
|
||
notification_plugin = self.plugin_manager.get_plugin("notification")
|
||
if notification_plugin and notification_plugin.enabled:
|
||
try:
|
||
await notification_plugin.send_error(
|
||
error=error,
|
||
context={
|
||
"endpoint": f"{request.method} {request.url.path}",
|
||
"request_id": getattr(request.state, "request_id", ""),
|
||
"duration": duration,
|
||
},
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Notification plugin failed: {e}")
|