xiaozhi-esp32-server/main/xiaozhi-server/core/protocols/mqtt_protocol.py

import asyncio
from typing import Dict, Any, Callable
from config.logger import setup_logging

logger = setup_logging()


# MQTT 固定头部的类型
class PacketType:
    CONNECT = 1
    CONNACK = 2
    PUBLISH = 3
    SUBSCRIBE = 8
    SUBACK = 9
    PINGREQ = 12
    PINGRESP = 13
    DISCONNECT = 14


class MQTTProtocol:
    """
    MQTT协议处理器：负责MQTT协议的解析和封装
    """
    
    def __init__(self, socket):
        self.socket = socket
        self.buffer = b''
        self.event_handlers = {}
        self.is_connected = False
        self.keep_alive_interval = 0
        self.last_activity = 0
        
        # 启动消息处理任务
        self._processing_task = asyncio.create_task(self._process_messages())
    
    def on(self, event: str, handler: Callable):
        """注册事件处理器"""
        self.event_handlers[event] = handler
    
    def emit(self, event: str, *args, **kwargs):
        """触发事件"""
        handler = self.event_handlers.get(event)
        if handler:
            if asyncio.iscoroutinefunction(handler):
                asyncio.create_task(handler(*args, **kwargs))
            else:
                handler(*args, **kwargs)
    
    async def _process_messages(self):
        """处理消息的主循环"""
        try:
            while True:
                # 从socket读取数据
                data = await self._read_socket()
                if not data:
                    break
                
                # 添加到缓冲区
                self.buffer += data
                
                # 处理缓冲区中的消息
                await self._process_buffer()
                
        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.error(f"MQTT消息处理循环出错: {e}")
            self.emit('error', e)
        finally:
            self.emit('close')
    
    async def _read_socket(self) -> bytes:
        """从socket读取数据"""
        try:
            # 使用asyncio的socket读取
            loop = asyncio.get_event_loop()
            data = await loop.sock_recv(self.socket, 4096)
            return data
        except Exception as e:
            logger.error(f"读取socket数据失败: {e}")
            return b''
    
    async def _process_buffer(self):
        """处理缓冲区中的消息"""
        while len(self.buffer) >= 2:  # 至少需要2字节开始解析
            try:
                # 解析消息
                message_length, message = self._parse_message()
                if message_length == 0:
                    break  # 消息不完整，等待更多数据
                
                # 从缓冲区移除已处理的消息
                self.buffer = self.buffer[message_length:]
                
                # 处理消息
                await self._handle_message(message)
                
            except Exception as e:
                logger.error(f"处理MQTT消息失败: {e}")
                self.emit('protocolError', e)
                break
    
    def _parse_message(self) -> tuple[int, Dict[str, Any]]:
        """解析MQTT消息"""
        if len(self.buffer) < 2:
            return 0, {}
        
        # 获取消息类型
        first_byte = self.buffer[0]
        packet_type = (first_byte >> 4)
        
        # 解析剩余长度
        remaining_length, bytes_read = self._decode_remaining_length()
        if remaining_length == -1:
            return 0, {}  # 长度解析失败，等待更多数据
        
        # 计算完整消息长度
        total_length = 1 + bytes_read + remaining_length
        
        if len(self.buffer) < total_length:
            return 0, {}  # 消息不完整
        
        # 提取消息数据
        message_data = self.buffer[:total_length]
        
        # 根据消息类型解析
        if packet_type == PacketType.CONNECT:
            message = self._parse_connect(message_data)
        elif packet_type == PacketType.PUBLISH:
            message = self._parse_publish(message_data)
        elif packet_type == PacketType.SUBSCRIBE:
            message = self._parse_subscribe(message_data)
        elif packet_type == PacketType.PINGREQ:
            message = {'type': 'pingreq'}
        elif packet_type == PacketType.DISCONNECT:
            message = {'type': 'disconnect'}
        else:
            logger.warning(f"未处理的MQTT消息类型: {packet_type}")
            message = {'type': 'unknown', 'packet_type': packet_type}
        
        return total_length, message
    
    def _decode_remaining_length(self) -> tuple[int, int]:
        """解码剩余长度字段"""
        multiplier = 1
        value = 0
        bytes_read = 0
        
        while bytes_read < 4 and bytes_read + 1 < len(self.buffer):
            digit = self.buffer[bytes_read + 1]
            bytes_read += 1
            
            value += (digit & 127) * multiplier
            multiplier *= 128
            
            if (digit & 128) == 0:
                break
        else:
            if bytes_read >= 4:
                return -1, 0  # 长度字段过长
            return -1, 0  # 数据不完整
        
        return value, bytes_read
    
    def _encode_remaining_length(self, length: int) -> bytes:
        """编码剩余长度字段"""
        result = bytearray()
        
        while True:
            digit = length % 128
            length = length // 128
            
            if length > 0:
                digit |= 0x80
            
            result.append(digit)
            
            if length == 0:
                break
        
        return bytes(result)
    
    def _parse_connect(self, message_data: bytes) -> Dict[str, Any]:
        """解析CONNECT消息"""
        try:
            # 跳过固定头部和剩余长度
            _, bytes_read = self._decode_remaining_length()
            pos = 1 + bytes_read
            
            # 协议名长度
            protocol_length = int.from_bytes(message_data[pos:pos+2], 'big')
            pos += 2
            
            # 协议名
            protocol = message_data[pos:pos+protocol_length].decode('utf-8')
            pos += protocol_length
            
            # 协议级别
            protocol_level = message_data[pos]
            pos += 1
            
            # 连接标志
            connect_flags = message_data[pos]
            has_username = (connect_flags & 0x80) != 0
            has_password = (connect_flags & 0x40) != 0
            pos += 1
            
            # 保持连接时间
            keep_alive = int.from_bytes(message_data[pos:pos+2], 'big')
            pos += 2
            
            # 客户端ID
            client_id_length = int.from_bytes(message_data[pos:pos+2], 'big')
            pos += 2
            client_id = message_data[pos:pos+client_id_length].decode('utf-8')
            pos += client_id_length
            
            # 用户名（如果存在）
            username = ''
            if has_username:
                username_length = int.from_bytes(message_data[pos:pos+2], 'big')
                pos += 2
                username = message_data[pos:pos+username_length].decode('utf-8')
                pos += username_length
            
            # 密码（如果存在）
            password = ''
            if has_password:
                password_length = int.from_bytes(message_data[pos:pos+2], 'big')
                pos += 2
                password = message_data[pos:pos+password_length].decode('utf-8')
                pos += password_length
            
            return {
                'type': 'connect',
                'protocol': protocol,
                'protocolLevel': protocol_level,
                'clientId': client_id,
                'keepAlive': keep_alive,
                'username': username,
                'password': password
            }
            
        except Exception as e:
            logger.error(f"解析CONNECT消息失败: {e}")
            raise
    
    def _parse_publish(self, message_data: bytes) -> Dict[str, Any]:
        """解析PUBLISH消息"""
        try:
            # 获取QoS等标志
            first_byte = message_data[0]
            qos = (first_byte & 0x06) >> 1
            dup = (first_byte & 0x08) != 0
            retain = (first_byte & 0x01) != 0
            
            # 跳过固定头部和剩余长度
            _, bytes_read = self._decode_remaining_length()
            pos = 1 + bytes_read
            
            # 主题长度
            topic_length = int.from_bytes(message_data[pos:pos+2], 'big')
            pos += 2
            
            # 主题
            topic = message_data[pos:pos+topic_length].decode('utf-8')
            pos += topic_length
            
            # 消息ID（QoS > 0时存在）
            packet_id = None
            if qos > 0:
                packet_id = int.from_bytes(message_data[pos:pos+2], 'big')
                pos += 2
            
            # 有效载荷
            payload = message_data[pos:].decode('utf-8')
            
            return {
                'type': 'publish',
                'topic': topic,
                'payload': payload,
                'qos': qos,
                'dup': dup,
                'retain': retain,
                'packetId': packet_id
            }
            
        except Exception as e:
            logger.error(f"解析PUBLISH消息失败: {e}")
            raise
    
    def _parse_subscribe(self, message_data: bytes) -> Dict[str, Any]:
        """解析SUBSCRIBE消息"""
        try:
            # 跳过固定头部和剩余长度
            _, bytes_read = self._decode_remaining_length()
            pos = 1 + bytes_read
            
            # 消息ID
            packet_id = int.from_bytes(message_data[pos:pos+2], 'big')
            pos += 2
            
            # 主题长度
            topic_length = int.from_bytes(message_data[pos:pos+2], 'big')
            pos += 2
            
            # 主题
            topic = message_data[pos:pos+topic_length].decode('utf-8')
            pos += topic_length
            
            # QoS
            qos = message_data[pos]
            
            return {
                'type': 'subscribe',
                'packetId': packet_id,
                'topic': topic,
                'qos': qos
            }
            
        except Exception as e:
            logger.error(f"解析SUBSCRIBE消息失败: {e}")
            raise
    
    async def _handle_message(self, message: Dict[str, Any]):
        """处理解析后的消息"""
        message_type = message.get('type')
        
        if message_type == 'connect':
            self.keep_alive_interval = message.get('keepAlive', 0)
            self.is_connected = True
            self.emit('connect', message)
        elif message_type == 'publish':
            self.emit('publish', message)
        elif message_type == 'subscribe':
            self.emit('subscribe', message)
        elif message_type == 'pingreq':
            await self.send_pingresp()
        elif message_type == 'disconnect':
            self.emit('disconnect')
        else:
            logger.warning(f"未处理的消息类型: {message_type}")
    
    async def send_connack(self, return_code: int = 0, session_present: bool = False):
        """发送CONNACK消息"""
        packet = bytearray([
            PacketType.CONNACK << 4,  # 固定头部
            2,  # 剩余长度
            1 if session_present else 0,  # 连接确认标志
            return_code  # 返回码
        ])
        
        await self._send_packet(packet)
    
    async def send_publish(self, topic: str, payload: str, qos: int = 0, 
                          dup: bool = False, retain: bool = False, packet_id: int = None):
        """发送PUBLISH消息"""
        # 构造固定头部
        first_byte = PacketType.PUBLISH << 4
        if dup:
            first_byte |= 0x08
        if qos > 0:
            first_byte |= (qos << 1)
        if retain:
            first_byte |= 0x01
        
        # 构造可变头部和载荷
        topic_bytes = topic.encode('utf-8')
        payload_bytes = payload.encode('utf-8')
        
        variable_header = bytearray()
        variable_header.extend(len(topic_bytes).to_bytes(2, 'big'))
        variable_header.extend(topic_bytes)
        
        if qos > 0 and packet_id is not None:
            variable_header.extend(packet_id.to_bytes(2, 'big'))
        
        # 计算剩余长度
        remaining_length = len(variable_header) + len(payload_bytes)
        remaining_length_bytes = self._encode_remaining_length(remaining_length)
        
        # 构造完整消息
        packet = bytearray([first_byte])
        packet.extend(remaining_length_bytes)
        packet.extend(variable_header)
        packet.extend(payload_bytes)
        
        await self._send_packet(packet)
    
    async def send_suback(self, packet_id: int, return_code: int = 0):
        """发送SUBACK消息"""
        packet = bytearray([
            PacketType.SUBACK << 4,  # 固定头部
            3,  # 剩余长度
            packet_id >> 8,  # 消息ID高字节
            packet_id & 0xFF,  # 消息ID低字节
            return_code  # 返回码
        ])
        
        await self._send_packet(packet)
    
    async def send_pingresp(self):
        """发送PINGRESP消息"""
        packet = bytearray([
            PacketType.PINGRESP << 4,  # 固定头部
            0  # 剩余长度
        ])
        
        await self._send_packet(packet)
    
    async def _send_packet(self, packet: bytearray):
        """发送数据包"""
        try:
            loop = asyncio.get_event_loop()
            await loop.sock_sendall(self.socket, bytes(packet))
        except Exception as e:
            logger.error(f"发送MQTT数据包失败: {e}")
            raise
    
    async def close(self):
        """关闭协议处理器"""
        if hasattr(self, '_processing_task') and not self._processing_task.done():
            self._processing_task.cancel()
            try:
                await self._processing_task
            except asyncio.CancelledError:
                pass
        
        try:
            self.socket.close()
        except Exception as e:
            logger.error(f"关闭socket失败: {e}")