tzt
b2b7a64cb3
- auth.RateLimiter.allow_tokens：TPM 令牌桶（容量=tpm_cap，连续回流 tpm/60 每秒， 与 RPM 桶共用 per-key 锁、判定相互独立；tpm_cap<=0 不限） - ProxyConfig.limits.tpm_per_key（默认 60000，0=不限）；chat_completions 在并发槽 之后做 TPM 预扣判定（est=字符/3 与计费同口径），拒绝时释放并发槽并回 429 （文案区分并发/RPM/TPM 三种原因） - semcache.stats 增 misses 与 hit_rate；/admin/stats 新增 semcache 段透出 （entries/hits_exact/hits_semantic/misses/hit_rate，与账本 h_g 互补） - 新增 tests/test_tpm_limiter.py 6 项（容量/回流/禁用/双桶独立/key 隔离/命中率）
198 lines
7.9 KiB
Python
198 lines
7.9 KiB
Python
"""鉴权与限流（T-P1）：key 签发/校验/注销 + 令牌桶 + 并发信号量 + 热路径缓存。
|
||
|
||
工程要点：
|
||
- D-P8：明文 key 只在 issue_key 返回一次；库中仅存 sha256 哈希 + 前 12 位前缀（展示用）。
|
||
- D-P10 热路径：哈希 -> (key_id, student 上下文) 的进程内 LRU，TTL 30s，
|
||
注销/停用后写失效（revoke 时主动失效，最坏 30s 内仍可能命中旧缓存——
|
||
注销语义取"最终生效"，符合校园场景）。
|
||
- 限流：令牌桶（rpm，按 key），日请求上限走 ledger.check_and_count（持久），
|
||
并发上限 per-key 信号量（concurrent_per_key，进程内，D-P9 单进程前提）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import hmac
|
||
import secrets
|
||
import threading
|
||
import time
|
||
from collections import OrderedDict
|
||
from typing import Any, Dict, Optional
|
||
|
||
from gateway.proxy.errors import ProxyAuthError, QuotaError, SuspendedError
|
||
from gateway.proxy.ledgerutil import _today
|
||
|
||
KEY_PREFIX = "sk-campus-"
|
||
_AUTH_TTL = 30.0
|
||
_AUTH_CACHE_MAX = 4096
|
||
|
||
|
||
def _hash_key(plaintext: str) -> str:
|
||
return hashlib.sha256(plaintext.encode("utf-8")).hexdigest()
|
||
|
||
|
||
def issue_key(ledger, student_id: int, rpm_cap: Optional[int] = None,
|
||
day_cap_req: Optional[int] = None) -> Dict[str, Any]:
|
||
"""签发学生代理 key：明文只此一次返回；库中仅存哈希与前缀。
|
||
|
||
返回 {key_id, key(明文), prefix}；学生不存在时 ledger.create_key 返回 None
|
||
-> 抛 ProxyAuthError。
|
||
"""
|
||
plaintext = KEY_PREFIX + secrets.token_urlsafe(24)
|
||
kid = ledger.create_key(
|
||
student_id, _hash_key(plaintext), plaintext[:12],
|
||
rpm_cap=rpm_cap if rpm_cap is not None else 10,
|
||
day_cap_req=day_cap_req if day_cap_req is not None else 200)
|
||
if kid is None:
|
||
raise ProxyAuthError(f"学生不存在: {student_id}")
|
||
return {"key_id": kid, "key": plaintext, "prefix": plaintext[:12]}
|
||
|
||
|
||
class _AuthCache:
|
||
"""热路径鉴权缓存（LRU + TTL 30s；D-P10）。"""
|
||
|
||
def __init__(self, ttl: float = _AUTH_TTL, maxsize: int = _AUTH_CACHE_MAX):
|
||
self._ttl = ttl
|
||
self._max = maxsize
|
||
self._data: "OrderedDict[str, tuple[float, Any]]" = OrderedDict()
|
||
self._lock = threading.Lock()
|
||
|
||
def get(self, key_hash: str) -> Optional[Any]:
|
||
now = time.monotonic()
|
||
with self._lock:
|
||
item = self._data.get(key_hash)
|
||
if item is None:
|
||
return None
|
||
ts, ctx = item
|
||
if now - ts > self._ttl:
|
||
self._data.pop(key_hash, None)
|
||
return None
|
||
self._data.move_to_end(key_hash)
|
||
return ctx
|
||
|
||
def put(self, key_hash: str, ctx: Any) -> None:
|
||
with self._lock:
|
||
self._data[key_hash] = (time.monotonic(), ctx)
|
||
self._data.move_to_end(key_hash)
|
||
while len(self._data) > self._max:
|
||
self._data.popitem(last=False)
|
||
|
||
def invalidate(self, key_hash: str) -> None:
|
||
with self._lock:
|
||
self._data.pop(key_hash, None)
|
||
|
||
|
||
class RateLimiter:
|
||
"""令牌桶（rpm，按 key，进程内）+ per-key 并发信号量。
|
||
|
||
D-P9：全部为进程内状态，uvicorn workers=1 是正确性前提。
|
||
"""
|
||
|
||
def __init__(self, concurrent_per_key: int = 2):
|
||
self._tokens: Dict[int, tuple[float, float]] = {} # key_id -> (tokens, last_ts)
|
||
self._tpm: Dict[int, tuple[float, float]] = {} # key_id -> (tokens, last_ts)
|
||
self._locks: Dict[int, threading.Lock] = {}
|
||
self._sems: Dict[int, threading.Semaphore] = {}
|
||
self._global = threading.Lock()
|
||
self._concurrent = max(1, int(concurrent_per_key))
|
||
|
||
def allow(self, key_id: int, rpm_cap: int) -> bool:
|
||
"""令牌桶放行判定（容量 = rpm_cap，速率 = rpm/60 每秒）。"""
|
||
now = time.monotonic()
|
||
with self._global:
|
||
lock = self._locks.setdefault(key_id, threading.Lock())
|
||
with lock:
|
||
tokens, last = self._tokens.get(key_id, (float(rpm_cap), now))
|
||
tokens = min(float(rpm_cap), tokens + (now - last) * (rpm_cap / 60.0))
|
||
if tokens < 1.0:
|
||
self._tokens[key_id] = (tokens, now)
|
||
return False
|
||
self._tokens[key_id] = (tokens - 1.0, now)
|
||
return True
|
||
|
||
def allow_tokens(self, key_id: int, est_tokens: int, tpm_cap: int) -> bool:
|
||
"""TPM 令牌桶放行判定（T-X12，采纳 enterprise-ai-gateway 双桶设计）。
|
||
|
||
容量 = tpm_cap，回流速率 = tpm/60 每秒（连续时间回流补充）；
|
||
est_tokens 为本次预扣估算（字符/3，与预扣计费同口径）。
|
||
tpm_cap <= 0 视为不限流；与 RPM 桶共用 per-key 锁，判定相互独立。
|
||
"""
|
||
if tpm_cap <= 0:
|
||
return True
|
||
now = time.monotonic()
|
||
with self._global:
|
||
lock = self._locks.setdefault(key_id, threading.Lock())
|
||
with lock:
|
||
tokens, last = self._tpm.get(key_id, (float(tpm_cap), now))
|
||
tokens = min(float(tpm_cap), tokens + (now - last) * (tpm_cap / 60.0))
|
||
if tokens < est_tokens:
|
||
self._tpm[key_id] = (tokens, now)
|
||
return False
|
||
self._tpm[key_id] = (tokens - est_tokens, now)
|
||
return True
|
||
|
||
def acquire_slot(self, key_id: int) -> bool:
|
||
"""并发槽（非阻塞）；返回 False = 超并发上限（429）。"""
|
||
with self._global:
|
||
sem = self._sems.setdefault(
|
||
key_id, threading.Semaphore(self._concurrent))
|
||
return sem.acquire(blocking=False)
|
||
|
||
def release_slot(self, key_id: int) -> None:
|
||
sem = self._sems.get(key_id)
|
||
if sem is not None:
|
||
sem.release()
|
||
|
||
|
||
def authenticate(authorization: str, ledger, limiter: RateLimiter,
|
||
now: Optional[float] = None) -> Dict[str, Any]:
|
||
"""校验 Bearer key -> 学生/key 上下文（§6 签名）。
|
||
|
||
校验链（任一失败即短路与对应异常）：
|
||
形态(401) -> 哈希存在且未注销(401) -> 学生 active(403)
|
||
-> 日请求上限(429, 持久计数) -> rpm 令牌桶(429, 进程内)
|
||
返回 {key_id, student_id, rpm_cap, day_cap_req, balance_milli, ...}。
|
||
"""
|
||
now_ts = now if now is not None else time.time()
|
||
token = (authorization or "").strip()
|
||
if not token.lower().startswith("bearer "):
|
||
raise ProxyAuthError("缺少 Bearer 凭据")
|
||
plaintext = token[7:].strip()
|
||
if not plaintext.startswith(KEY_PREFIX):
|
||
raise ProxyAuthError("无效 key")
|
||
key_hash = _hash_key(plaintext)
|
||
|
||
ctx = _AUTH_SINGLETON.get(key_hash)
|
||
if ctx is None:
|
||
row = ledger.find_key(key_hash)
|
||
if row is None or row["revoked"]:
|
||
raise ProxyAuthError("无效或已注销的 key")
|
||
if row["student_status"] != "active":
|
||
raise SuspendedError("学生账户已停用")
|
||
ctx = row
|
||
_AUTH_SINGLETON.put(key_hash, ctx)
|
||
|
||
if not ledger.check_and_count(ctx["key_id"], ctx["student_id"], now_ts):
|
||
raise QuotaError("已达当日请求上限")
|
||
if not _AUTH_SINGLETON_LIMITS.allow(ctx["key_id"], int(ctx["rpm_cap"])):
|
||
raise QuotaError("请求过于频繁（rpm）")
|
||
return dict(ctx)
|
||
|
||
|
||
# 进程内单例（热缓存 + 限流器；D-P9 单进程）
|
||
_AUTH_SINGLETON = _AuthCache()
|
||
_AUTH_SINGLETON_LIMITS = RateLimiter(concurrent_per_key=2)
|
||
|
||
|
||
def reset_auth_state() -> None:
|
||
"""测试用：清空热缓存与限流器。"""
|
||
global _AUTH_SINGLETON, _AUTH_SINGLETON_LIMITS
|
||
_AUTH_SINGLETON = _AuthCache()
|
||
_AUTH_SINGLETON_LIMITS = RateLimiter(concurrent_per_key=2)
|
||
|
||
|
||
def verify_admin(requested_key: str, admin_key: str, client_host: str = "") -> bool:
|
||
"""管理面鉴权（§5.2）：hmac.compare_digest 防时序；未配置 admin_key 时仅 loopback。"""
|
||
if admin_key:
|
||
return hmac.compare_digest(str(requested_key or ""), str(admin_key))
|
||
return client_host in ("127.0.0.1", "::1", "localhost", "testclient")
|