diff --git a/backend/app/services/mp_webview_ticket.py b/backend/app/services/mp_webview_ticket.py index 320256acf..124ff65d9 100644 --- a/backend/app/services/mp_webview_ticket.py +++ b/backend/app/services/mp_webview_ticket.py @@ -2,7 +2,9 @@ from __future__ import annotations import json -from typing import Optional +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 @@ -10,41 +12,70 @@ 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 issue_ticket(user_id: str, session_id: str, refresh_jti: str = "") -> str: +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("Redis 不可用,无法签发 WebView ticket") - 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 + 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 - r = get_sync_redis_client() - if r is None: - raise RuntimeError("Redis 不可用,无法兑换 WebView ticket") - key = _key(raw) - payload = r.get(key) - r.delete(key) - 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 + + 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") diff --git a/backend/tests/test_mp_webview_ticket.py b/backend/tests/test_mp_webview_ticket.py index e419b0df1..f567782e5 100644 --- a/backend/tests/test_mp_webview_ticket.py +++ b/backend/tests/test_mp_webview_ticket.py @@ -1,4 +1,3 @@ -import json from unittest.mock import patch import pytest @@ -22,6 +21,9 @@ class FakeRedis: def delete(self, key): return 1 if self.data.pop(key, None) is not None else 0 + def getdel(self, key): + return self.data.pop(key, None) + def test_issue_and_consume_once(): fake = FakeRedis() @@ -44,3 +46,9 @@ def test_issue_without_redis_raises(): with patch.object(svc, "get_sync_redis_client", return_value=None): with pytest.raises(RuntimeError, match="Redis"): svc.issue_ticket("u1", "sid-1") + + +def test_consume_without_redis_raises(): + with patch.object(svc, "get_sync_redis_client", return_value=None): + with pytest.raises(RuntimeError, match="Redis"): + svc.consume_ticket("any-ticket")