mirror of
https://github.com/xinnan-tech/xiaozhi-esp32-server.git
synced 2026-07-24 08:03:53 +08:00
315 lines
11 KiB
Python
315 lines
11 KiB
Python
import asyncio
|
|
import socket
|
|
import time
|
|
from typing import Dict, Any, Set
|
|
from config.logger import setup_logging
|
|
from core.protocols.mqtt_connection import MQTTConnection
|
|
from core.transport.mqtt_transport import MQTTTransport, UDPAudioHandler
|
|
from core.services.connection_service import ConnectionService
|
|
|
|
logger = setup_logging()
|
|
|
|
|
|
class MQTTServer:
|
|
"""
|
|
原生MQTT服务器：直接处理MQTT协议连接
|
|
集成到xiaozhi-server架构中
|
|
"""
|
|
|
|
def __init__(self, config: Dict[str, Any]):
|
|
self.config = config
|
|
self.logger = setup_logging()
|
|
|
|
# 服务器配置
|
|
server_config = config.get('mqtt_server', {})
|
|
self.mqtt_port = server_config.get('port', 1883)
|
|
self.udp_port = server_config.get('udp_port', self.mqtt_port)
|
|
self.host = server_config.get('host', '0.0.0.0')
|
|
self.public_ip = server_config.get('public_ip', 'localhost')
|
|
|
|
# 连接管理
|
|
self.connections: Dict[int, MQTTConnection] = {}
|
|
self.udp_handlers: Dict[int, UDPAudioHandler] = {}
|
|
self.connection_id_counter = 0
|
|
|
|
# 服务器实例
|
|
self.mqtt_server = None
|
|
self.udp_server = None
|
|
|
|
# 连接服务
|
|
self.connection_service = ConnectionService(config)
|
|
|
|
# 活跃连接管理
|
|
self.active_transports: Set[MQTTTransport] = set()
|
|
|
|
# 心跳检查
|
|
self.heartbeat_task = None
|
|
self.heartbeat_interval = 30 # 30秒检查一次
|
|
|
|
async def start(self):
|
|
"""启动MQTT服务器"""
|
|
try:
|
|
# 启动MQTT TCP服务器
|
|
await self._start_mqtt_server()
|
|
|
|
# 启动UDP服务器
|
|
await self._start_udp_server()
|
|
|
|
# 启动心跳检查
|
|
self.heartbeat_task = asyncio.create_task(self._heartbeat_check())
|
|
|
|
logger.info(f"MQTT服务器启动成功: {self.host}:{self.mqtt_port}")
|
|
logger.info(f"UDP服务器启动成功: {self.host}:{self.udp_port}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"启动MQTT服务器失败: {e}")
|
|
raise
|
|
|
|
async def _start_mqtt_server(self):
|
|
"""启动MQTT TCP服务器"""
|
|
self.mqtt_server = await asyncio.start_server(
|
|
self._handle_mqtt_connection,
|
|
self.host,
|
|
self.mqtt_port
|
|
)
|
|
|
|
async def _start_udp_server(self):
|
|
"""启动UDP服务器"""
|
|
loop = asyncio.get_event_loop()
|
|
|
|
# 创建UDP socket
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
sock.bind((self.host, self.udp_port))
|
|
sock.setblocking(False)
|
|
|
|
# 创建UDP协议处理器
|
|
transport, protocol = await loop.create_datagram_endpoint(
|
|
lambda: UDPProtocol(self),
|
|
sock=sock
|
|
)
|
|
|
|
self.udp_server = (transport, protocol)
|
|
|
|
async def _handle_mqtt_connection(self, reader, writer):
|
|
"""处理新的MQTT连接"""
|
|
connection_id = self._generate_connection_id()
|
|
|
|
try:
|
|
# 获取客户端地址
|
|
client_addr = writer.get_extra_info('peername')
|
|
logger.info(f"新MQTT连接: {client_addr}, connection_id: {connection_id}")
|
|
|
|
# 创建MQTT连接处理器
|
|
mqtt_connection = MQTTConnection(
|
|
writer.get_extra_info('socket'),
|
|
connection_id,
|
|
self
|
|
)
|
|
|
|
# 创建UDP音频处理器
|
|
udp_handler = UDPAudioHandler(
|
|
connection_id,
|
|
self,
|
|
{} # 加密配置将在hello消息中设置
|
|
)
|
|
|
|
# 创建MQTT传输层
|
|
transport = MQTTTransport(mqtt_connection, udp_handler)
|
|
|
|
# 注册连接
|
|
self.connections[connection_id] = mqtt_connection
|
|
self.udp_handlers[connection_id] = udp_handler
|
|
self.active_transports.add(transport)
|
|
|
|
# 提取连接头信息
|
|
headers = {
|
|
'x-real-ip': client_addr[0] if client_addr else 'unknown',
|
|
'connection-type': 'mqtt'
|
|
}
|
|
|
|
try:
|
|
# 使用ConnectionService处理连接
|
|
await self.connection_service.handle_connection(transport, headers)
|
|
|
|
except Exception as e:
|
|
logger.error(f"ConnectionService处理MQTT连接失败: {e}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"处理MQTT连接失败: {e}")
|
|
finally:
|
|
# 清理连接
|
|
await self._cleanup_connection(connection_id)
|
|
|
|
async def _cleanup_connection(self, connection_id: int):
|
|
"""清理连接资源"""
|
|
try:
|
|
# 移除连接
|
|
if connection_id in self.connections:
|
|
connection = self.connections.pop(connection_id)
|
|
await connection.close()
|
|
|
|
# 移除UDP处理器
|
|
if connection_id in self.udp_handlers:
|
|
udp_handler = self.udp_handlers.pop(connection_id)
|
|
await udp_handler.close()
|
|
|
|
# 移除传输层（通过连接ID查找）
|
|
transports_to_remove = []
|
|
for transport in self.active_transports:
|
|
if hasattr(transport, '_mqtt_connection') and \
|
|
transport._mqtt_connection.connection_id == connection_id:
|
|
transports_to_remove.append(transport)
|
|
|
|
for transport in transports_to_remove:
|
|
self.active_transports.discard(transport)
|
|
await transport.close()
|
|
|
|
logger.info(f"MQTT连接清理完成: {connection_id}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"清理MQTT连接失败: {e}")
|
|
|
|
def _generate_connection_id(self) -> int:
|
|
"""生成连接ID"""
|
|
self.connection_id_counter += 1
|
|
return self.connection_id_counter
|
|
|
|
async def on_client_connected(self, mqtt_connection: MQTTConnection):
|
|
"""客户端连接回调"""
|
|
logger.info(f"MQTT客户端已连接: {mqtt_connection.client_id}")
|
|
|
|
async def on_client_disconnected(self, mqtt_connection: MQTTConnection):
|
|
"""客户端断开连接回调"""
|
|
logger.info(f"MQTT客户端已断开: {mqtt_connection.client_id}")
|
|
|
|
async def send_udp_message(self, data: bytes, remote_addr: tuple):
|
|
"""发送UDP消息"""
|
|
if self.udp_server:
|
|
transport, protocol = self.udp_server
|
|
transport.sendto(data, remote_addr)
|
|
|
|
async def send_encrypted_audio(self, connection_id: int, audio_data: bytes,
|
|
timestamp: int, remote_addr: tuple, encryption_config: Dict[str, Any]):
|
|
"""发送加密音频数据"""
|
|
try:
|
|
# 这里应该实现音频数据加密逻辑
|
|
# 暂时直接发送原始数据
|
|
header = self._generate_udp_header(connection_id, len(audio_data), timestamp, 0)
|
|
message = header + audio_data
|
|
|
|
await self.send_udp_message(message, remote_addr)
|
|
|
|
except Exception as e:
|
|
logger.error(f"发送加密音频失败: {e}")
|
|
|
|
def _generate_udp_header(self, connection_id: int, length: int, timestamp: int, sequence: int) -> bytes:
|
|
"""生成UDP消息头"""
|
|
header = bytearray(16)
|
|
header[0] = 1 # type
|
|
header[2:4] = length.to_bytes(2, 'big') # payload length
|
|
header[4:8] = connection_id.to_bytes(4, 'big') # connection id
|
|
header[8:12] = timestamp.to_bytes(4, 'big') # timestamp
|
|
header[12:16] = sequence.to_bytes(4, 'big') # sequence
|
|
return bytes(header)
|
|
|
|
async def _heartbeat_check(self):
|
|
"""心跳检查任务"""
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(self.heartbeat_interval)
|
|
|
|
# 检查所有连接的状态
|
|
dead_connections = []
|
|
for connection_id, connection in self.connections.items():
|
|
if not connection.is_connected():
|
|
dead_connections.append(connection_id)
|
|
|
|
# 清理死连接
|
|
for connection_id in dead_connections:
|
|
logger.info(f"清理死连接: {connection_id}")
|
|
await self._cleanup_connection(connection_id)
|
|
|
|
# 记录活跃连接数
|
|
active_count = len(self.connections)
|
|
if active_count > 0:
|
|
logger.info(f"MQTT活跃连接数: {active_count}")
|
|
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as e:
|
|
logger.error(f"心跳检查任务出错: {e}")
|
|
|
|
async def stop(self):
|
|
"""停止MQTT服务器"""
|
|
logger.info("正在停止MQTT服务器...")
|
|
|
|
# 停止心跳检查
|
|
if self.heartbeat_task and not self.heartbeat_task.done():
|
|
self.heartbeat_task.cancel()
|
|
try:
|
|
await self.heartbeat_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
# 关闭所有连接
|
|
for connection_id in list(self.connections.keys()):
|
|
await self._cleanup_connection(connection_id)
|
|
|
|
# 关闭UDP服务器
|
|
if self.udp_server:
|
|
transport, protocol = self.udp_server
|
|
transport.close()
|
|
|
|
# 关闭MQTT服务器
|
|
if self.mqtt_server:
|
|
self.mqtt_server.close()
|
|
await self.mqtt_server.wait_closed()
|
|
|
|
logger.info("MQTT服务器已停止")
|
|
|
|
def get_server_status(self) -> Dict[str, Any]:
|
|
"""获取服务器状态"""
|
|
return {
|
|
'type': 'mqtt',
|
|
'host': self.host,
|
|
'mqtt_port': self.mqtt_port,
|
|
'udp_port': self.udp_port,
|
|
'active_connections': len(self.connections),
|
|
'active_transports': len(self.active_transports)
|
|
}
|
|
|
|
|
|
class UDPProtocol(asyncio.DatagramProtocol):
|
|
"""UDP协议处理器"""
|
|
|
|
def __init__(self, mqtt_server: MQTTServer):
|
|
self.mqtt_server = mqtt_server
|
|
self.transport = None
|
|
|
|
def connection_made(self, transport):
|
|
self.transport = transport
|
|
|
|
def datagram_received(self, data: bytes, addr: tuple):
|
|
"""接收UDP数据报"""
|
|
try:
|
|
# 解析UDP消息头
|
|
if len(data) < 16:
|
|
return
|
|
|
|
connection_id = int.from_bytes(data[4:8], 'big')
|
|
timestamp = int.from_bytes(data[8:12], 'big')
|
|
sequence = int.from_bytes(data[12:16], 'big')
|
|
payload = data[16:]
|
|
|
|
# 找到对应的UDP处理器
|
|
udp_handler = self.mqtt_server.udp_handlers.get(connection_id)
|
|
if udp_handler:
|
|
udp_handler.on_udp_message(payload, timestamp, addr)
|
|
|
|
except Exception as e:
|
|
logger.error(f"处理UDP数据报失败: {e}")
|
|
|
|
def error_received(self, exc):
|
|
logger.error(f"UDP协议错误: {exc}")
|
|
|