82 lines
2.1 KiB
Python
82 lines
2.1 KiB
Python
"""小程序 WebView 一次性短时 ticket（不把 JWT 放进 URL）。"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Callable, Optional, TypeVar
|
|
|
|
from redis import exceptions as redis_exceptions
|
|
|
|
from app.core.redis_client import get_sync_redis_client
|
|
from app.core.utils import generate_uuid
|
|
|
|
TICKET_TTL_SEC = 60
|
|
KEY_PREFIX = "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 _atomic_getdel(r, key: str):
|
|
if hasattr(r, "getdel"):
|
|
return r.getdel(key)
|
|
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_ticket(user_id: str, session_id: str, refresh_jti: str = "") -> str:
|
|
def _issue(r) -> str:
|
|
ticket = generate_uuid()
|
|
body = json.dumps(
|
|
{"userId": user_id, "sessionId": session_id, "refreshJti": refresh_jti or ""},
|
|
ensure_ascii=False,
|
|
)
|
|
ok = r.set(_key(ticket), body, ex=TICKET_TTL_SEC, nx=True)
|
|
if not ok:
|
|
raise RuntimeError("签发 WebView ticket 失败")
|
|
return ticket
|
|
|
|
return _redis_call(_issue, "Redis 不可用，无法签发 WebView ticket")
|
|
|
|
|
|
def consume_ticket(ticket: str) -> Optional[dict]:
|
|
raw = (ticket or "").strip()
|
|
if not raw:
|
|
return None
|
|
|
|
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
|
|
|
|
return _redis_call(_consume, "Redis 不可用，无法兑换 WebView ticket")
|