第一次提交

app/ratelimit.py

@@ -0,0 +1,57 @@
+"""限流配置。
+
+链路是 Caddy → frp → Traefik → 本服务，真实 IP 依赖 uvicorn 的
+--proxy-headers（见 entrypoint.sh），否则所有请求会被算作同一来源。
+
+两种维度：
+- 认证端点按 IP 限流：防止密码爆破（此时还没有用户身份）。
+- AI 端点按用户限流：AI 调用花的是用户自己配的 API 额度，且多人可能
+  共用出口 IP，按 IP 限会互相误伤。
+"""
+from slowapi import Limiter
+from slowapi.util import get_remote_address
+from starlette.requests import Request
+
+from app.config import settings
+
+
+def _user_or_ip(request: Request) -> str:
+    """优先用 JWT 的 sub 作为限流键，回退到 IP。
+
+    这里只做轻量解析，不校验签名——限流键不是安全边界，真正的鉴权
+    仍由 get_current_user 负责。解析失败就退回 IP。
+    """
+    auth = request.headers.get("authorization", "")
+    if auth.lower().startswith("bearer "):
+        token = auth[7:]
+        try:
+            import jwt
+
+            payload = jwt.decode(token, options={"verify_signature": False})
+            sub = payload.get("sub")
+            if sub:
+                return f"user:{sub}"
+        except Exception:  # noqa: BLE001 - 解析失败就按 IP 限
+            pass
+    return f"ip:{get_remote_address(request)}"
+
+
+# 默认按 IP 限流（认证端点用：此时还没有用户身份）。
+# AI / 上传端点在装饰器上传 key_func=user_key 改为按用户限流。
+limiter = Limiter(
+    key_func=get_remote_address,
+    enabled=settings.rate_limit_enabled,
+    # 内存存储：单副本部署够用（SQLite 本来也只能单副本）
+    storage_uri="memory://",
+    # 必须按 endpoint 计数，不能用默认的 "url"：像 /attempts/{id}/judge
+    # 这种路径带 ID，按 url 计数会让每次请求各占一个桶，限流形同虚设。
+    key_style="endpoint",
+)
+
+# 供 AI / 上传端点使用的按用户限流键
+user_key = _user_or_ip
+
+# 各类端点的限流策略，集中在这里方便调整
+AUTH_LIMIT = settings.rate_limit_auth
+AI_LIMIT = settings.rate_limit_ai
+UPLOAD_LIMIT = settings.rate_limit_upload