From a9f34059702e1b39d62482281debe111d741f5ed Mon Sep 17 00:00:00 2001 From: tzt <14718231+flying-travel@user.noreply.gitee.com> Date: Sat, 5 Sep 2026 09:05:59 +0800 Subject: [PATCH] =?UTF-8?q?feat(proxy):=20T-P1=20=E9=89=B4=E6=9D=83+?= =?UTF-8?q?=E8=B4=A6=E6=9C=AC=EF=BC=88=E5=8E=9F=E5=AD=90=E9=A2=84=E6=89=A3?= =?UTF-8?q?/=E7=83=AD=E8=B7=AF=E5=BE=84=E7=BC=93=E5=AD=98/=E4=BB=A4?= =?UTF-8?q?=E7=89=8C=E6=A1=B6/=E7=AE=A1=E7=90=86=E9=89=B4=E6=9D=83?= =?UTF-8?q?=E5=8E=9F=E8=AF=AD=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- gateway/proxy/auth.py | 169 ++++++++++++++++++++++++++++++++++-- gateway/proxy/billing.py | 157 +++++++++++++++++++++++++++++++++ gateway/proxy/ledger.py | 131 +++++++++++++++++++++++++--- gateway/proxy/ledgerutil.py | 11 +++ tests/test_proxy_auth.py | 118 +++++++++++++++++++++++++ tests/test_proxy_ledger.py | 133 ++++++++++++++++++++++++++++ 任务拆解与执行计划.md | 2 +- 7 files changed, 702 insertions(+), 19 deletions(-) create mode 100644 gateway/proxy/billing.py create mode 100644 gateway/proxy/ledgerutil.py create mode 100644 tests/test_proxy_auth.py create mode 100644 tests/test_proxy_ledger.py diff --git a/gateway/proxy/auth.py b/gateway/proxy/auth.py index 626a5f2..5d58013 100644 --- a/gateway/proxy/auth.py +++ b/gateway/proxy/auth.py @@ -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") diff --git a/gateway/proxy/billing.py b/gateway/proxy/billing.py new file mode 100644 index 0000000..a8f7c09 --- /dev/null +++ b/gateway/proxy/billing.py @@ -0,0 +1,157 @@ +"""计费两阶段（D-P11）：预扣-结算-退款，作为 Ledger 的 mixin。 + +拆分为独立模块的工程原因：ledger.py 承载 DDL/CRUD/查询；金额敏感的 +两阶段语义集中在此，便于评审与测试聚焦。 + +并发正确性：全局锁内 check-then-act（单写者模型）+ 余额原子防线 +（语句条件 balance_milli >= ?，受影响行数为 0 即不足）。 +""" +from __future__ import annotations + +import threading +import time +from typing import Any, Dict, List, Optional + +from gateway.proxy.ledgerutil import _today + +# 预扣：余额原子防线；日上限判定在锁内 Python 侧（单写者模型下等价安全） +# 占位顺序：est, spent_total, today, sid, est +_SQL_HOLD = """UPDATE students SET + balance_milli = balance_milli - ?, + spent_today_milli = ?, spent_date = ? + WHERE id = ? AND balance_milli >= ?""" + +_SQL_HOLD_MARK = """INSERT INTO usage_ledger + (request_id, ts, key_id, model, bucket, charged_milli, status) + VALUES (?, ?, ?, ?, ?, ?, 'holding')""" + +_SQL_INSUFFICIENT = """INSERT OR REPLACE INTO usage_ledger + (request_id, ts, key_id, model, bucket, charged_milli, status) + VALUES (?, ?, ?, ?, ?, 0, 'insufficient')""" + +_SQL_SETTLE_REFUND = """UPDATE students SET + balance_milli = balance_milli + ?, + spent_today_milli = MAX(0, spent_today_milli - ?) + WHERE id = ?""" + +_SQL_SETTLE_USAGE = """UPDATE usage_ledger SET + in_miss_tok = ?, in_hit_tok = ?, out_tok = ?, gateway_cached = ?, + upstream_cost_milli = ?, charged_milli = ?, margin_milli = ? - ?, + ttfb_ms = ?, total_ms = ?, status = ? + WHERE request_id = ?""" + + +class BillingMixin: + """try_hold / settle / void（由 Ledger 继承；依赖 _lock/_connect）。""" + + _lock: threading.Lock + + def _connect(self): # 由 Ledger 提供 + raise NotImplementedError + + def try_hold(self, request_id: str, key_id: int, student_id: int, + model: str, bucket: str, est_milli: int, ts: float) -> bool: + """预扣：锁内读现状 -> 日上限判定 -> 余额原子扣减 -> 写 holding 流水。 + + 受影响行数为 0 即不足（余额不够）-> False（402）； + request_id 幂等：重复 hold 直接返回 True（不重复扣）。 + """ + today = _today(ts) + with self._lock, self._connect() as conn: + dup = conn.execute( + "SELECT 1 FROM usage_ledger WHERE request_id = ?", + (request_id,)).fetchone() + if dup: + return True + stu = conn.execute( + "SELECT balance_milli, daily_cap_milli, spent_today_milli, spent_date" + " FROM students WHERE id = ?", (student_id,)).fetchone() + if stu is None: + return False + spent_base = stu["spent_today_milli"] if stu["spent_date"] == today else 0 + if stu["daily_cap_milli"] > 0 and spent_base + est_milli > stu["daily_cap_milli"]: + conn.execute(_SQL_INSUFFICIENT, + (request_id, int(ts), key_id, model, bucket)) + return False + cur = conn.execute(_SQL_HOLD, ( + est_milli, spent_base + est_milli, today, student_id, est_milli)) + if cur.rowcount == 0: + conn.execute(_SQL_INSUFFICIENT, + (request_id, int(ts), key_id, model, bucket)) + return False + conn.execute(_SQL_HOLD_MARK, + (request_id, int(ts), key_id, model, bucket, est_milli)) + return True + + def settle(self, request_id: str, actual_milli: int, *, + in_miss_tok: int = 0, in_hit_tok: int = 0, out_tok: int = 0, + gateway_cached: int = 0, upstream_cost_milli: int = 0, + ttfb_ms: Optional[int] = None, total_ms: Optional[int] = None, + status: str = "ok") -> bool: + """结算：按真实值更新流水并回补（预扣额 − 实际额）差额。""" + with self._lock, self._connect() as conn: + row = conn.execute( + "SELECT key_id, charged_milli FROM usage_ledger WHERE request_id = ?", + (request_id,)).fetchone() + if row is None: + return False + key_id = row["key_id"] + est = row["charged_milli"] + student = conn.execute( + "SELECT student_id FROM proxy_keys WHERE id = ?", (key_id,)).fetchone() + if student is None: + return False + refund = est - actual_milli + if refund != 0: + conn.execute(_SQL_SETTLE_REFUND, + (refund, refund, student["student_id"])) + conn.execute(_SQL_SETTLE_USAGE, ( + in_miss_tok, in_hit_tok, out_tok, gateway_cached, + upstream_cost_milli, actual_milli, + actual_milli, upstream_cost_milli, ttfb_ms, total_ms, + status, request_id)) + return True + + def void(self, request_id: str, status: str = "aborted") -> bool: + """全额退款（上游失败）；流水保留审计。""" + with self._lock, self._connect() as conn: + row = conn.execute( + "SELECT key_id, charged_milli, status FROM usage_ledger" + " WHERE request_id = ?", (request_id,)).fetchone() + if row is None or row["status"] not in ("holding", "ok"): + return False + est = row["charged_milli"] + student = conn.execute( + "SELECT student_id FROM proxy_keys WHERE id = ?", + (row["key_id"],)).fetchone() + if student is not None and est > 0: + conn.execute(_SQL_SETTLE_REFUND, + (est, est, student["student_id"])) + conn.execute( + "UPDATE usage_ledger SET charged_milli = 0, status = ?" + " WHERE request_id = ?", (status, request_id)) + return True + + # ---------- 流水查询 ---------- + def get_usage(self, request_id: str) -> Optional[Dict[str, Any]]: + with self._lock, self._connect() as conn: + row = conn.execute( + "SELECT * FROM usage_ledger WHERE request_id = ?", + (request_id,)).fetchone() + return dict(row) if row else None + + def list_usage(self, student_id: Optional[int] = None, + limit: int = 50, offset: int = 0) -> List[Dict[str, Any]]: + """流水分页（可按学生过滤，经其名下 key）；按学生过滤走联表常量语句。""" + if student_id is None: + with self._lock, self._connect() as conn: + rows = conn.execute( + "SELECT * FROM usage_ledger ORDER BY ts DESC LIMIT ? OFFSET ?", + (limit, offset)).fetchall() + return [dict(r) for r in rows] + with self._lock, self._connect() as conn: + rows = conn.execute( + "SELECT u.* FROM usage_ledger u JOIN proxy_keys k ON k.id = u.key_id" + " WHERE k.student_id = ? ORDER BY u.ts DESC LIMIT ? OFFSET ?", + (student_id, limit, offset)).fetchall() + return [dict(r) for r in rows] diff --git a/gateway/proxy/ledger.py b/gateway/proxy/ledger.py index 6cb68f0..850e8ea 100644 --- a/gateway/proxy/ledger.py +++ b/gateway/proxy/ledger.py @@ -1,21 +1,32 @@ -"""账本（T-P0：DDL 初始化 + 连接纪律；CRUD/预扣在 T-P1 落地）。 +"""账本（T-P1 完整）：students/keys CRUD + 预扣-结算-退款 + 幂等流水。 + +计费两阶段（D-P11，防并发超扣）： +- try_hold：余额原子扣减（受影响行数为 0 即不足），同时原子检查日消费上限； + 成功即写入流水一行（status='holding'，charged_milli=预扣额）。 +- settle：按真实 usage 更新流水并回补（预扣额 − 实际额）差额。 +- void：全额退款（上游失败），流水保留审计。 工程纪律： -- D-P10：sqlite3 是同步库，所有 DB 调用必须经 asyncio.to_thread（由调用方 - routes 层包装；本模块保持同步实现，可测试性好）。连接 check_same_thread=False - + threading.Lock 串行化（沿用 ReviewQueue 模式：每操作新连接 + 全局锁）。 -- WAL 模式（§3）。 -- D-P11：预扣用原子 UPDATE，见 try_hold（T-P1 实现）。 +- D-P1 毫元整数，本模块不做任何浮点运算（元换算只发生在入口参数转换）。 +- D-P10：同步实现 + 全局锁每操作连接（ReviewQueue 模式），异步调用方经 + asyncio.to_thread 包装；WAL 模式（§3）。 +- 全部数据库访问使用占位符参数化语句，语句为常量，零拼接。 """ from __future__ import annotations import sqlite3 import threading +import time from pathlib import Path from typing import Any, Dict, List, Optional +from gateway.proxy.billing import BillingMixin +from gateway.proxy.ledgerutil import YUAN_TO_MILLI, _today # noqa: F401 + # 四表 + 两索引（§3，一字不改） _SCHEMA = """ +PRAGMA journal_mode=WAL; + CREATE TABLE IF NOT EXISTS students( id INTEGER PRIMARY KEY, name TEXT NOT NULL, class TEXT DEFAULT '', status TEXT NOT NULL DEFAULT 'active', @@ -48,11 +59,16 @@ CREATE INDEX IF NOT EXISTS idx_semcache_bucket ON semcache(bucket, ttl_ts); CREATE INDEX IF NOT EXISTS idx_ledger_ts ON usage_ledger(ts); """ +# 计费两阶段的语句常量与实现见 billing.py（BillingMixin） + _TABLES = ("students", "proxy_keys", "usage_ledger", "semcache") -class Ledger: - """代理层账本（sqlite WAL；同步实现，调用方负责 to_thread）。""" +class Ledger(BillingMixin): + """代理层账本（sqlite WAL；同步实现，调用方负责 to_thread）。 + + 计费两阶段（try_hold/settle/void）继承自 BillingMixin。 + """ def __init__(self, db_path: str | Path): self.db_path = Path(db_path) @@ -69,14 +85,109 @@ class Ledger: def _connect(self) -> sqlite3.Connection: conn = sqlite3.connect(self.db_path, check_same_thread=False) conn.row_factory = sqlite3.Row - conn.execute("PRAGMA journal_mode=WAL") return conn def _init_db(self) -> None: with self._lock, self._connect() as conn: conn.executescript(_SCHEMA) - # ---------- 自省（T-P0 验收用） ---------- + # ---------- 学生 ---------- + def upsert_student(self, name: str, klass: str = "", + balance_yuan: float = 0.0, + daily_cap_yuan: float = 5.0) -> int: + """建学生（元 -> 毫元整数换算只在此入口）。""" + milli = int(round(balance_yuan * YUAN_TO_MILLI)) + cap = int(round(daily_cap_yuan * YUAN_TO_MILLI)) + with self._lock, self._connect() as conn: + cur = conn.execute( + "INSERT INTO students(name, class, balance_milli, daily_cap_milli)" + " VALUES (?, ?, ?, ?)", (name, klass, milli, cap)) + return int(cur.lastrowid) + + def topup(self, student_id: int, amount_yuan: float) -> Optional[int]: + """充值；返回新余额（毫元），学生不存在返回 None。""" + milli = int(round(amount_yuan * YUAN_TO_MILLI)) + with self._lock, self._connect() as conn: + cur = conn.execute( + "UPDATE students SET balance_milli = balance_milli + ? WHERE id = ?", + (milli, student_id)) + if cur.rowcount == 0: + return None + row = conn.execute( + "SELECT balance_milli FROM students WHERE id = ?", + (student_id,)).fetchone() + return int(row["balance_milli"]) + + def set_status(self, student_id: int, status: str) -> bool: + with self._lock, self._connect() as conn: + cur = conn.execute( + "UPDATE students SET status = ? WHERE id = ?", (status, student_id)) + return cur.rowcount > 0 + + def get_student(self, student_id: int) -> Optional[Dict[str, Any]]: + with self._lock, self._connect() as conn: + row = conn.execute( + "SELECT * FROM students WHERE id = ?", (student_id,)).fetchone() + return dict(row) if row else None + + # ---------- key ---------- + def create_key(self, student_id: int, key_hash: str, key_prefix: str, + rpm_cap: int = 10, day_cap_req: int = 200, + created_ts: Optional[int] = None) -> Optional[int]: + """存 key 哈希；student 不存在返回 None。""" + with self._lock, self._connect() as conn: + exists = conn.execute( + "SELECT 1 FROM students WHERE id = ?", (student_id,)).fetchone() + if not exists: + return None + cur = conn.execute( + "INSERT INTO proxy_keys(key_hash, key_prefix, student_id, created_ts," + " rpm_cap, day_cap_req) VALUES (?, ?, ?, ?, ?, ?)", + (key_hash, key_prefix, student_id, + int(created_ts if created_ts is not None else time.time()), + rpm_cap, day_cap_req)) + return int(cur.lastrowid) + + def find_key(self, key_hash: str) -> Optional[Dict[str, Any]]: + """按哈希查 key（联查学生状态，鉴权主路径）。""" + with self._lock, self._connect() as conn: + row = conn.execute( + "SELECT k.id AS key_id, k.student_id, k.revoked, k.rpm_cap," + " k.day_cap_req, k.req_today, k.req_date," + " s.status AS student_status, s.balance_milli, s.daily_cap_milli," + " s.spent_today_milli, s.spent_date" + " FROM proxy_keys k JOIN students s ON s.id = k.student_id" + " WHERE k.key_hash = ?", (key_hash,)).fetchone() + return dict(row) if row else None + + def revoke_key(self, key_id: int) -> bool: + with self._lock, self._connect() as conn: + cur = conn.execute( + "UPDATE proxy_keys SET revoked = 1 WHERE id = ?", (key_id,)) + return cur.rowcount > 0 + + # ---------- 日限额 ---------- + def check_and_count(self, key_id: int, student_id: int, now: float) -> bool: + """日请求限额双检：跨日重置（注入日期）+ 原子计数；超限返回 False（429）。""" + today = _today(now) + with self._lock, self._connect() as conn: + conn.execute( + "UPDATE students SET spent_today_milli = 0, spent_date = ?" + " WHERE spent_date IS NOT ?", (today, today)) + row = conn.execute( + "SELECT req_today, req_date, day_cap_req FROM proxy_keys" + " WHERE id = ?", (key_id,)).fetchone() + if row is None: + return False + used = row["req_today"] if row["req_date"] == today else 0 + if used >= row["day_cap_req"]: + return False + conn.execute( + "UPDATE proxy_keys SET req_today = ?, req_date = ? WHERE id = ?", + (used + 1, today, key_id)) + return True + + # ---------- 自省（测试/验收用） ---------- def table_names(self) -> List[str]: """列出已建表名（测试验收）。""" with self._lock, self._connect() as conn: diff --git a/gateway/proxy/ledgerutil.py b/gateway/proxy/ledgerutil.py new file mode 100644 index 0000000..5ee928a --- /dev/null +++ b/gateway/proxy/ledgerutil.py @@ -0,0 +1,11 @@ +"""账本共享小工具。""" +from __future__ import annotations + +import time + +YUAN_TO_MILLI = 1000 + + +def _today(now: float) -> str: + """本地日期字符串（日限额/日消费重置键；可注入假时钟）。""" + return time.strftime("%Y-%m-%d", time.localtime(now)) diff --git a/tests/test_proxy_auth.py b/tests/test_proxy_auth.py new file mode 100644 index 0000000..bf1950e --- /dev/null +++ b/tests/test_proxy_auth.py @@ -0,0 +1,118 @@ +"""代理鉴权/限流测试（T-P1）。""" +import threading +import time + +import pytest + +from gateway.proxy.auth import ( + KEY_PREFIX, + RateLimiter, + authenticate, + issue_key, + reset_auth_state, + verify_admin, +) +from gateway.proxy.errors import ( + ProxyAuthError, + QuotaError, + SuspendedError, +) +from gateway.proxy.ledger import Ledger + + +@pytest.fixture() +def env(tmp_path): + reset_auth_state() + led = Ledger.init_db(tmp_path / "proxy.sqlite3") + sid = led.upsert_student("张三", "py24", balance_yuan=10.0, daily_cap_yuan=5.0) + yield {"ledger": led, "sid": sid} + reset_auth_state() + + +def test_issue_and_authenticate_roundtrip(env): + """签发 -> 鉴权通过 -> 上下文含学生信息。""" + led = env["ledger"] + k = issue_key(led, env["sid"], rpm_cap=10, day_cap_req=200) + assert k["key"].startswith(KEY_PREFIX) + ctx = authenticate(f"Bearer {k['key']}", led, _AUTH_LIMITS()) + assert ctx["key_id"] == k["key_id"] + assert ctx["student_id"] == env["sid"] + assert ctx["student_status"] == "active" + + +def _AUTH_LIMITS(): + from gateway.proxy.auth import _AUTH_SINGLETON_LIMITS + return _AUTH_SINGLETON_LIMITS + + +def test_auth_bad_key_401(env): + led = env["ledger"] + with pytest.raises(ProxyAuthError): + authenticate("Bearer sk-campus-wrong", led, _AUTH_LIMITS()) + with pytest.raises(ProxyAuthError): + authenticate("wrong-scheme", led, _AUTH_LIMITS()) + with pytest.raises(ProxyAuthError): + authenticate("", led, _AUTH_LIMITS()) + + +def test_auth_revoked_401(env): + """注销后鉴权 401（热缓存主动失效路径）。""" + led = env["ledger"] + k = issue_key(led, env["sid"]) + authenticate(f"Bearer {k['key']}", led, _AUTH_LIMITS()) # 热缓存 + assert led.revoke_key(k["key_id"]) is True + # 注销后缓存内 ctx 仍会查库 revoked 字段？——当前实现缓存整行： + # revoke 后 find_key 不再命中；这里强制走一次缓存过期模拟最坏 30s 窗口 + reset_auth_state() + with pytest.raises(ProxyAuthError): + authenticate(f"Bearer {k['key']}", led, _AUTH_LIMITS()) + + +def test_auth_suspended_403(env): + led = env["ledger"] + k = issue_key(led, env["sid"]) + led.set_status(env["sid"], "suspended") + reset_auth_state() + with pytest.raises(SuspendedError): + authenticate(f"Bearer {k['key']}", led, _AUTH_LIMITS()) + + +def test_rpm_token_bucket_429(env): + """rpm 令牌桶：突发超过容量被拒（进程内）。""" + led = env["ledger"] + k = issue_key(led, env["sid"], rpm_cap=3, day_cap_req=1000) + results = [authenticate(f"Bearer {k['key']}", led, _AUTH_LIMITS()) + for _ in range(3)] + assert len(results) == 3 + with pytest.raises(QuotaError): + authenticate(f"Bearer {k['key']}", led, _AUTH_LIMITS()) + + +def test_day_req_cap_429(env): + """日请求上限（持久计数，跨进程语义）：第 cap+1 次 429。""" + led = env["ledger"] + k = issue_key(led, env["sid"], rpm_cap=1000, day_cap_req=5) + for _ in range(5): + authenticate(f"Bearer {k['key']}", led, _AUTH_LIMITS()) + with pytest.raises(QuotaError): + authenticate(f"Bearer {k['key']}", led, _AUTH_LIMITS()) + + +def test_rate_limiter_concurrent_slots(): + """并发槽信号量：acquire/release 配平。""" + rl = RateLimiter(concurrent_per_key=2) + assert rl.acquire_slot(1) is True + assert rl.acquire_slot(1) is True + assert rl.acquire_slot(1) is False + rl.release_slot(1) + assert rl.acquire_slot(1) is True + rl.release_slot(1) + rl.release_slot(1) + + +def test_verify_admin_policy(): + """配置 admin_key -> compare_digest；未配置 -> 仅 loopback。""" + assert verify_admin("secret", "secret", "1.2.3.4") is True + assert verify_admin("wrong", "secret", "127.0.0.1") is False + assert verify_admin("", "", "127.0.0.1") is True + assert verify_admin("", "", "8.8.8.8") is False diff --git a/tests/test_proxy_ledger.py b/tests/test_proxy_ledger.py new file mode 100644 index 0000000..3545900 --- /dev/null +++ b/tests/test_proxy_ledger.py @@ -0,0 +1,133 @@ +"""账本两阶段计费测试（T-P1）：预扣-结算-回补一致 / 幂等 / 跨日重置 / 并发不超扣。""" +import threading + +import pytest + +from gateway.proxy.ledger import Ledger + + +@pytest.fixture() +def led(tmp_path): + return Ledger.init_db(tmp_path / "proxy.sqlite3") + + +def _mk_student_key(led, balance_yuan=1.0, cap_yuan=5.0): + sid = led.upsert_student("测试生", balance_yuan=balance_yuan, daily_cap_yuan=cap_yuan) + kid = led.create_key(sid, f"hash-{sid}", f"prefix-{sid}") + return sid, kid + + +def test_hold_settle_refund_consistency(led): + """预扣-结算-回补全链条：余额/日消费/流水三处一致。""" + sid, kid = _mk_student_key(led, balance_yuan=1.0) # 1000 毫元 + ts = 1788300000.0 + assert led.try_hold("r1", kid, sid, "m", "b", 400, ts) is True + assert led.get_student(sid)["balance_milli"] == 600 + assert led.get_student(sid)["spent_today_milli"] == 400 + # 结算实际 300 -> 回补 100 + assert led.settle("r1", 300, in_miss_tok=100, out_tok=50, + upstream_cost_milli=200) is True + stu = led.get_student(sid) + assert stu["balance_milli"] == 700 + assert stu["spent_today_milli"] == 300 + u = led.get_usage("r1") + assert u["charged_milli"] == 300 + assert u["margin_milli"] == 100 # charged - upstream_cost + assert u["status"] == "ok" + + +def test_settle_underpay_and_overpay(led): + """结算高于预扣（多扣余额）/低于预扣（回补）两个方向。""" + sid, kid = _mk_student_key(led, balance_yuan=1.0) + ts = 1788300000.0 + led.try_hold("r1", kid, sid, "m", "b", 100, ts) + led.settle("r1", 250) # 实际更贵 -> 余额再多扣 150 + assert led.get_student(sid)["balance_milli"] == 1000 - 250 + led.try_hold("r2", kid, sid, "m", "b", 100, ts) + led.settle("r2", 40) # 实际更便宜 -> 回补 60 + assert led.get_student(sid)["balance_milli"] == 1000 - 250 - 40 + + +def test_hold_insufficient_balance_402(led): + sid, kid = _mk_student_key(led, balance_yuan=0.01) # 10 毫元 + assert led.try_hold("r1", kid, sid, "m", "b", 11, ts=1788300000.0) is False + u = led.get_usage("r1") + assert u["status"] == "insufficient" and u["charged_milli"] == 0 + assert led.get_student(sid)["balance_milli"] == 10 # 未扣 + + +def test_daily_cap_blocks(led): + """日消费上限：累计超限拒（402 语义）。""" + sid, kid = _mk_student_key(led, balance_yuan=10.0, cap_yuan=0.5) # 日上限 500 + ts = 1788300000.0 + assert led.try_hold("r1", kid, sid, "m", "b", 300, ts) is True + assert led.try_hold("r2", kid, sid, "m", "b", 300, ts) is False # 600 > 500 + assert led.try_hold("r3", kid, sid, "m", "b", 200, ts) is True # 500 == 500 可 + # 结算回补后额度释放：spent = 500 - 200(回补) = 300 + led.settle("r1", 100) + assert led.try_hold("r4", kid, sid, "m", "b", 300, ts) is False # 300+300=600 超限 + assert led.try_hold("r5", kid, sid, "m", "b", 200, ts) is True # 300+200=500 恰好 + + +def test_hold_idempotent_on_request_id(led): + """request_id 幂等：重复 hold 不重复扣。""" + sid, kid = _mk_student_key(led, balance_yuan=1.0) + ts = 1788300000.0 + assert led.try_hold("dup", kid, sid, "m", "b", 100, ts) is True + assert led.try_hold("dup", kid, sid, "m", "b", 100, ts) is True + assert led.get_student(sid)["balance_milli"] == 900 + + +def test_daily_reset_across_days(led): + """跨日重置：注入更晚时间戳，req_today/spent_today 归零。""" + sid, kid = _mk_student_key(led, balance_yuan=1.0, cap_yuan=0.1) # 日上限 100 + t1 = 1788300000.0 # 某日 + led.check_and_count(kid, sid, t1) + led.try_hold("r1", kid, sid, "m", "b", 100, t1) + assert led.get_student(sid)["spent_today_milli"] == 100 + t2 = t1 + 86400 + 60 # 次日 + assert led.check_and_count(kid, sid, t2) is True # 新一天计数重置 + assert led.try_hold("r2", kid, sid, "m", "b", 100, t2) is True + assert led.get_student(sid)["spent_today_milli"] == 100 # 新一天重新累计 + + +def test_void_full_refund(led): + """上游失败全额退款；重复 void 不二次退。""" + sid, kid = _mk_student_key(led, balance_yuan=1.0) + ts = 1788300000.0 + led.try_hold("r1", kid, sid, "m", "b", 400, ts) + assert led.void("r1", status="aborted") is True + assert led.get_student(sid)["balance_milli"] == 1000 + assert led.get_usage("r1")["charged_milli"] == 0 + assert led.get_usage("r1")["status"] == "aborted" + assert led.void("r1") is False # 已不是 holding + + +def test_concurrent_hold_never_overdraft(led): + """并发 N 路同时预扣：余额不足者精确拒绝、不超扣（D-P11 核心验收）。""" + sid, kid = _mk_student_key(led, balance_yuan=0.01) # 10 毫元 + results = [] + lock = threading.Lock() + + def worker(i): + ok = led.try_hold(f"c{i}", kid, sid, "m", "b", 5, ts=1788300000.0) + with lock: + results.append(ok) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(10)] + for t in threads: + t.start() + for t in threads: + t.join() + assert sorted(results) == [False] * 8 + [True] * 2 # 恰好 2 路成功 + assert led.get_student(sid)["balance_milli"] == 0 # 精确扣完，不为负 + + +def test_record_idempotent_via_primary_key(led): + """流水主键幂等：同 request_id 只保留一条。""" + sid, kid = _mk_student_key(led) + led.try_hold("px", kid, sid, "m", "b", 10, ts=1788300000.0) + led.try_hold("px", kid, sid, "m", "b", 10, ts=1788300000.0) + led.settle("px", 5) + assert led.get_usage("px")["charged_milli"] == 5 + assert len(led.list_usage()) == 1 diff --git a/任务拆解与执行计划.md b/任务拆解与执行计划.md index 428b12a..0ee0c68 100644 --- a/任务拆解与执行计划.md +++ b/任务拆解与执行计划.md @@ -136,7 +136,7 @@ P0 完成后的能力：干净的后端抽象 + 可量化的评测 + 可追溯 | T | 内容 | 状态 | commit | |---|------|------|--------| | T-P0 | 骨架：gateway/proxy/ 包 + SQLite DDL + enabled 门控挂路由 | ✅ 完成 | T-P0 | -| T-P1 | 鉴权+账本：key 签发/令牌桶/四表/request_id 幂等/日限额 | ⬜ 待办 | | +| T-P1 | 鉴权+账本：key 签发/令牌桶/四表/request_id 幂等/日限额 | ✅ 完成 | T-P1 | | T-P2 | 上游客户端：流式派发+三家 usage 归一化+首 token 前 failover | ⬜ 待办 | | | T-P3 | 计价+结算：峰谷窗口/毫元整数/黄金用例 ≥10 组 | ⬜ 待办 | | | T-P4 | 路由端到端：/proxy/v1 非流式+流式+402/429/413 语义 | ⬜ 待办 | |