80 lines
2.3 KiB
Python
80 lines
2.3 KiB
Python
from __future__ import annotations
|
||
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from typing import Dict, Tuple
|
||
|
||
from fastapi import HTTPException, Request
|
||
|
||
|
||
@dataclass
|
||
class FixedWindowRateLimiter:
|
||
"""
|
||
固定窗口限流（内存版）。
|
||
|
||
约束：
|
||
- 适用于单进程/单实例；多进程/多实例下不共享计数（V1 可接受）
|
||
- 窗口粒度：按分钟 bucket（window_seconds 建议为 60）
|
||
"""
|
||
|
||
limit: int
|
||
window_seconds: int
|
||
_counters: Dict[Tuple[str, int], int] = field(default_factory=dict)
|
||
_last_gc_bucket: int = 0
|
||
|
||
def _bucket(self, now_ts: float) -> int:
|
||
return int(now_ts // float(self.window_seconds))
|
||
|
||
def _gc(self, current_bucket: int) -> None:
|
||
# 每隔一段时间清理一次，避免 dict 无限增长（保留最近 3 个 bucket）
|
||
if self._last_gc_bucket == current_bucket:
|
||
return
|
||
self._last_gc_bucket = current_bucket
|
||
keep_from = current_bucket - 2
|
||
to_delete = [k for k in self._counters.keys() if k[1] < keep_from]
|
||
for k in to_delete:
|
||
self._counters.pop(k, None)
|
||
|
||
def allow(self, *, key: str, now_ts: float) -> None:
|
||
bucket = self._bucket(now_ts)
|
||
self._gc(bucket)
|
||
|
||
k = (str(key), int(bucket))
|
||
n = int(self._counters.get(k, 0)) + 1
|
||
self._counters[k] = n
|
||
if n > int(self.limit):
|
||
raise HTTPException(status_code=429, detail="rate_limited")
|
||
|
||
|
||
_reco_rate_limiter = FixedWindowRateLimiter(limit=10, window_seconds=60)
|
||
_push_rate_limiter = FixedWindowRateLimiter(limit=30, window_seconds=60)
|
||
|
||
|
||
async def rate_limit_reco_by_ip(request: Request) -> None:
|
||
"""
|
||
推荐接口限流：按 IP，1 分钟 10 次。
|
||
"""
|
||
|
||
ip = "unknown"
|
||
if request.client and request.client.host:
|
||
ip = str(request.client.host)
|
||
|
||
_reco_rate_limiter.allow(key=ip, now_ts=time.time())
|
||
|
||
|
||
async def rate_limit_push_by_ip(request: Request) -> None:
|
||
"""
|
||
推送相关接口限流：按 IP，1 分钟 30 次。
|
||
|
||
说明：
|
||
- register/preferences 等接口可能在客户端反复重试
|
||
- 本期先用内存固定窗口限流做基础保护
|
||
"""
|
||
|
||
ip = "unknown"
|
||
if request.client and request.client.host:
|
||
ip = str(request.client.host)
|
||
|
||
_push_rate_limiter.allow(key=ip, now_ts=time.time())
|
||
|