chongming-core/test/test_jwt_cache.py

"""
JWT 解析共享缓存测试（auth.py _cached_decode / jwt_auth_dependency）

验证:
  1. 同 token 重复鉴权只解密一次（缓存命中），不同 token 独立缓存
  2. TTL 过期后重新解码（条目自然失效）
  3. 缓存条目 exp 已过期：命中即拒绝，不重解密（缓存不延长 token 有效期）
  4. 无效 token 不缓存失败：每次都走完整验证（decode 调用计数递增）
  5. 缓存 key 含密钥/算法指纹：同 token 不同密钥不串用
  6. 返回值为独立对象：调用方修改不影响缓存
  7. 端到端：真实 chongming_jwt 签发 -> 鉴权成功 -> 二次命中缓存

运行: uv run python test_jwt_cache.py
"""
import sys
import tempfile
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent / "src"))
sys.path.insert(0, str(Path(__file__).resolve().parent / "utils" / "worker" / "src"))

import orjson
from fastapi.security import HTTPAuthorizationCredentials

from chongming.app.core import auth as auth_mod
from chongming.app.core.auth import jwt_auth_dependency
from chongming_cache import Cache
from chongming_jwt import JWT
from chongming_config import get_config

failures = []


def check(name, cond, detail=""):
    status = "PASS" if cond else "FAIL"
    print(f"[{status}] {name}" + (f" — {detail}" if detail else ""))
    if not cond:
        failures.append(name)


config = get_config()
SECRET = config["security"]["secret_key"]
ALGORITHM = config["security"]["algorithm"]

# 每个测试使用独立临时 diskcache，隔离真实缓存目录
_tmp = tempfile.TemporaryDirectory()


def fresh_cache():
    auth_mod._jwt_cache = Cache(backend="diskcache", cache_store_path=_tmp.name)
    auth_mod._jwt_cache.clear()


def make_creds(token):
    return HTTPAuthorizationCredentials(scheme="Bearer", credentials=token)


# ---- 1. 缓存命中与隔离 ----
def test_cache_hit_and_isolation():
    fresh_cache()

    class CountingJWT:
        calls = 0

        def __init__(self, secret, algorithm):
            pass

        def decode(self, token):
            CountingJWT.calls += 1
            return orjson.dumps({"sub": token, "exp": int(time.time()) + 3600}).decode()

    auth_mod.JWT = CountingJWT
    t1, t2 = "token-1", "token-2"
    p1 = jwt_auth_dependency(credentials=make_creds(t1))
    p2 = jwt_auth_dependency(credentials=make_creds(t1))
    p3 = jwt_auth_dependency(credentials=make_creds(t2))
    check("同 token 二次鉴权命中缓存（只解密一次）", CountingJWT.calls == 2, f"calls={CountingJWT.calls}")
    check("payload 解析正确", p1["sub"] == "token-1" and p2["sub"] == "token-1", str(p1))
    check("不同 token 独立缓存", p3["sub"] == "token-2", str(p3))
    check("缓存中已写入条目", len(auth_mod._jwt_cache) >= 2, str(len(auth_mod._jwt_cache)))


# ---- 2. TTL 过期后重新解码 ----
def test_ttl_expiry_redecodes():
    fresh_cache()
    auth_mod._JWT_CACHE_TTL = 1  # 缩短 TTL 加速测试

    class TtlJWT:
        calls = 0

        def __init__(self, secret, algorithm):
            pass

        def decode(self, token):
            TtlJWT.calls += 1
            return orjson.dumps({"sub": "u1", "exp": int(time.time()) + 3600}).decode()

    auth_mod.JWT = TtlJWT
    creds = make_creds("ttl-token")
    p1 = jwt_auth_dependency(credentials=creds)
    p2 = jwt_auth_dependency(credentials=creds)
    check("TTL 内命中", TtlJWT.calls == 1 and p2["sub"] == "u1", f"calls={TtlJWT.calls}")
    time.sleep(1.2)  # 等待 TTL 过期
    p3 = jwt_auth_dependency(credentials=creds)
    check("TTL 过期后重新解码", TtlJWT.calls == 2 and p3["sub"] == "u1", f"calls={TtlJWT.calls}")


# ---- 3. 缓存条目 exp 已过期：命中即拒绝，不重解密 ----
def test_cached_expired_rejected():
    fresh_cache()

    class NeverJWT:
        calls = 0

        def __init__(self, secret, algorithm):
            pass

        def decode(self, token):
            NeverJWT.calls += 1
            raise AssertionError("不应触发解码（命中缓存即应拒绝）")

    auth_mod.JWT = NeverJWT
    # 预置一个 exp 已过期的缓存条目（模拟缓存写入后 token 到期）
    key = auth_mod._cache_key("expired-token", SECRET, ALGORITHM)
    auth_mod._jwt_cache.set(
        key, {"payload": {"sub": "u1"}, "exp": int(time.time()) - 10}, expire=60
    )
    for i in range(2):
        try:
            jwt_auth_dependency(credentials=make_creds("expired-token"))
            check(f"缓存过期条目第 {i+1} 次应 401", False, "未抛异常")
        except Exception as e:
            check(f"缓存过期条目第 {i+1} 次 401", getattr(e, "status_code", None) == 401, str(e))
    check("过期条目命中后未重解密（两次都直接拒绝）", NeverJWT.calls == 0, f"calls={NeverJWT.calls}")
    check("过期条目保留在缓存（TTL 内持续拒绝）", key in auth_mod._jwt_cache)


# ---- 4. 无效 token 失败不缓存 ----
def test_invalid_token_not_cached():
    fresh_cache()

    class BadJWT:
        calls = 0

        def __init__(self, secret, algorithm):
            pass

        def decode(self, token):
            BadJWT.calls += 1
            raise ValueError("Invalid signature")

    auth_mod.JWT = BadJWT
    creds = make_creds("bad-token")
    for _ in range(2):
        try:
            jwt_auth_dependency(credentials=creds)
            check("无效 token 应 401", False)
        except Exception as e:
            ok = getattr(e, "status_code", None) == 401
            check("无效 token 401", ok, str(e))
    check("无效 token 每次都完整验证（失败不缓存）", BadJWT.calls == 2, f"calls={BadJWT.calls}")
    check("缓存无失败条目", len(auth_mod._jwt_cache) == 0, str(len(auth_mod._jwt_cache)))


# ---- 5. 缓存 key 含密钥/算法指纹 ----
def test_key_scoped_to_secret():
    fresh_cache()

    class KeyJWT:
        calls = 0

        def __init__(self, secret, algorithm):
            pass

        def decode(self, token):
            KeyJWT.calls += 1
            return orjson.dumps({"sub": token, "exp": int(time.time()) + 3600}).decode()

    auth_mod.JWT = KeyJWT
    creds = make_creds("same-token")
    jwt_auth_dependency(credentials=creds)                       # secret A
    jwt_auth_dependency(credentials=creds)                       # secret A -> 命中
    auth_mod._cached_decode("same-token", "secret-B", ALGORITHM)  # secret B -> miss
    check("同 token 不同密钥独立缓存", KeyJWT.calls == 2, f"calls={KeyJWT.calls}")


# ---- 6. 返回值独立对象：修改不影响缓存 ----
def test_payload_object_independent():
    fresh_cache()

    class FixedJWT:
        def __init__(self, secret, algorithm):
            pass

        def decode(self, token):
            return orjson.dumps({"sub": "u1", "role": "admin", "exp": int(time.time()) + 3600}).decode()

    auth_mod.JWT = FixedJWT
    p1 = jwt_auth_dependency(credentials=make_creds("t-copy"))
    p1["role"] = "hacker"
    p2 = jwt_auth_dependency(credentials=make_creds("t-copy"))
    check("修改返回值不污染缓存", p2["role"] == "admin", str(p2))


# ---- 7. 端到端：真实 JWT 签发 -> 鉴权 -> 命中缓存 ----
def test_e2e_real_jwt():
    fresh_cache()
    auth_mod.JWT = JWT  # 恢复真实 JWT 类
    jwt = JWT(SECRET, ALGORITHM)
    token = jwt.encode(orjson.dumps({"sub": "1", "username": "alice"}).decode(), expires_in_seconds=3600)
    creds = make_creds(token)
    p1 = jwt_auth_dependency(credentials=creds)
    p2 = jwt_auth_dependency(credentials=creds)
    check("真实 JWT 鉴权成功", p1["sub"] == "1" and p1["username"] == "alice", str(p1))
    check("二次调用命中缓存", auth_mod._cache_key(token, SECRET, ALGORITHM) in auth_mod._jwt_cache)
    try:
        jwt_auth_dependency(credentials=None)
        check("缺 header 应 401", False)
    except Exception as e:
        check("缺 header 401", getattr(e, "status_code", None) == 401)


test_cache_hit_and_isolation()
test_ttl_expiry_redecodes()
test_cached_expired_rejected()
test_invalid_token_not_cached()
test_key_scoped_to_secret()
test_payload_object_independent()
test_e2e_real_jwt()

_tmp.cleanup()

print()
if failures:
    print(f"共 {len(failures)} 项失败: {failures}")
    sys.exit(1)
print("全部通过")