Archived
snowy-yang
4356c5f817
为什么：/api/task/heartbeat 同时出现在两处—— 1. guard() 开头的 _EXEMPT_PATHS 豁免表里（命中就直接 return，不进桶） 2. task_heartbeat 自己的 guard(bucket='hb', limit=60) 调用里 结果 guard 在进入桶之前就被豁免表放行，hb 桶的 60/分限额写了等于没写， 心跳端点实际处于无限制状态。 改法：把 /api/task/heartbeat 从 _EXEMPT_PATHS 移除，让它真正走到 hb 桶。 保留 /api/queue 与 /api/batch_progress 的豁免（它们被高频轮询且无自带桶）。 顺带修正上一轮的一处误判：原以为「手写滑动窗口在多 worker 下会放大限额」， 实测 Dockerfile / docker-compose.yml / entrypoint.sh 全是单进程 python3 web/main.py（uvicorn.run 无 --workers），进程内内存窗口是正确实现， 该顾虑不成立。
72 lines
2.7 KiB
Python
72 lines
2.7 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""公网限流——按 IP 滑动窗口。
|
|
|
|
限流桶集中在 core.state.RATE_BUCKETS/_RATE_LOCK，使所有端点共享同一批桶
|
|
（同一 IP 一律计入同一窗口），避免出现绕过限流的口子。
|
|
"""
|
|
import collections
|
|
import time
|
|
|
|
from fastapi import Request
|
|
|
|
from core.config import RESULT_RETENTION_SECONDS # noqa: F401 - 再导出，与 core.config 口径同源
|
|
from core.state import RATE_BUCKETS, RATE_LOCK
|
|
from web.errors import AppError
|
|
|
|
# 免于限流的端点（试听流/进度轮询）。
|
|
# 注意 /api/task/heartbeat 不在此列：它自带 hb 桶（60/分），此前的豁免表
|
|
# 让 guard 在进入桶之前就直接返回，限额写了等于没写——心跳端点是唯一
|
|
# 在豁免与限流两处同时出现的路径（/api/queue、/api/batch_progress 本就不调 guard）。
|
|
_EXEMPT_PREFIXES = ('/audio_stream',)
|
|
_EXEMPT_PATHS = {'/api/queue', '/api/batch_progress'}
|
|
|
|
|
|
def client_ip(request: Request) -> str:
|
|
"""客户端 IP：仅当直连方是回环（可信反代）时采信 XFF/Real-IP。
|
|
|
|
规则见 core.net（与原实现同源）：XFF 最右段优先，
|
|
最右段为回环时回落 X-Real-IP，全部无效回落直连地址。
|
|
"""
|
|
from core.net import is_loopback_ip, trusted_proxy_ip
|
|
remote = request.client.host if request.client else '-'
|
|
if remote and remote != '-' and is_loopback_ip(remote):
|
|
xff = request.headers.get('x-forwarded-for')
|
|
if xff:
|
|
ip = trusted_proxy_ip(xff)
|
|
if ip and not is_loopback_ip(ip):
|
|
return ip
|
|
xreal = request.headers.get('x-real-ip')
|
|
if xreal:
|
|
ip = trusted_proxy_ip(xreal)
|
|
if ip:
|
|
return ip
|
|
return remote or '-'
|
|
|
|
|
|
def rate_ok(key: str, limit: int, window: int) -> bool:
|
|
"""滑动窗口限流：True=放行。"""
|
|
now = time.time()
|
|
with RATE_LOCK:
|
|
dq = RATE_BUCKETS.get(key)
|
|
if dq is None:
|
|
if len(RATE_BUCKETS) > 8000:
|
|
for k in [k for k, v in RATE_BUCKETS.items() if not v or now - v[-1] > window]:
|
|
RATE_BUCKETS.pop(k, None)
|
|
dq = collections.deque()
|
|
RATE_BUCKETS[key] = dq
|
|
while dq and now - dq[0] > window:
|
|
dq.popleft()
|
|
if len(dq) >= limit:
|
|
return False
|
|
dq.append(now)
|
|
return True
|
|
|
|
|
|
async def guard(request: Request, *, bucket: str, limit: int, window: int) -> None:
|
|
"""限流闸：超限抛 429。豁免端点直接放行。"""
|
|
path = request.url.path
|
|
if path in _EXEMPT_PATHS or path.startswith(_EXEMPT_PREFIXES):
|
|
return
|
|
if not rate_ok(f'{bucket}:' + client_ip(request), limit, window):
|
|
raise AppError('请求过于频繁，请稍后再试', 429)
|