Procházet zdrojové kódy

fix(backend): consume webview tickets atomically and fail closed on Redis errors.

Co-authored-by: Cursor <cursoragent@cursor.com>
wuzilong před 1 měsícem
rodič
revize
98292cb914

+ 57 - 26
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")

+ 9 - 1
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")