diff --git a/pipeline_service/llm_proxy.py b/pipeline_service/llm_proxy.py index 877691b..cc1c2a0 100644 --- a/pipeline_service/llm_proxy.py +++ b/pipeline_service/llm_proxy.py @@ -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"