155 lines
4.6 KiB
Python
155 lines
4.6 KiB
Python
"""小程序 WebView 一次性短时 ticket（不把 JWT 放进 URL）。
|
|
|
|
默认签发短时签名票（S. 前缀）：不依赖 Redis 跨实例共享，避免生产多机
|
|
「签发在 A、兑换在 B」导致 ticket 无效。
|
|
仍兼容旧版 UUID+Redis 票；签名票在 Redis 可用时用 jti 做记录（允许 TTL 内复兑）。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Callable, Optional, TypeVar
|
|
|
|
from jose import JWTError, jwt
|
|
from redis import exceptions as redis_exceptions
|
|
|
|
from app.core.config import settings
|
|
from app.core.redis_client import get_sync_redis_client
|
|
from app.core.utils import generate_uuid
|
|
|
|
logger = logging.getLogger("sqxp")
|
|
|
|
TICKET_TTL_SEC = 60
|
|
KEY_PREFIX = "mp_wv_ticket:"
|
|
JTI_USED_PREFIX = "mp_wv_jti:"
|
|
SIGNED_PREFIX = "S."
|
|
SIGNED_TYP = "mp_wv_ticket"
|
|
|
|
_GETDEL_LUA = """
|
|
local val = redis.call('GET', KEYS[1])
|
|
if val then
|
|
redis.call('DEL', KEYS[1])
|
|
end
|
|
return val
|
|
"""
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def _key(ticket: str) -> str:
|
|
return f"{KEY_PREFIX}{ticket}"
|
|
|
|
|
|
def _jti_key(jti: str) -> str:
|
|
return f"{JTI_USED_PREFIX}{jti}"
|
|
|
|
|
|
def _atomic_getdel(r, key: str):
|
|
"""原子 GET+DEL。
|
|
|
|
不用 redis-py 的 ``getdel``：客户端有该方法，但 Redis < 6.2 不支持 GETDEL 命令，
|
|
生产常见 5.x 会直接报错并被包装成 503。Lua 在旧版 Redis 上可用。
|
|
"""
|
|
return r.eval(_GETDEL_LUA, 1, key)
|
|
|
|
|
|
def _redis_call(fn: Callable[..., T], unavailable_msg: str) -> T:
|
|
r = get_sync_redis_client()
|
|
if r is None:
|
|
raise RuntimeError(unavailable_msg)
|
|
try:
|
|
return fn(r)
|
|
except (
|
|
redis_exceptions.TimeoutError,
|
|
redis_exceptions.ConnectionError,
|
|
redis_exceptions.RedisError,
|
|
) as exc:
|
|
raise RuntimeError(unavailable_msg) from exc
|
|
|
|
|
|
def _issue_signed_ticket(user_id: str, session_id: str, refresh_jti: str = "") -> str:
|
|
now = datetime.now(timezone.utc)
|
|
payload = {
|
|
"typ": SIGNED_TYP,
|
|
"userId": user_id,
|
|
"sessionId": session_id,
|
|
"refreshJti": refresh_jti or "",
|
|
"jti": generate_uuid(),
|
|
"iat": now,
|
|
"exp": now + timedelta(seconds=TICKET_TTL_SEC),
|
|
}
|
|
token = jwt.encode(payload, settings.SECRET_KEY, algorithm=settings.ALGORITHM)
|
|
if isinstance(token, bytes):
|
|
token = token.decode("utf-8")
|
|
return f"{SIGNED_PREFIX}{token}"
|
|
|
|
|
|
def _try_mark_jti_once(jti: str) -> Optional[bool]:
|
|
"""Redis 可用时标记 jti 已用。返回 True=首次, False=重复, None=Redis 不可用。"""
|
|
if not jti:
|
|
return None
|
|
|
|
def _mark(r) -> bool:
|
|
return bool(r.set(_jti_key(jti), "1", ex=TICKET_TTL_SEC, nx=True))
|
|
|
|
try:
|
|
return _redis_call(_mark, "Redis 不可用")
|
|
except RuntimeError:
|
|
return None
|
|
|
|
|
|
def _consume_signed_ticket(raw_jwt: str) -> Optional[dict]:
|
|
try:
|
|
data = jwt.decode(raw_jwt, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
|
|
except JWTError:
|
|
return None
|
|
if str(data.get("typ") or "") != SIGNED_TYP:
|
|
return None
|
|
user_id = str(data.get("userId") or "").strip()
|
|
session_id = str(data.get("sessionId") or "").strip()
|
|
if not user_id or not session_id:
|
|
return None
|
|
jti = str(data.get("jti") or "").strip()
|
|
marked = _try_mark_jti_once(jti)
|
|
if marked is False:
|
|
# 已兑换过；微信 web-view 可能二次加载，同 TTL 内允许再兑
|
|
logger.info("mp-webview signed ticket jti already used, allow within TTL: %s", jti[:8])
|
|
return {
|
|
"userId": user_id,
|
|
"sessionId": session_id,
|
|
"refreshJti": str(data.get("refreshJti") or ""),
|
|
}
|
|
|
|
|
|
def issue_ticket(user_id: str, session_id: str, refresh_jti: str = "") -> str:
|
|
"""签发签名票（跨实例不依赖 Redis 读到对方写入的 key）。"""
|
|
return _issue_signed_ticket(user_id, session_id, refresh_jti)
|
|
|
|
|
|
def consume_ticket(ticket: str) -> Optional[dict]:
|
|
raw = (ticket or "").strip()
|
|
if not raw:
|
|
return None
|
|
|
|
if raw.startswith(SIGNED_PREFIX):
|
|
return _consume_signed_ticket(raw[len(SIGNED_PREFIX) :])
|
|
|
|
# 兼容部署过渡期仍在流通的 UUID+Redis 票
|
|
def _consume(r) -> Optional[dict]:
|
|
payload = _atomic_getdel(r, _key(raw))
|
|
if not payload:
|
|
return None
|
|
if isinstance(payload, bytes):
|
|
payload = payload.decode("utf-8")
|
|
data = json.loads(payload)
|
|
if not isinstance(data, dict):
|
|
return None
|
|
return data
|
|
|
|
try:
|
|
return _redis_call(_consume, "Redis 不可用，无法兑换 WebView ticket")
|
|
except RuntimeError as exc:
|
|
logger.warning("mp-webview ticket Redis 兑换失败: %s", exc)
|
|
raise
|