mirror of
https://github.com/xinnan-tech/xiaozhi-esp32-server.git
synced 2026-07-30 13:53:56 +08:00
refactor: 重构底层代码，抽离conn，调整消息处理器并创建传输层接口。
feature: 支持mqtt非桥接版本。
This commit is contained in:
caixypromise
@@ -0,0 +1,313 @@
|
||||
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}")
|
||||
Reference in New Issue
Block a user