From c5d9bac0edd507c4fe48556d59f1a35e73dc399c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?cls=5F=E5=AE=81=E6=B3=A2=E6=9C=AC=E6=9C=BA?= <908705107@qq.com> Date: Mon, 22 Jun 2026 11:34:35 +0800 Subject: [PATCH] fix: avoid websocket token leakage in logs --- .../core/websocket/consumers/base.py | 68 +++++++++++++------ web/apps/web-ele/src/api/core/websocket.ts | 19 +++--- 2 files changed, 57 insertions(+), 30 deletions(-) diff --git a/backend-fastapi/core/websocket/consumers/base.py b/backend-fastapi/core/websocket/consumers/base.py index 106c8f6..5ca7a66 100644 --- a/backend-fastapi/core/websocket/consumers/base.py +++ b/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): """断开连接""" diff --git a/web/apps/web-ele/src/api/core/websocket.ts b/web/apps/web-ele/src/api/core/websocket.ts index 4a52f16..8edb3d6 100644 --- a/web/apps/web-ele/src/api/core/websocket.ts +++ b/web/apps/web-ele/src/api/core/websocket.ts @@ -169,21 +169,22 @@ export class WebSocketManager { return; } - // 构建带token的URL - const separator = this.config.url.includes('?') ? '&' : '?'; - const wsUrl = `${this.config.url}${separator}token=${encodeURIComponent(token)}`; + const wsUrl = this.config.url; + const protocols = Array.isArray(this.config.protocols) + ? [...this.config.protocols] + : this.config.protocols + ? [this.config.protocols] + : []; + protocols.unshift('access_token', token); - console.log( - 'Connecting to WebSocket:', - wsUrl.replace(/token=[^&]+/, 'token=***'), - ); + console.log('Connecting to WebSocket:', wsUrl); console.log('WebSocket URL详情:', { originalUrl: this.config.url, - finalUrl: wsUrl.replace(/token=[^&]+/, 'token=***'), + finalUrl: wsUrl, isDev: import.meta.env.DEV, }); - this.ws = new WebSocket(wsUrl, this.config.protocols); + this.ws = new WebSocket(wsUrl, protocols); this.ws.addEventListener('open', (event) => { console.log('WebSocket连接已建立');