58 lines
2.1 KiB
Python
58 lines
2.1 KiB
Python
"""限流配置。
|
|
|
|
链路是 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
|