fix: 增加请求前限额拦截(async_pre_call_hook),修复超限滞后击穿漏洞

This commit is contained in:
root
2026-08-10 08:07:04 +00:00
parent fbd2ab1ede
commit 6b51d74e72
+83 -1
View File
@@ -17,7 +17,11 @@ from urllib.parse import urlparse
try:
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.exceptions import ServiceUnavailableError, ContextWindowExceededError
from litellm.exceptions import (
ServiceUnavailableError,
ContextWindowExceededError,
RateLimitError,
)
LITELLM_AVAILABLE = True
except ImportError:
@@ -33,6 +37,9 @@ except ImportError:
class ContextWindowExceededError(Exception):
pass
class RateLimitError(Exception):
pass
# ==========================================
# 🛠️ Configuration & Constants
@@ -626,6 +633,55 @@ def startup_audit():
log(f"❌ Startup audit failed: {e}", level="ERROR")
def _check_request_block(model_name):
"""
请求前闸门:对外模型名 → 判定是否超限。
只拦截 MONITORED_MODEL_PAIRS 内的火山监控条目;newapi 等非监控渠道天然放行。
返回 (clean_model, usage, limit) 表示超限,None 表示放行。任何异常 fail-open。
"""
if not redis_client or not LITELLM_AVAILABLE:
return None
try:
from litellm.proxy.proxy_server import llm_router
if not llm_router or not llm_router.model_list:
return None
for entry in llm_router.model_list:
if entry.get("model_name") != model_name:
continue
# 已被限额踢出的条目跳过(路由层已匹配不到)
if "_daily_limit_backup" in entry:
continue
params = entry.get("litellm_params", {})
raw_model = params.get("model", "")
api_base = params.get("api_base", "")
clean = clean_model_name(raw_model)
h = generate_api_base_hash(api_base)
# 只拦已纳入监控的火山模型(模型 + 端点粒度)
if MONITORED_MODEL_PAIRS and (clean, h) not in MONITORED_MODEL_PAIRS:
continue
effective_limit = get_effective_limit(clean, api_base)
if effective_limit is None:
continue
usage_key = get_usage_key(clean, api_base)
val = redis_client.get(usage_key)
usage = float(val or 0)
if usage >= float(effective_limit):
return (clean, usage, effective_limit)
return None
except Exception as e:
log(f"❌ Pre-block check failed: {e}", level="ERROR")
return None
def quota_interceptor(kwargs, user_api_key_dict=None):
try:
# 1. 尝试恢复(定期)
@@ -646,6 +702,32 @@ class DailyLimitHandler(CustomLogger):
litellm.callbacks.append(self)
dprint("✅ DailyLimitHandler registered")
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
"""
请求前限额闸门:在 route_request 之前拒绝超限请求(429),
不进入 router → 不触发 num_retries / fallback 放大。
"""
if call_type not in ("completion", "acompletion", "text_completion"):
return None
if not data or not data.get("model"):
return None
blocked = _check_request_block(data["model"])
if not blocked:
return None
clean, usage, limit = blocked
log(
f"🚫 [PRE-BLOCK] {data['model']} daily limit reached ({int(usage):,}/{int(limit):,})",
level="WARNING",
)
raise RateLimitError(
message=f"Daily limit reached for model '{data['model']}' ({int(usage):,}/{int(limit):,})",
model=data["model"],
llm_provider="volcengine",
request_data=data,
)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self._record_usage(kwargs, response_obj)