feat(auth): add wx-qr session state machine for web login.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
wuzilongandCursor committed 2026-09-18 19:55:00 +08:00
1 parent 7f2b5d06c1
commit 56a3311c45
2 files changed
+225

No files matched your search

+172
View File
@@ -0,0 +1,172 @@
"""web-admin 扫码登录会话:短码映射 + 状态机;优先 Redis,否则进程内存。"""
from __future__ import annotations
import json
import secrets
import time
import uuid
from typing import Any, Dict, Optional
from app.core.config import settings
_PREFIX = "wxqr:sess:"
_SHORT_PREFIX = "wxqr:short:"
_memory: Dict[str, Dict[str, Any]] = {}
_short_map: Dict[str, str] = {}
def _memory_clear_for_tests() -> None:
_memory.clear()
_short_map.clear()
def _ttl() -> int:
return int(getattr(settings, "WX_QR_SESSION_TTL_SEC", 180) or 180)
def _redis():
"""复用 auth 侧 Redis 客户端模式;无 REDIS_HOST 返回 None。"""
host = (settings.REDIS_HOST or "").strip()
if not host:
return None
try:
import redis
return redis.Redis(
host=host,
port=int(settings.REDIS_PORT or 6379),
db=int(settings.REDIS_DB or 0),
password=settings.REDIS_PASSWORD or None,
decode_responses=True,
socket_timeout=2,
)
except Exception:
return None
def create_session(*, ttl_sec: Optional[int] = None) -> tuple:
ttl = int(ttl_sec if ttl_sec is not None else _ttl())
session_id = str(uuid.uuid4())
short = secrets.token_urlsafe(16)[:32]
now = int(time.time())
data = {
"sessionId": session_id,
"shortCode": short,
"status": "pending",
"message": "",
"createdAt": now,
"expiresAt": now + ttl,
"tokens": None,
}
r = _redis()
if r:
pipe = r.pipeline()
pipe.setex(_PREFIX + session_id, ttl, json.dumps(data, ensure_ascii=False))
pipe.setex(_SHORT_PREFIX + short, ttl, session_id)
pipe.execute()
else:
_memory[session_id] = data
_short_map[short] = session_id
return session_id, short
def resolve_session_id(short_or_id: str) -> Optional[str]:
key = (short_or_id or "").strip()
if not key:
return None
r = _redis()
if r:
sid = r.get(_SHORT_PREFIX + key)
if sid:
return sid
if r.exists(_PREFIX + key):
return key
return None
if key in _short_map:
return _short_map[key]
if key in _memory:
return key
return None
def get_session(session_id: str) -> Optional[Dict[str, Any]]:
sid = (session_id or "").strip()
if not sid:
return None
r = _redis()
if r:
raw = r.get(_PREFIX + sid)
if not raw:
return None
data = json.loads(raw)
else:
data = _memory.get(sid)
if not data:
return None
data = dict(data)
if int(time.time()) >= int(data.get("expiresAt") or 0):
data["status"] = "expired"
return data
return data
def _save(data: Dict[str, Any]) -> None:
sid = data["sessionId"]
ttl = max(1, int(data["expiresAt"]) - int(time.time()))
r = _redis()
if r:
r.setex(_PREFIX + sid, ttl, json.dumps(data, ensure_ascii=False))
else:
_memory[sid] = data
def mark_scanned(session_id: str) -> bool:
data = get_session(session_id)
if not data or data["status"] in ("expired", "success", "consumed", "failed"):
return False
if data["status"] == "pending":
data["status"] = "scanned"
_save(data)
return True
def mark_success(session_id: str, tokens: Dict[str, Any]) -> bool:
data = get_session(session_id)
if not data or data["status"] in ("expired", "consumed", "failed"):
return False
data["status"] = "success"
data["tokens"] = tokens
data["message"] = ""
_save(data)
return True
def mark_failed(session_id: str, message: str) -> bool:
data = get_session(session_id)
if not data or data["status"] in ("expired", "consumed", "success"):
return False
data["status"] = "failed"
data["message"] = message or "登录失败"
data["tokens"] = None
_save(data)
return True
def consume_success(session_id: str) -> Optional[Dict[str, Any]]:
data = get_session(session_id)
if not data or data["status"] != "success":
return None
tokens = data.get("tokens")
data["status"] = "consumed"
data["tokens"] = None
_save(data)
return tokens
def invalidate(session_id: str) -> None:
data = get_session(session_id)
if not data:
return
data["status"] = "expired"
data["tokens"] = None
_save(data)
+53
View File
@@ -0,0 +1,53 @@
"""wx_qr_session 状态机与一次性消费。"""
import time
from app.services import wx_qr_session as s
def setup_function():
s._memory_clear_for_tests()
def test_create_pending_and_get():
sid, short = s.create_session(ttl_sec=60)
assert sid and short and len(short) <= 32
data = s.get_session(sid)
assert data["status"] == "pending"
assert s.resolve_session_id(short) == sid
def test_mark_scanned():
sid, _ = s.create_session(ttl_sec=60)
assert s.mark_scanned(sid) is True
assert s.get_session(sid)["status"] == "scanned"
def test_mark_success_consume_once():
sid, _ = s.create_session(ttl_sec=60)
payload = {"token": "t", "refreshToken": "r", "expiresIn": 1, "userInfo": {}}
assert s.mark_success(sid, payload) is True
first = s.consume_success(sid)
assert first["token"] == "t"
assert s.consume_success(sid) is None
assert s.get_session(sid)["status"] == "consumed"
def test_mark_failed():
sid, _ = s.create_session(ttl_sec=60)
assert s.mark_failed(sid, "未找到员工") is True
assert s.get_session(sid)["status"] == "failed"
assert s.get_session(sid)["message"] == "未找到员工"
def test_expire_old_session():
sid, _ = s.create_session(ttl_sec=1)
time.sleep(1.1)
data = s.get_session(sid)
assert data is None or data["status"] == "expired"
def test_invalidate():
sid, _ = s.create_session(ttl_sec=60)
s.invalidate(sid)
data = s.get_session(sid)
assert data is None or data["status"] == "expired"