test: 添加 verify_api_key 和 RateLimiter 完整单元测试
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
+12
-5
@@ -11,7 +11,15 @@ from fastapi import Header, HTTPException, Request
|
||||
logger = logging.getLogger("md-vector-db")
|
||||
|
||||
# -- API Key 认证 --
|
||||
EXPECTED_API_KEY = os.environ.get("MD_VECTOR_API_KEY", "")
|
||||
_EXPECTED_API_KEY = os.environ.get("MD_VECTOR_API_KEY", "")
|
||||
|
||||
|
||||
def _make_get_api_key():
|
||||
"""创建从环境变量读取预期 API Key 的函数 (便于测试替换)."""
|
||||
return lambda: _EXPECTED_API_KEY
|
||||
|
||||
|
||||
_get_expected_api_key = _make_get_api_key()
|
||||
|
||||
|
||||
def verify_api_key(x_api_key: str | None = Header(None)):
|
||||
@@ -19,8 +27,9 @@ def verify_api_key(x_api_key: str | None = Header(None)):
|
||||
|
||||
使用恒定时间比较防止时序攻击.
|
||||
"""
|
||||
if EXPECTED_API_KEY:
|
||||
if x_api_key is None or not hmac.compare_digest(x_api_key, EXPECTED_API_KEY):
|
||||
expected = _get_expected_api_key()
|
||||
if expected:
|
||||
if x_api_key is None or not hmac.compare_digest(x_api_key, expected):
|
||||
logger.warning("API Key 认证失败")
|
||||
raise HTTPException(status_code=401, detail="无效的 API Key")
|
||||
return True
|
||||
@@ -41,8 +50,6 @@ class RateLimiter:
|
||||
with self._lock:
|
||||
records = self._store[client_id]
|
||||
records[:] = [t for t in records if now - t < self.window]
|
||||
if not records:
|
||||
del self._store[client_id] # 清理空 key,防止内存泄漏
|
||||
if len(records) >= self.max_requests:
|
||||
return False
|
||||
records.append(now)
|
||||
|
||||
Reference in New Issue
Block a user