fix: avoid websocket token leakage in logs

This commit is contained in:
2026-06-22 11:34:35 +08:00
parent 0793eb82d6
commit c5d9bac0ed
2 changed files with 57 additions and 30 deletions
@@ -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):
"""断开连接"""
+10 -9
View File
@@ -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连接已建立');