feat(proxy): T-P1 鉴权+账本（原子预扣/热路径缓存/令牌桶/管理鉴权原语）
- ledger.py：students/keys CRUD + 日限额 check_and_count（跨日重置，注入时钟）+ 流水分页；计费两阶段拆分至 billing.py（BillingMixin，评审聚焦）： try_hold 锁内读现状->日上限判定->余额原子防线->holding 流水（request_id 幂等）、 settle 按真实值回补（预扣-实际）差额、void 全额退款（上游失败） - auth.py：issue_key（sk-campus- 前缀，明文只返回一次，库存 sha256+前缀）、 authenticate 校验链（形态/注销/停用/日额/rpm，热路径 LRU TTL30s D-P10）、 RateLimiter（令牌桶 rpm + per-key 并发信号量，D-P9 单进程）、 verify_admin（hmac.compare_digest；未配置仅 loopback） - 测试 +17：签发/错key/注销/停用/403/rpm 429/日额 429/并发槽/admin 策略 + 预扣-结算-回补一致/双向差额/余额不足/日上限/幂等/跨日重置/void/并发10路不超扣 - 全量 339 passed
This commit is contained in:
tzt
+161
-8
@@ -1,22 +1,175 @@
|
||||
"""鉴权与限流（T-P1 落地；本文件先立签名）。"""
|
||||
"""鉴权与限流（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：明文只返回一次（T-P1 实现）。"""
|
||||
raise NotImplementedError("T-P1")
|
||||
"""签发学生代理 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]}
|
||||
|
||||
|
||||
def authenticate(authorization: str, ledger, limits, now: float) -> Dict[str, Any]:
|
||||
"""校验 Bearer key -> 学生上下文（T-P1 实现；带热路径缓存）。"""
|
||||
raise NotImplementedError("T-P1")
|
||||
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 + per-key 并发信号量（T-P1 实现）。"""
|
||||
"""令牌桶（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._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:
|
||||
raise NotImplementedError("T-P1")
|
||||
"""令牌桶放行判定（容量 = 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 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")
|
||||
|
||||
Reference in New Issue
Block a user