feat(auth): add wx-qr session state machine for web login.
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
1 parent
7f2b5d06c1
commit
56a3311c45
2 files changed
+225
No files matched your search
@@ -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)
|
||||
@@ -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"
|
||||
Reference in new issue
Block a user