239 lines
8.1 KiB
Python
239 lines
8.1 KiB
Python
"""
|
||
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("全部通过")
|