security: LLM 代理端点加失败限速(防 token 枚举)

窗口 300s 内鉴权失败 >= 20 次即拒绝,不再查库(避免枚举放大 DB 压力)。
限速键取 token 前 12 字符,内存不留全量密钥;条目超 5000 清理过期窗口。
This commit is contained in:
ymq 2026-08-25 14:47:24 +08:00
parent 7404ab1ecd
commit 9bb1c5ce65

View File

@ -36,6 +36,51 @@ S_REVOKED = "revoked"
DEFAULT_TTL_HOURS = 8
DEFAULT_MAX_CALLS = 500
# 失败限速:防 token 枚举。按来源标识token 前缀/调用方)计数,
# 窗口内失败次数超阈值则直接拒绝,不再查库(避免枚举放大 DB 压力)。
_FAIL_WINDOW_SEC = 300
_FAIL_MAX = 20
_fail_counters: dict = {}
def _fail_key(token):
"""限速键:取 token 前 12 字符(足够区分调用方,又不在内存留全量密钥)。"""
t = (token or '').strip()
if t.lower().startswith('bearer '):
t = t[7:].strip()
return t[:12] or 'anon'
def _rate_limited(token):
"""是否已被限速。顺带清理过期窗口。"""
import time
now = time.time()
k = _fail_key(token)
rec = _fail_counters.get(k)
if not rec:
return False
start, cnt = rec
if now - start > _FAIL_WINDOW_SEC:
_fail_counters.pop(k, None)
return False
return cnt >= _FAIL_MAX
def _record_fail(token):
"""记一次鉴权失败。"""
import time
now = time.time()
k = _fail_key(token)
rec = _fail_counters.get(k)
if not rec or now - rec[0] > _FAIL_WINDOW_SEC:
_fail_counters[k] = (now, 1)
else:
_fail_counters[k] = (rec[0], rec[1] + 1)
# 防内存膨胀:条目过多时清理已过期窗口
if len(_fail_counters) > 5000:
for kk in [x for x, v in _fail_counters.items() if now - v[0] > _FAIL_WINDOW_SEC]:
_fail_counters.pop(kk, None)
def _get_db():
db = DBPools()
@ -214,8 +259,13 @@ async def proxy_chat_completion(token, payload):
返回 (True, 上游响应 dict) (False, 错误信息)
api_key llm_bridge token 绑定的 org_id llm 表解析不下发给调用方
"""
# 失败限速前置:窗口内鉴权失败过多直接拒绝(防 token 枚举)
if _rate_limited(token):
logger.warning("proxy_chat_completion: 限速拒绝 key=%s", _fail_key(token))
return False, "鉴权失败次数过多,请稍后再试"
ok, info = await verify_llm_token(token)
if not ok:
_record_fail(token)
return False, info
if not isinstance(payload, dict) or not payload.get('messages'):
return False, "请求体缺 messages"