shanxia2.0/backend/app/services/mp_webview_ticket.py

"""小程序 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")