fix: avoid websocket token leakage in logs

backend-fastapi/core/websocket/consumers/base.py

@@ -17,6 +17,8 @@ from utils.security import verify_access_token
 
 logger = logging.getLogger(__name__)
 
+WS_AUTH_PROTOCOL = "access_token"
+
 
 class ConnectionManager:
     """WebSocket 连接管理器"""
@@ -27,9 +29,14 @@ class ConnectionManager:
         # 组连接: {group_name: {websocket1, websocket2, ...}}
         self.groups: Dict[str, Set[WebSocket]] = {}
     
-    async def connect(self, websocket: WebSocket, user_id: str):
+    async def connect(
+        self,
+        websocket: WebSocket,
+        user_id: str,
+        subprotocol: Optional[str] = None,
+    ):
         """添加连接"""
-        await websocket.accept()
+        await websocket.accept(subprotocol=subprotocol)
         if user_id not in self.active_connections:
             self.active_connections[user_id] = set()
         self.active_connections[user_id].add(websocket)
@@ -101,27 +108,45 @@ class TokenAuthWebSocketConsumer:
         self.user_id: Optional[str] = None
         self.is_authenticated = False
         self._token: Optional[str] = None  # 保存原始token用于心跳校验
+        self._accept_subprotocol: Optional[str] = None
+
+    def _get_token_from_protocol(self) -> Optional[str]:
+        protocol_header = self.websocket.headers.get("sec-websocket-protocol") or ""
+        protocols = [
+            item.strip()
+            for item in protocol_header.split(",")
+            if item.strip()
+        ]
+        for index, protocol in enumerate(protocols):
+            if protocol == WS_AUTH_PROTOCOL and index + 1 < len(protocols):
+                self._accept_subprotocol = WS_AUTH_PROTOCOL
+                return protocols[index + 1]
+        return None
+
+    def _get_token_from_query(self) -> Optional[str]:
+        query_string = self.websocket.scope.get('query_string', b'').decode('utf-8')
+        if not query_string:
+            return None
+
+        query_params = parse_qs(query_string)
+        token_list = query_params.get('token', [])
+        return token_list[0] if token_list else None
+
+    async def _accept_then_close(self, code: int):
+        await self.websocket.accept(subprotocol=self._accept_subprotocol)
+        await self.websocket.close(code=code)
     
     async def authenticate(self) -> bool:
         """
         进行Token认证
-        从查询参数中获取token并验证
+        优先从 WebSocket 子协议中获取 token，兼容旧 query 参数方式。
         """
-        # 获取查询参数中的token
-        query_string = self.websocket.scope.get('query_string', b'').decode('utf-8')
-        token = None
-        
-        if query_string:
-            query_params = parse_qs(query_string)
-            token_list = query_params.get('token', [])
-            if token_list:
-                token = token_list[0]
+        token = self._get_token_from_protocol() or self._get_token_from_query()
         
         if not token:
             logger.warning("WebSocket connection rejected: No token provided")
             # 必须先accept才能close
-            await self.websocket.accept()
-            await self.websocket.close(code=4001)
+            await self._accept_then_close(code=4001)
             return False
         
         # 验证token
@@ -130,15 +155,13 @@ class TokenAuthWebSocketConsumer:
             
             if not payload:
                 logger.warning("WebSocket connection rejected: Invalid token")
-                await self.websocket.accept()
-                await self.websocket.close(code=4001)
+                await self._accept_then_close(code=4001)
                 return False
             
             user_id = payload.get('sub')
             if not user_id:
                 logger.warning("WebSocket connection rejected: Invalid token payload")
-                await self.websocket.accept()
-                await self.websocket.close(code=4001)
+                await self._accept_then_close(code=4001)
                 return False
             
             self.user_id = user_id
@@ -149,14 +172,17 @@ class TokenAuthWebSocketConsumer:
             
         except Exception as e:
             logger.error(f"WebSocket authentication failed: {str(e)}")
-            await self.websocket.accept()
-            await self.websocket.close(code=4001)
+            await self._accept_then_close(code=4001)
             return False
     
     async def connect(self):
         """连接时进行Token认证"""
         if await self.authenticate():
-            await manager.connect(self.websocket, self.user_id)
+            await manager.connect(
+                self.websocket,
+                self.user_id,
+                subprotocol=self._accept_subprotocol,
+            )
     
     async def disconnect(self, close_code: int = 1000):
         """断开连接"""