fix websocket auth

api/controllers/console/__init__.py

@@ -59,6 +59,7 @@ from .app import (
     mcp_server,
     message,
     model_config,
+    online_user,
     ops_trace,
     site,
     statistic,

api/controllers/console/app/online_user.py

@@ -1,21 +1,21 @@
 import json
 
 from flask import request
-from flask_login import current_user, login_required
 
 from extensions.ext_redis import redis_client
 from extensions.ext_socketio import ext_socketio
 
 
+
 @ext_socketio.on("user_connect")
-@login_required
 def handle_user_connect(data):
     """
     Handle user connect event, check login and get user info.
     """
-
     sid = request.sid
     workflow_id = data.get("workflow_id")
+    if not (current_user := request.environ.get("ws_user")):
+        return {"msg": "unauthorized"}, 401
 
     old_info_json = redis_client.hget(f"workflow_online_users:{workflow_id}", current_user.id)
     if old_info_json:

api/extensions/ext_login.py

@@ -1,13 +1,15 @@
 import json
 
 import flask_login  # type: ignore
-from flask import Response, request
+from flask import Response, request, g
+from flask_socketio import disconnect
 from flask_login import user_loaded_from_request, user_logged_in
 from werkzeug.exceptions import NotFound, Unauthorized
 
 from configs import dify_config
 from dify_app import DifyApp
 from extensions.ext_database import db
+from extensions.ext_socketio import ext_socketio
 from libs.passport import PassportService
 from models.account import Account, Tenant, TenantAccountJoin
 from models.model import AppMCPServer, EndUser
@@ -113,5 +115,35 @@ def unauthorized_handler():
     )
 
 
+@ext_socketio.on('connect')
+def socket_connect(auth):
+    """
+    WebSocket connect event, do authentication here.
+    """
+    token = None
+    if auth and isinstance(auth, dict):
+        token = auth.get('token')
+    if not token:
+        disconnect()
+        return False
+
+    try:
+        decoded = PassportService().verify(token)
+        user_id = decoded.get("user_id")
+        if not user_id:
+            disconnect()
+            return False
+        user = AccountService.load_logged_in_account(account_id=user_id)
+        if not user:
+            disconnect()
+            return False
+
+        request.environ['ws_user'] = user
+
+    except Exception:
+        disconnect()
+        return False
+
+
 def init_app(app: DifyApp):
     login_manager.init_app(app)