zhaolei
16f134e328
核心功能: - 6 层检测管线: L0 结构分析 → L1 ClamAV → L2 文档深度检测 → L3 关键词匹配 → L4 YARA → L5 PE/二进制 - FastAPI + PostgreSQL + Redis + Docker Compose 架构 - 客户端 HMAC-SHA256 签名鉴权 + 限流 + 幂等性 - 控制台 UI (登录/Dashboard/诊断中心/策略管理/任务列表/测试中心/产品介绍页) - 三层诊断定时器 (error_scan 30s / db_aggregate 30s / container_check 5min) 并发与稳定性修复: - 诊断定时器 Redis 时间槽位锁 (多 worker 同一槽位只有一个执行) - JSONL 存储 fcntl 跨进程排他锁 + 容错解码 + 原子重写 - 移除 SQLite fallback, 统一 PostgreSQL - API_CLIENTS 脱敏为环境变量, .env 已 .gitignore
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
|