import asyncio import os import sys # ============================================================================== # 解决 Windows 下 psycopg 异步模式不支持 ProactorEventLoop 的问题 # 注意:这段代码必须放在应用的极早期,最好在导入 FastAPI 或初始化数据库之前 # ============================================================================== if sys.platform == "win32": # 把当前文件 (main.py) 的上一级的上一级 (即根目录 Yuxi) 加入到 sys.path sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) import time from collections import defaultdict, deque import uvicorn from fastapi import FastAPI, Request, status from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from starlette.middleware.base import BaseHTTPMiddleware from yuxi.external_systems.exceptions import ExternalSystemError from yuxi.scheduler.exceptions import SchedulerError from server.routers import router from server.utils.lifespan import lifespan from server.utils.auth_middleware import is_public_path from server.utils.common_utils import setup_logging from server.utils.access_log_middleware import AccessLogMiddleware # 设置日志配置 setup_logging() RATE_LIMIT_MAX_ATTEMPTS = 10 RATE_LIMIT_WINDOW_SECONDS = 60 RATE_LIMIT_ENDPOINTS = {("/api/auth/token", "POST")} # In-memory login attempt tracker to reduce brute-force exposure per worker _login_attempts: defaultdict[str, deque[float]] = defaultdict(deque) _attempt_lock = asyncio.Lock() app = FastAPI(lifespan=lifespan) # 所有业务接口统一挂载到 /api,具体分组在 server.routers 中集中注册。 app.include_router(router, prefix="/api") @app.exception_handler(ExternalSystemError) async def external_system_error_handler(_request: Request, exc: ExternalSystemError) -> JSONResponse: """外部系统限界上下文异常统一映射为 HTTP 响应。 - ``status_code`` 来自异常类。 - ``details`` 透传到响应体,供前端结构化处理。 - 429/503 且 ``details`` 含 ``retry_after`` 时,自动设置 ``Retry-After`` 响应头。 """ headers: dict[str, str] = {} if exc.status_code in (429, 503): retry_after = exc.details.get("retry_after") if retry_after is not None: headers["Retry-After"] = str(retry_after) return JSONResponse( status_code=exc.status_code, content={"detail": exc.message, "details": exc.details}, headers=headers, ) @app.exception_handler(SchedulerError) async def scheduler_error_handler(_request: Request, exc: SchedulerError) -> JSONResponse: """定时任务调度限界上下文异常统一映射为 HTTP 响应。 - ``status_code`` 来自异常类(400 / 404 / 409 / 413 / 500)。 - ``details`` 透传到响应体,供前端结构化处理。 """ return JSONResponse( status_code=exc.status_code, content={"detail": exc.message, "details": exc.details}, ) # CORS 设置 app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) def _extract_client_ip(request: Request) -> str: forwarded_for = request.headers.get("x-forwarded-for") if forwarded_for: return forwarded_for.split(",")[0].strip() if request.client: return request.client.host return "unknown" class LoginRateLimitMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next): normalized_path = request.url.path.rstrip("/") or "/" request_signature = (normalized_path, request.method.upper()) if request_signature in RATE_LIMIT_ENDPOINTS: client_ip = _extract_client_ip(request) now = time.monotonic() async with _attempt_lock: attempt_history = _login_attempts[client_ip] while attempt_history and now - attempt_history[0] > RATE_LIMIT_WINDOW_SECONDS: attempt_history.popleft() if len(attempt_history) >= RATE_LIMIT_MAX_ATTEMPTS: retry_after = int(max(1, RATE_LIMIT_WINDOW_SECONDS - (now - attempt_history[0]))) return JSONResponse( status_code=status.HTTP_429_TOO_MANY_REQUESTS, content={"detail": "登录尝试过于频繁,请稍后再试"}, headers={"Retry-After": str(retry_after)}, ) attempt_history.append(now) response = await call_next(request) if response.status_code < 400: async with _attempt_lock: _login_attempts.pop(client_ip, None) return response return await call_next(request) # 鉴权中间件 class AuthMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next): # 获取请求路径 path = request.url.path # 检查是否为公开路径,公开路径无需身份验证 if is_public_path(path): return await call_next(request) if not path.startswith("/api"): # 非API路径,可能是前端路由或静态资源 return await call_next(request) # # 提取Authorization头 # auth_header = request.headers.get("Authorization") # if not auth_header or not auth_header.startswith("Bearer "): # return JSONResponse( # status_code=status.HTTP_401_UNAUTHORIZED, # content={"detail": f"请先登录。Path: {path}"}, # headers={"WWW-Authenticate": "Bearer"} # ) # # 获取token # token = auth_header.split("Bearer ")[1] # # 添加token到请求状态,后续路由可以直接使用 # request.state.token = token # 继续处理请求 return await call_next(request) # 添加访问日志中间件(记录请求处理时间) app.add_middleware(AccessLogMiddleware) # 添加鉴权中间件 app.add_middleware(LoginRateLimitMiddleware) app.add_middleware(AuthMiddleware) if __name__ == "__main__": # uvicorn.run(app, host="0.0.0.0", port=5050, threads=10, workers=10, reload=True) uvicorn.run( "server.main:app", host="0.0.0.0", port=5050, reload=True, # 与 docker-compose 开发环境保持一致,避免 package 下代码变更不触发热重载。 reload_dirs=["server", "package"], )