shaguabijia-app-server/app/core/ratelimit.py

"""轻量内存限流(固定窗口计数)。

定位:防脚本刷接口的**安全网**,不是精确配额。局限——进程内存、单 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