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 typing import NamedTuple

from fastapi import HTTPException, Request, status

from app.core.config import settings

# key -> (window_start_ts, count, window_sec)
# 存每个 key 自己的 window_sec:_buckets 混着不同窗口(60s 广告 / 3600s 登录 / 86400s 日闸)的 key,
# GC 必须按各 key 自己的窗口判过期(见 [_purge_expired]),否则短窗口调用触发的 GC 会误删长窗口 key。
_buckets: dict[str, tuple[float, int, float]] = {}
_lock = threading.Lock()
_GC_THRESHOLD = 10000  # _buckets 超此阈值才顺手清过期 key(仿 sms.py;测试可 monkeypatch 调小强制每次扫)


def _purge_expired(now: float) -> None:
    """清过期 key(**仅在持有 _lock 时调用**)。按每个 key 自己存的 window_sec 判过期,而非调用方的窗口
    —— _buckets 是全局共享、混着 60s(广告)/3600s(登录)/86400s(日闸)不同窗口的 key;若用调用方窗口,
    高频的 60s 广告端点触发 GC 时会把本该活 3600s/86400s 的登录/日闸计数一并删掉,使其在规模上(超阈值才
    触发本清理)被反复清零而失效。仅在超阈值时扫,低频、开销可忽略。"""
    if len(_buckets) <= _GC_THRESHOLD:
        return
    for k in [k for k, (s, _, w) in _buckets.items() if now - s >= w]:
        _buckets.pop(k, None)


def _hit(key: str, limit: int, window_sec: float) -> bool:
    """记一次访问。返回 True=放行,False=超限。"""
    now = time.monotonic()
    with _lock:
        start, count, _ = _buckets.get(key, (now, 0, window_sec))
        if now - start >= window_sec:  # 窗口过期,重置
            start, count = now, 0
        count += 1
        _buckets[key] = (start, count, window_sec)
        _purge_expired(now)  # 顺手清过期 key(按各自窗口),防内存无限涨
    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,
        )


# ===================== 先判 / 后记(只按「成功」计数)=====================
# _hit 是原子「判+记」:一调用就 +1,适合登录爆破(失败尝试也要计)。但对「短信发码」这类
# **只想给成功动作计数**的场景不合适 —— 被单号冷却挡下的重发没真发、没烧钱,不该占额度。
# 故拆成 _peek(只判不记)+ _commit(只记):check_rate_limits 先判 → 动作 → 成功后 record。


class RateLimitRule(NamedTuple):
    """一条限流规则。scope 区分不同闸(不同 key 前缀);同一 (subject, IP) 在 window_sec
    内最多 limit 次,超限抛 429 用 detail 文案。

    (scope, window_sec) 成对绑在一条规则里 —— check(先判)与 record(计数)复用同一条,
    避免两处把窗口/scope 写歪导致 key 对不上。
    """

    scope: str
    limit: int
    window_sec: float
    detail: str = "操作过于频繁,请稍后再试"


def _peek(key: str, limit: int, window_sec: float) -> bool:
    """只读:当前窗口内是否还没到上限(count < limit)。**不改计数**。
    与 [_commit] 配对实现「先判后记」——只在动作成功后才 _commit。"""
    now = time.monotonic()
    with _lock:
        start, count, _ = _buckets.get(key, (now, 0, window_sec))
        if now - start >= window_sec:  # 窗口已过期 → 视作已重置(count 归零)
            count = 0
    return count < limit


def _commit(key: str, window_sec: float) -> None:
    """记一次访问(+1)。窗口过期则以本次为起点重置。仅在动作成功后调用。"""
    now = time.monotonic()
    with _lock:
        start, count, _ = _buckets.get(key, (now, 0, window_sec))
        if now - start >= window_sec:  # 窗口过期,重置
            start, count = now, 0
        _buckets[key] = (start, count + 1, window_sec)
        _purge_expired(now)  # 顺手清过期 key(按各自窗口,同 [_hit])


def check_rate_limits(request: Request, subject: str, rules: list[RateLimitRule]) -> None:
    """【先判】一组限流:任一规则已达上限即抛 429,且**不改计数**。

    配合 [record_rate_limits] 实现「只按成功计数」:先 check 所有闸(全未超才继续)→ 执行动作
    → 动作**成功后**再 record。动作被下游挡下(如短信单号冷却)、没真正发生时不 record → 不占额度。
    key = `scope:subject:client_ip`(与 [enforce_rate_limit] 同款)。
    """
    if not settings.RATE_LIMIT_ENABLED:
        return
    ip = _client_ip(request)
    for rule in rules:
        if not _peek(f"{rule.scope}:{subject}:{ip}", rule.limit, rule.window_sec):
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail=rule.detail,
            )


def record_rate_limits(request: Request, subject: str, rules: list[RateLimitRule]) -> None:
    """【记一次】一组限流(每条规则 +1)。仅在动作成功后调用,与 [check_rate_limits] 配对。

    ⚠️ check→动作→record 非原子:并发突发下计数可能略超 limit(每个在途请求各 +1)。对
    「防脚本/防轰炸」的安全网定位可接受;要精确配额需迁 Redis(见模块 docstring)。
    """
    if not settings.RATE_LIMIT_ENABLED:
        return
    ip = _client_ip(request)
    for rule in rules:
        _commit(f"{rule.scope}:{subject}:{ip}", rule.window_sec)