projectAIpopular/gateway/proxy/auth.py

"""鉴权与限流（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._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 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")