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

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