zhuzihao
2edebfad83
Co-authored-by: zzhyyyyy <2685922758@qq.com> Reviewed-on: #78 Co-authored-by: zhuzihao <zhuzihao@wonderable.ai> Co-committed-by: zhuzihao <zhuzihao@wonderable.ai>
86 lines
3.0 KiB
Python
86 lines
3.0 KiB
Python
"""轻量内存限流(固定窗口计数)。
|
|
|
|
定位:防脚本刷接口的**安全网**,不是精确配额。局限——进程内存、单 worker 有效、重启清零、
|
|
多 worker/多机不共享。上线放量后应换 Redis 后端 + 按用户维度。当前单 worker uvicorn 够用。
|
|
|
|
用法:把 `rate_limit(limit, window_sec, scope)` 作为路由依赖挂上,默认按客户端 IP 计数。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
|
|
from fastapi import HTTPException, Request, status
|
|
|
|
from app.core.config import settings
|
|
|
|
# key -> (window_start_ts, count)
|
|
_buckets: dict[str, tuple[float, int]] = {}
|
|
_lock = threading.Lock()
|
|
|
|
|
|
def _hit(key: str, limit: int, window_sec: float) -> bool:
|
|
"""记一次访问。返回 True=放行,False=超限。"""
|
|
now = time.monotonic()
|
|
with _lock:
|
|
start, count = _buckets.get(key, (now, 0))
|
|
if now - start >= window_sec: # 窗口过期,重置
|
|
start, count = now, 0
|
|
count += 1
|
|
_buckets[key] = (start, count)
|
|
# 顺手清理过期 key,防内存无限涨(低频访问足够)
|
|
if len(_buckets) > 10000:
|
|
for k in [k for k, (s, _) in _buckets.items() if now - s >= window_sec]:
|
|
_buckets.pop(k, None)
|
|
return count <= limit
|
|
|
|
|
|
def _client_ip(request: Request) -> str:
|
|
# nginx 反代后真实 IP 在 X-Forwarded-For 首段;直连用 request.client
|
|
xff = request.headers.get("x-forwarded-for")
|
|
if xff:
|
|
return xff.split(",")[0].strip()
|
|
return request.client.host if request.client else "unknown"
|
|
|
|
|
|
def rate_limit(limit: int, window_sec: float, scope: str):
|
|
"""生成一个 FastAPI 依赖:同一 IP 在 window_sec 内对该 scope 超过 limit 次 → 429。"""
|
|
|
|
def _dep(request: Request) -> None:
|
|
if not settings.RATE_LIMIT_ENABLED:
|
|
return
|
|
key = f"{scope}:{_client_ip(request)}"
|
|
if not _hit(key, limit, window_sec):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail="操作过于频繁,请稍后再试",
|
|
)
|
|
|
|
return _dep
|
|
|
|
|
|
def enforce_rate_limit(
|
|
request: Request,
|
|
scope: str,
|
|
subject: str,
|
|
limit: int,
|
|
window_sec: float,
|
|
*,
|
|
detail: str = "操作过于频繁,请稍后再试",
|
|
) -> None:
|
|
"""在路由内部手动限流,按 (subject, 客户端 IP) 计数。
|
|
|
|
用于限流 key 需要请求体字段(如手机号)、Depends 阶段还拿不到 body 时
|
|
—— 此时无法用 [rate_limit] 依赖,改在 handler 解析完 body 后调用本函数。
|
|
key = `scope:subject:client_ip`;同一 (subject, IP) 在 window_sec 内超过 limit 次 → 抛 429。
|
|
受 [settings.RATE_LIMIT_ENABLED] 总开关控制(与 [rate_limit] 一致)。
|
|
"""
|
|
if not settings.RATE_LIMIT_ENABLED:
|
|
return
|
|
key = f"{scope}:{subject}:{_client_ip(request)}"
|
|
if not _hit(key, limit, window_sec):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail=detail,
|
|
)
|