Blame for app/cache.py

"""内存缓存与限流计数，减少对数据库的依赖。

key 缓存：首次校验后把 key 的校验信息与配额缓存在内存，
后续请求命中缓存即不查库；后台改 key 时清缓存。
RPM/TPM：用内存滑动窗口计数，不查 request_logs 表。
"""
from __future__ import annotations

import asyncio
import hashlib
import hmac
import time
from dataclasses import dataclass
from typing import Dict, List


@dataclass
class CachedKey:
    """key 的缓存校验信息（不含敏感的完整 key）。"""
    id: int
    name: str
    prefix: str
    key_hash: str
    is_active: bool
    expires_at: float | None  # epoch 秒
    allowed_models: list[str] | None
    rpm_limit: int | None
    tpm_limit: int | None
    token_quota: int | None
    tokens_used: int
    # 负载快照时间，用于判断是否需要回库刷新 tokens_used
    loaded_at: float


class KeyCache:
    """按 prefix 索引的 key 内存缓存。"""

    def __init__(self) -> None:
        self._by_prefix: Dict[str, CachedKey] = {}
        self._by_id: Dict[int, str] = {}  # id -> prefix，用于按 id 失效
        self._lock = asyncio.Lock()

    def get(self, prefix: str) -> CachedKey | None:
        return self._by_prefix.get(prefix)

    async def put(self, k: CachedKey) -> None:
        async with self._lock:
            self._by_prefix[k.prefix] = k
            self._by_id[k.id] = k.prefix

    async def invalidate_id(self, key_id: int) -> None:
        """按 id 失效（后台改/删 key 时调用）。"""
        async with self._lock:
            prefix = self._by_id.pop(key_id, None)
            if prefix:
                self._by_prefix.pop(prefix, None)

    async def invalidate_prefix(self, prefix: str) -> None:
        async with self._lock:
            ck = self._by_prefix.pop(prefix, None)
            if ck:
                self._by_id.pop(ck.id, None)

    async def clear(self) -> None:
        async with self._lock:
            self._by_prefix.clear()
            self._by_id.clear()

    async def add_usage(self, key_id: int, tokens: int) -> None:
        """成功转发后把用量累加到内存缓存（不回库），保持配额判断近实时。"""
        async with self._lock:
            prefix = self._by_id.get(key_id)
            if prefix:
                ck = self._by_prefix.get(prefix)
                if ck:
                    ck.tokens_used += tokens


class RPMCounter:
    """内存滑动窗口 RPM 计数，按 key_id 维度。"""

    def __init__(self, window: float = 60.0) -> None:
        self._window = window
        self._hits: Dict[int, List[float]] = {}
        self._lock = asyncio.Lock()

    async def hit_and_check(self, key_id: int, limit: int | None) -> bool:
        """记录一次命中并返回是否在限额内（True=放行，False=超限）。"""
        if not limit:
            return True
        now = time.time()
        cutoff = now - self._window
        async with self._lock:
            hits = self._hits.get(key_id)
            if hits is None:
                hits = []
                self._hits[key_id] = hits
            # 清掉过期命中
            hits[:] = [t for t in hits if t > cutoff]
            if len(hits) >= limit:
                return False
            hits.append(now)
            return True

    async def gc(self) -> None:
        """清理已无命中的 key 条目，避免长跑时 _hits 缓慢膨胀。"""
        cutoff = time.time() - self._window
        async with self._lock:
            for kid in [kid for kid, hits in self._hits.items() if not hits or hits[-1] <= cutoff]:
                self._hits.pop(kid, None)


class TPMCounter:
    """内存滑动窗口 TPM 计数，按 key_id 维度记录窗口内 token 消耗。

    请求前 check：窗口内累计 tokens >= limit 即拒绝；
    请求完成后 record 实际用量。在途请求的 token 不可预知，
    因此 TPM 是软限流（防持续超速，不防瞬时超额）。
    """

    def __init__(self, window: float = 60.0) -> None:
        self._window = window
        self._recs: Dict[int, List[tuple[float, int]]] = {}
        self._lock = asyncio.Lock()

    async def check(self, key_id: int, limit: int | None) -> bool:
        """检查窗口内已消耗 tokens 是否已达上限（True=放行）。顺带清理过期记录。"""
        if not limit:
            return True
        cutoff = time.time() - self._window
        async with self._lock:
            recs = self._recs.get(key_id)
            if not recs:
                self._recs.pop(key_id, None)  # 顺带清理空条目
                return True
            recs[:] = [(t, n) for t, n in recs if t > cutoff]
            if not recs:
                self._recs.pop(key_id, None)
                return True
            return sum(n for _, n in recs) < limit

    async def record(self, key_id: int, tokens: int) -> None:
        """请求完成后记录实际 token 消耗。"""
        if tokens <= 0:
            return
        async with self._lock:
            self._recs.setdefault(key_id, []).append((time.time(), tokens))


class AttemptLimiter:
    """通用失败尝试限流（如后台登录防爆破）：window 秒内失败超过 max_attempts 次即拒绝。"""

    def __init__(self, window: float = 600.0, max_attempts: int = 10) -> None:
        self._window = window
        self._max = max_attempts
        self._fails: Dict[str, List[float]] = {}
        self._lock = asyncio.Lock()

    async def allowed(self, ident: str) -> bool:
        """只检查不记录：是否已被限流（True=未被限流）。"""
        cutoff = time.time() - self._window
        async with self._lock:
            fails = self._fails.get(ident)
            if not fails:
                self._fails.pop(ident, None)
                return True
            fails[:] = [t for t in fails if t > cutoff]
            return len(fails) < self._max

    async def hit(self, ident: str) -> bool:
        """记录一次失败并返回是否仍在限额内（True=允许继续尝试，False=已锁定）。"""
        now = time.time()
        cutoff = now - self._window
        async with self._lock:
            fails = self._fails.setdefault(ident, [])
            fails[:] = [t for t in fails if t > cutoff]
            if len(fails) >= self._max:
                return False
            fails.append(now)
            return True

    async def reset(self, ident: str) -> None:
        """登录成功后清除失败计数。"""
        async with self._lock:
            self._fails.pop(ident, None)


key_cache = KeyCache()
rpm_counter = RPMCounter()
tpm_counter = TPMCounter()
login_limiter = AttemptLimiter()


def verify_key_hash(raw: str, digest: str) -> bool:
    """常数时间比较哈希，避免前缀时序侧信道。"""
    return hmac.compare_digest(hashlib.sha256(raw.encode()).hexdigest(), digest)