feat: 添加WS TOKEN 认证功能

core/auth.py

@@ -0,0 +1,60 @@
+# xiaozhi-esp32-server-main/core/auth.py
+import logging
+
+logger = logging.getLogger(__name__)
+
+class AuthenticationError(Exception):
+    """认证异常"""
+    pass
+
+class AuthMiddleware:
+    def __init__(self, config):
+        self.config = config
+        self.auth_config = config["server"].get("auth", {})
+        # 构建token查找表
+        self.tokens = {
+            item["token"]: item["name"] 
+            for item in self.auth_config.get("tokens", [])
+        }
+        # 设备白名单
+        self.allowed_devices = set(
+            self.auth_config.get("allowed_devices", [])
+        )
+        
+    async def authenticate(self, headers):
+        """验证连接请求"""
+        # 检查是否启用认证
+        if not self.auth_config.get("enabled", False):
+            return True
+            
+        # 验证Authorization header
+        auth_header = headers.get("Authorization", "")
+        if not auth_header.startswith("Bearer "):
+            logger.error("Missing or invalid Authorization header")
+            raise AuthenticationError("Missing or invalid Authorization header")
+        
+        token = auth_header.split(" ")[1]
+        if token not in self.tokens:
+            logger.error(f"Invalid token: {token}")
+            raise AuthenticationError("Invalid token")
+            
+        # 验证Device-Id
+        device_id = headers.get("Device-Id")
+        if not device_id:
+            logger.error("Missing Device-Id header")
+            raise AuthenticationError("Missing Device-Id header")
+            
+        # 检查设备白名单
+        if self.allowed_devices and device_id not in self.allowed_devices:
+            logger.error(f"Device not in whitelist: {device_id}")
+            raise AuthenticationError("Device not in whitelist")
+            
+        logger.info(f"Authentication successful - Device: {device_id}, Token: {self.tokens[token]}")
+        return True
+        
+    def get_token_name(self, token):
+        """获取token对应的设备名称"""
+        return self.tokens.get(token)
+
+
+

core/connection.py

@@ -15,12 +15,14 @@ from core.handle.textHandle import handleTextMessage
 from core.utils.util import get_string_no_punctuation_or_emoji
 from concurrent.futures import ThreadPoolExecutor, TimeoutError
 from core.handle.audioHandle import handleAudioMessage, sendAudioMessage
+from .auth import AuthMiddleware, AuthenticationError
 
 
 class ConnectionHandler:
     def __init__(self, config: Dict[str, Any], _vad, _asr, _llm, _tts):
         self.config = config
         self.logger = logging.getLogger(__name__)
+        self.auth = AuthMiddleware(config)
 
         self.websocket = None
         self.headers = None
@@ -67,27 +69,42 @@ class ConnectionHandler:
         self.tts_duration = 0
 
     async def handle_connection(self, ws):
-        self.websocket = ws
-        """处理单个WebSocket连接"""
-        self.headers = dict(self.websocket.request.headers)
-        self.logger.info(f"连接建立，请求头：\n{self.headers}")
-
-        self.welcome_msg = self.config["xiaozhi"]
-        self.session_id = str(uuid.uuid4())
-        self.welcome_msg["session_id"] = self.session_id
-        await self.websocket.send(json.dumps(self.welcome_msg))
-
-        await self.loop.run_in_executor(None, self._initialize_components)
-
-        tts_priority = threading.Thread(target=self._priority_thread, daemon=True)
-        tts_priority.start()
-
         try:
-            async for message in self.websocket:
-                await self._route_message(message)
-        except websockets.exceptions.ConnectionClosed:
-            self.logger.info("客户端断开连接")
-            await self.close()
+            # 获取并验证headers
+            self.headers = dict(ws.request.headers)
+            self.logger.info(f"New connection request - Headers: {self.headers}")
+            
+            # 进行认证
+            await self.auth.authenticate(self.headers)
+            
+            # 认证通过,继续处理
+            self.websocket = ws
+            self.session_id = str(uuid.uuid4())
+            
+            self.welcome_msg = self.config["xiaozhi"]
+            self.welcome_msg["session_id"] = self.session_id
+            await self.websocket.send(json.dumps(self.welcome_msg))
+
+            await self.loop.run_in_executor(None, self._initialize_components)
+
+            tts_priority = threading.Thread(target=self._priority_thread, daemon=True)
+            tts_priority.start()
+
+            try:
+                async for message in self.websocket:
+                    await self._route_message(message)
+            except websockets.exceptions.ConnectionClosed:
+                self.logger.info("客户端断开连接")
+                await self.close()
+                
+        except AuthenticationError as e:
+            self.logger.error(f"Authentication failed: {str(e)}")
+            await ws.close()
+            return
+        except Exception as e:
+            self.logger.error(f"Connection error: {str(e)}")
+            await ws.close()
+            return
 
     async def _route_message(self, message):
         """消息路由"""