From ed14031d65eed9b82e04a28081444f58f8ccff10 Mon Sep 17 00:00:00 2001
From: zzlgreat <zzlgreat@gmail.com>
Date: Tue, 2 Dec 2025 11:07:45 +0800
Subject: [PATCH] update pay ui

---
 app_vx.py | 109 ++++++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 97 insertions(+), 12 deletions(-)

diff --git a/app_vx.py b/app_vx.py
index d3188681..1e51f6c6 100644
--- a/app_vx.py
+++ b/app_vx.py
@@ -559,7 +559,86 @@ app.config['COMPRESS_MIMETYPES'] = [
     'application/javascript',
     'application/x-javascript'
 ]
-user_tokens = {}
+
+# ===================== Token 存储（支持多 worker 共享） =====================
+class TokenStore:
+    """
+    Token 存储类 - 支持 Redis（多 worker 共享）或内存（单 worker）
+    """
+    def __init__(self):
+        self._redis_client = None
+        self._memory_store = {}
+        self._prefix = 'vf_token:'
+        self._initialized = False
+
+    def _ensure_initialized(self):
+        """延迟初始化，确保在 fork 后才连接 Redis"""
+        if self._initialized:
+            return
+        self._initialized = True
+
+        redis_url = os.environ.get('REDIS_URL', 'redis://localhost:6379/0')
+        try:
+            import redis
+            self._redis_client = redis.from_url(redis_url)
+            self._redis_client.ping()
+            logger.info(f"✅ Token 存储: Redis ({redis_url})")
+        except Exception as e:
+            logger.warning(f"⚠️ Redis 不可用 ({e})，Token 使用内存存储（多 worker 模式下会有问题！）")
+            self._redis_client = None
+
+    def get(self, token):
+        """获取 token 数据"""
+        self._ensure_initialized()
+        if self._redis_client:
+            try:
+                data = self._redis_client.get(f"{self._prefix}{token}")
+                if data:
+                    return json.loads(data)
+                return None
+            except Exception as e:
+                logger.error(f"Redis get error: {e}")
+                return self._memory_store.get(token)
+        return self._memory_store.get(token)
+
+    def set(self, token, data, expire_seconds=30*24*3600):
+        """设置 token 数据"""
+        self._ensure_initialized()
+        if self._redis_client:
+            try:
+                # 将 datetime 转为字符串存储
+                store_data = data.copy()
+                if 'expires' in store_data and isinstance(store_data['expires'], datetime):
+                    store_data['expires'] = store_data['expires'].isoformat()
+                self._redis_client.setex(
+                    f"{self._prefix}{token}",
+                    expire_seconds,
+                    json.dumps(store_data)
+                )
+                return
+            except Exception as e:
+                logger.error(f"Redis set error: {e}")
+        self._memory_store[token] = data
+
+    def delete(self, token):
+        """删除 token"""
+        self._ensure_initialized()
+        if self._redis_client:
+            try:
+                self._redis_client.delete(f"{self._prefix}{token}")
+                return
+            except Exception as e:
+                logger.error(f"Redis delete error: {e}")
+        self._memory_store.pop(token, None)
+
+    def __contains__(self, token):
+        """支持 'in' 操作符"""
+        return self.get(token) is not None
+
+
+# 使用 TokenStore 替代内存字典
+user_tokens = TokenStore()
+
 app.config['SECRET_KEY'] = 'vf7891574233241'
 app.config['SQLALCHEMY_DATABASE_URI'] = 'mysql+pymysql://root:Zzl5588161!@222.128.1.157:33060/stock?charset=utf8mb4'
 app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
@@ -683,9 +762,12 @@ def token_required(f):
         if not token_data:
             return jsonify({'message': 'Token无效', 'code': 401}), 401
 
-        # 检查是否过期
-        if token_data['expires'] < datetime.now():
-            del user_tokens[token]
+        # 检查是否过期（expires 可能是字符串或 datetime）
+        expires = token_data['expires']
+        if isinstance(expires, str):
+            expires = datetime.fromisoformat(expires)
+        if expires < datetime.now():
+            user_tokens.delete(token)
             return jsonify({'message': 'Token已过期'}), 401
 
         # 获取用户对象并添加到请求上下文
@@ -3460,7 +3542,7 @@ def logout_with_token():
         token = data.get('token') if data else None
 
     if token and token in user_tokens:
-        del user_tokens[token]
+        user_tokens.delete(token)
 
     # 清除session
     session.clear()
@@ -3617,10 +3699,10 @@ def login_with_phone():
         token = generate_token(32)
 
         # 存储token映射（30天有效期）
-        user_tokens[token] = {
+        user_tokens.set(token, {
             'user_id': user.id,
             'expires': datetime.now() + timedelta(days=30)
-        }
+        })
 
         # 清除验证码
         del verification_codes[phone]
@@ -3670,9 +3752,12 @@ def verify_token():
     if not token_data:
         return jsonify({'valid': False, 'message': 'Token无效', 'code': 401}), 401
 
-    # 检查是否过期
-    if token_data['expires'] < datetime.now():
-        del user_tokens[token]
+    # 检查是否过期（expires 可能是字符串或 datetime）
+    expires = token_data['expires']
+    if isinstance(expires, str):
+        expires = datetime.fromisoformat(expires)
+    if expires < datetime.now():
+        user_tokens.delete(token)
         return jsonify({'valid': False, 'message': 'Token已过期'}), 401
 
     # 获取用户信息
@@ -3905,10 +3990,10 @@ def api_login_wechat():
         token = generate_token(32)  # 使用相同的随机字符串生成器
 
         # 存储token映射（与手机登录保持一致）
-        user_tokens[token] = {
+        user_tokens.set(token, {
             'user_id': user.id,
             'expires': datetime.now() + timedelta(days=30)  # 30天有效期
-        }
+        })
 
         # 设置session（可选，保持与手机登录一致）
         session.permanent = True