114 lines
3.8 KiB
Python
114 lines
3.8 KiB
Python
|
"""接口限流（Token bucket 滑窗）。
|
|||
|
|
|
|||
|
|
设计：
|
|||
|
|
- 客户端维度（按 client_id 或 IP），每分钟请求配额
|
|||
|
|
- 优先用 Redis（多 worker 共享计数），Redis 不可用时降级为内存字典
|
|||
|
|
- 返回 429 + Retry-After header
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import threading
|
|||
|
|
import time
|
|||
|
|
from typing import Optional
|
|||
|
|
|
|||
|
|
from fastapi import Depends, HTTPException, Request, status
|
|||
|
|
|
|||
|
|
# ---------- 内存降级模式 ----------
|
|||
|
|
# key -> (window_start_timestamp, count)
|
|||
|
|
_MEM_BUCKET: dict[str, tuple[float, int]] = {}
|
|||
|
|
_MEM_LOCK = threading.Lock()
|
|||
|
|
_MEM_WINDOW = 60 # 秒
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _get_redis():
|
|||
|
|
"""懒加载 Redis（避免模块 import 时没装直接崩）。"""
|
|||
|
|
try:
|
|||
|
|
import redis # type: ignore
|
|||
|
|
except Exception: # noqa: BLE001
|
|||
|
|
return None
|
|||
|
|
from app.config import get_settings
|
|||
|
|
try:
|
|||
|
|
return redis.from_url(get_settings().redis_url, decode_responses=True)
|
|||
|
|
except Exception: # noqa: BLE001
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _memory_check(key: str, limit: int) -> bool:
|
|||
|
|
"""内存滑窗：返回 True 表示允许，False 表示超限。"""
|
|||
|
|
now = time.time()
|
|||
|
|
with _MEM_LOCK:
|
|||
|
|
start_ts, count = _MEM_BUCKET.get(key, (now, 0.0))
|
|||
|
|
if now - start_ts > _MEM_WINDOW:
|
|||
|
|
# 窗口外，重建
|
|||
|
|
_MEM_BUCKET[key] = (now, 1)
|
|||
|
|
return True
|
|||
|
|
if count >= limit:
|
|||
|
|
return False
|
|||
|
|
_MEM_BUCKET[key] = (start_ts, count + 1)
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _redis_check(client_key: str, limit: int) -> bool:
|
|||
|
|
"""Redis 滑窗（INCR + EXPIRE）。"""
|
|||
|
|
r = _get_redis()
|
|||
|
|
if r is None:
|
|||
|
|
return _memory_check(client_key, limit)
|
|||
|
|
try:
|
|||
|
|
key = f"meshsafe:rl:{client_key}"
|
|||
|
|
pipe = r.pipeline()
|
|||
|
|
pipe.incr(key)
|
|||
|
|
pipe.expire(key, _MEM_WINDOW)
|
|||
|
|
results = pipe.execute()
|
|||
|
|
count = results[0]
|
|||
|
|
return count <= limit
|
|||
|
|
except Exception: # noqa: BLE001
|
|||
|
|
return _memory_check(client_key, limit)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class RateLimiter:
|
|||
|
|
"""限流依赖：返回一个 callable 给 Depends。
|
|||
|
|
|
|||
|
|
用法：
|
|||
|
|
@app.post("/scan/sync")
|
|||
|
|
async def scan_sync(..., limiter=Depends(RateLimiter(max_requests_per_min=60))):
|
|||
|
|
...
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, max_requests_per_min: int = 60, label: str = "scan") -> None:
|
|||
|
|
self.max_requests = max_requests_per_min
|
|||
|
|
self.label = label
|
|||
|
|
|
|||
|
|
async def __call__(self, request: Request) -> None:
|
|||
|
|
# 优先用 client_id（客户端签名带上的），退而用 IP
|
|||
|
|
ctx = getattr(request.state, "client", None)
|
|||
|
|
client_id = (ctx.client_id if ctx else None) or getattr(request.state, "client_id", None) or request.client.host
|
|||
|
|
key = f"{self.label}:{client_id}"
|
|||
|
|
if not _redis_check(key, self.max_requests):
|
|||
|
|
raise HTTPException(
|
|||
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|||
|
|
detail=f"Rate limit exceeded ({self.label}: {self.max_requests}/min)",
|
|||
|
|
headers={"Retry-After": str(_MEM_WINDOW)},
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def rate_limit(max_requests_per_min: int = 60, label: str = "scan"):
|
|||
|
|
"""更简洁的限流装饰器（函数依赖）。
|
|||
|
|
|
|||
|
|
用法：
|
|||
|
|
@app.post("/scan/sync")
|
|||
|
|
async def scan_sync(..., _ = Depends(rate_limit(60, "scans"))):
|
|||
|
|
...
|
|||
|
|
"""
|
|||
|
|
async def _check(request: Request) -> None:
|
|||
|
|
ctx = getattr(request.state, "client", None)
|
|||
|
|
client_id = (ctx.client_id if ctx else None) or getattr(request.state, "client_id", None) or request.client.host
|
|||
|
|
key = f"{label}:{client_id}"
|
|||
|
|
if not _redis_check(key, max_requests_per_min):
|
|||
|
|
raise HTTPException(
|
|||
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|||
|
|
detail=f"Rate limit exceeded ({label}: {max_requests_per_min}/min)",
|
|||
|
|
headers={"Retry-After": str(_MEM_WINDOW)},
|
|||
|
|
)
|
|||
|
|
return _check
|