mirror of
https://github.com/xinnan-tech/xiaozhi-esp32-server.git
synced 2026-07-28 10:03:54 +08:00
feat: add optional native mqtt and udp transport
This commit is contained in:
caixypromise
@@ -0,0 +1,649 @@
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from typing import Dict, Any, Optional, Callable
|
||||
from config.logger import setup_logging
|
||||
from core.utils.mqtt_auth import validate_mqtt_credentials
|
||||
|
||||
logger = setup_logging()
|
||||
|
||||
|
||||
class MQTTConnection:
|
||||
"""
|
||||
MQTT连接处理类：管理单个MQTT客户端连接
|
||||
处理MQTT协议消息和会话管理
|
||||
"""
|
||||
|
||||
def __init__(self, socket, connection_id: int, mqtt_server, udp_handler=None, reader=None, writer=None):
|
||||
self.socket = socket
|
||||
self.reader = reader
|
||||
self.writer = writer
|
||||
self.connection_id = connection_id
|
||||
self.mqtt_server = mqtt_server
|
||||
self.udp_handler = udp_handler
|
||||
|
||||
# 连接信息
|
||||
self.client_id = None
|
||||
self.device_id = None
|
||||
self.username = None
|
||||
self.password = None
|
||||
self.session_id = None
|
||||
|
||||
# 协议状态
|
||||
self.is_connected_flag = False
|
||||
self.keep_alive_interval = 0
|
||||
self.last_activity = time.monotonic()
|
||||
|
||||
# 消息处理
|
||||
self.message_callback = None
|
||||
self.reply_topic = None
|
||||
|
||||
# UDP相关
|
||||
self.udp_config = None
|
||||
self.message_queue_size = int(
|
||||
getattr(mqtt_server, "message_queue_size", 128) or 128
|
||||
)
|
||||
self.business_ready_timeout = float(
|
||||
getattr(mqtt_server, "business_ready_timeout", 30) or 30
|
||||
)
|
||||
self.close_timeout = max(
|
||||
0.1, float(getattr(mqtt_server, "close_timeout", 2) or 2)
|
||||
)
|
||||
self.goodbye_timeout = max(
|
||||
0.1, float(getattr(mqtt_server, "goodbye_timeout", 1) or 1)
|
||||
)
|
||||
|
||||
# 任务管理
|
||||
self.keep_alive_task = None
|
||||
self.business_task = None
|
||||
self._closed = False
|
||||
self._close_task = None
|
||||
self._close_initiator = None
|
||||
self._close_complete = asyncio.Event()
|
||||
self.connect_processed_event = asyncio.Event()
|
||||
self.connect_accepted = False
|
||||
self.business_ready_event = asyncio.Event()
|
||||
self._hello_business_ready_event = None
|
||||
self._hello_business_session_id = None
|
||||
self._logical_hello_received = False
|
||||
self._startup_recovery_task = None
|
||||
self._session_transition_lock = asyncio.Lock()
|
||||
self._last_goodbye_session_id = None
|
||||
self._goodbye_lock = asyncio.Lock()
|
||||
|
||||
# 创建MQTT协议处理器
|
||||
from core.protocols.mqtt_protocol import MQTTProtocol
|
||||
self.protocol = MQTTProtocol(
|
||||
socket=socket,
|
||||
reader=reader,
|
||||
writer=writer,
|
||||
max_payload_size=getattr(mqtt_server, "max_payload_size", 8192),
|
||||
event_queue_size=self.message_queue_size,
|
||||
close_timeout=self.close_timeout,
|
||||
)
|
||||
self._setup_protocol_handlers()
|
||||
|
||||
def _setup_protocol_handlers(self):
|
||||
"""设置协议事件处理"""
|
||||
self.protocol.on('connect', self._handle_connect)
|
||||
self.protocol.on('publish', self._handle_publish)
|
||||
self.protocol.on('subscribe', self._handle_subscribe)
|
||||
self.protocol.on('disconnect', self._handle_disconnect)
|
||||
self.protocol.on('close', self._handle_close)
|
||||
self.protocol.on('error', self._handle_error)
|
||||
self.protocol.on('protocolError', self._handle_error)
|
||||
self.protocol.on('activity', self._handle_activity)
|
||||
|
||||
def _handle_activity(self):
|
||||
"""更新最近活动时间（用于心跳保活）"""
|
||||
self.last_activity = time.monotonic()
|
||||
|
||||
async def _handle_connect(self, connect_data: Dict[str, Any]):
|
||||
"""处理CONNECT消息"""
|
||||
try:
|
||||
if self.connect_processed_event.is_set() or self.is_connected_flag:
|
||||
logger.warning("同一TCP连接收到重复CONNECT，关闭连接")
|
||||
await self.close()
|
||||
return
|
||||
self.client_id = connect_data['clientId']
|
||||
self.username = connect_data.get('username')
|
||||
self.password = connect_data.get('password')
|
||||
self.keep_alive_interval = connect_data.get('keepAlive', 0) * 1000 # 转换为毫秒
|
||||
|
||||
logger.info(f"MQTT客户端连接: {self.client_id}")
|
||||
|
||||
try:
|
||||
validate_mqtt_credentials(
|
||||
self.client_id,
|
||||
self.username,
|
||||
self.password,
|
||||
self.mqtt_server.signature_key,
|
||||
)
|
||||
except ValueError as e:
|
||||
logger.warning(f"MQTT客户端认证失败: {self.client_id}, {e}")
|
||||
await self.protocol.send_connack(4)
|
||||
self.connect_processed_event.set()
|
||||
await self.close()
|
||||
return False
|
||||
|
||||
# 解析客户端ID获取设备信息
|
||||
if not self._parse_client_id():
|
||||
await self.protocol.send_connack(1) # 连接被拒绝
|
||||
self.connect_processed_event.set()
|
||||
await self.close()
|
||||
return False
|
||||
|
||||
# 生成会话ID
|
||||
self.session_id = str(uuid.uuid4())
|
||||
|
||||
# 设置回复主题
|
||||
self.reply_topic = f"devices/p2p/{self.device_id.replace(':', '_')}"
|
||||
|
||||
async def complete_acceptance():
|
||||
# Keep this inside the server's per-client takeover lock. A
|
||||
# newer connection cannot reclaim this clientId between the
|
||||
# success CONNACK and publishing the active owner.
|
||||
self.is_connected_flag = True
|
||||
try:
|
||||
await self.protocol.send_connack(0)
|
||||
except Exception:
|
||||
self.is_connected_flag = False
|
||||
raise
|
||||
if self.keep_alive_interval > 0:
|
||||
self.keep_alive_task = asyncio.create_task(
|
||||
self._keep_alive_check()
|
||||
)
|
||||
self.connect_accepted = True
|
||||
|
||||
accepted = await self.mqtt_server.on_client_connected(
|
||||
self, complete_acceptance
|
||||
)
|
||||
if not accepted:
|
||||
if not self._closed:
|
||||
await self.protocol.send_connack(3) # 服务端暂不可用
|
||||
self.connect_processed_event.set()
|
||||
await self.close()
|
||||
return False
|
||||
self.connect_processed_event.set()
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"处理CONNECT消息失败: {e}")
|
||||
self.connect_processed_event.set()
|
||||
await self.close()
|
||||
return False
|
||||
|
||||
def _parse_client_id(self) -> bool:
|
||||
"""解析客户端ID获取设备信息"""
|
||||
try:
|
||||
# 支持格式: GID_test@@@mac_address@@@uuid 或 GID_test@@@mac_address
|
||||
parts = self.client_id.split('@@@')
|
||||
|
||||
if len(parts) >= 2:
|
||||
self.group_id = parts[0]
|
||||
# 将设备标识统一为服务端使用的MAC地址格式
|
||||
self.device_id = (
|
||||
parts[1].replace('_', ':').replace('-', ':').lower()
|
||||
)
|
||||
|
||||
if len(parts) >= 3:
|
||||
self.uuid = parts[2]
|
||||
|
||||
return True
|
||||
else:
|
||||
logger.error(f"无效的客户端ID格式: {self.client_id}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"解析客户端ID失败: {e}")
|
||||
return False
|
||||
|
||||
async def _handle_publish(self, publish_data: Dict[str, Any]):
|
||||
"""处理PUBLISH消息"""
|
||||
try:
|
||||
if publish_data.get('qos', 0) != 0:
|
||||
logger.warning(
|
||||
f"不支持的MQTT QoS级别: {publish_data.get('qos')}"
|
||||
)
|
||||
await self.close()
|
||||
return
|
||||
topic = publish_data['topic']
|
||||
payload = publish_data['payload']
|
||||
|
||||
logger.debug(f"收到MQTT发布消息: topic={topic}, payload={payload}")
|
||||
|
||||
# 更新活动时间
|
||||
self.last_activity = time.monotonic()
|
||||
|
||||
# 解析JSON消息
|
||||
try:
|
||||
message_data = json.loads(payload)
|
||||
|
||||
# 处理不同类型的消息
|
||||
if message_data.get('type') == 'hello':
|
||||
if message_data.get('version', 3) != 3:
|
||||
logger.warning(
|
||||
f"不支持的MQTT协议版本: {message_data.get('version')}"
|
||||
)
|
||||
await self.close()
|
||||
return
|
||||
await self._handle_hello_message(message_data)
|
||||
else:
|
||||
# 其他消息通过回调处理
|
||||
if self.message_callback:
|
||||
self.message_callback(topic, payload)
|
||||
|
||||
except json.JSONDecodeError:
|
||||
logger.error(f"MQTT消息JSON解析失败: {payload}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"处理PUBLISH消息失败: {e}")
|
||||
|
||||
async def _handle_hello_message(self, message_data: Dict[str, Any]):
|
||||
"""处理hello消息，初始化UDP配置"""
|
||||
async with self._session_transition_lock:
|
||||
self._logical_hello_received = True
|
||||
try:
|
||||
handle_logical_hello = getattr(
|
||||
self.mqtt_server, "handle_logical_hello", None
|
||||
)
|
||||
if callable(handle_logical_hello):
|
||||
await handle_logical_hello(self)
|
||||
# Match the gateway contract: do not advertise a usable audio
|
||||
# channel until private config and runtime components are ready.
|
||||
await asyncio.wait_for(
|
||||
self.business_ready_event.wait(),
|
||||
timeout=self.business_ready_timeout,
|
||||
)
|
||||
if self._closed or not self.is_connected_flag:
|
||||
return
|
||||
hello_reply = self._prepare_hello_reply(
|
||||
message_data.get('audio_params', {}),
|
||||
message_data.get('version', 3),
|
||||
)
|
||||
hello_ready = asyncio.Event()
|
||||
self._hello_business_ready_event = hello_ready
|
||||
self._hello_business_session_id = self.session_id
|
||||
|
||||
# Enqueue the logical-session boundary before the device can react
|
||||
# to the reply with UDP audio.
|
||||
if self.message_callback:
|
||||
try:
|
||||
self.message_callback(self.reply_topic, json.dumps(message_data))
|
||||
except Exception as e:
|
||||
logger.error(f"转发hello消息失败: {e}")
|
||||
|
||||
# Long-lived MQTT connections can change Agent configuration at
|
||||
# each logical Hello. Do not expose the new UDP session until the
|
||||
# business runtime has either refreshed or deliberately retained
|
||||
# the previous healthy runtime.
|
||||
await asyncio.wait_for(
|
||||
hello_ready.wait(),
|
||||
timeout=self.business_ready_timeout,
|
||||
)
|
||||
if self._closed or not self.is_connected_flag:
|
||||
return
|
||||
await self.send_message(self.reply_topic, json.dumps(hello_reply))
|
||||
|
||||
logger.info(f"MQTT Hello消息处理完成: {self.client_id}")
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
logger.error(
|
||||
f"MQTT Hello等待业务运行时超时: {self.client_id}, "
|
||||
f"timeout={self.business_ready_timeout}s"
|
||||
)
|
||||
await self.close()
|
||||
except Exception as e:
|
||||
logger.error(f"处理hello消息失败: {e}")
|
||||
finally:
|
||||
self._hello_business_ready_event = None
|
||||
self._hello_business_session_id = None
|
||||
|
||||
def schedule_stale_session_recovery(self, delay: float = 1.0) -> None:
|
||||
"""Return a reconnected device with a stale UDP session to Idle."""
|
||||
if self._startup_recovery_task is not None:
|
||||
return
|
||||
self._startup_recovery_task = asyncio.create_task(
|
||||
self._recover_stale_session(delay)
|
||||
)
|
||||
|
||||
async def _recover_stale_session(self, delay: float) -> None:
|
||||
try:
|
||||
await asyncio.sleep(max(0.0, delay))
|
||||
async with self._session_transition_lock:
|
||||
if (
|
||||
self._closed
|
||||
or not self.is_connected_flag
|
||||
or self._logical_hello_received
|
||||
or not self.reply_topic
|
||||
):
|
||||
return
|
||||
# No session id is intentional: firmware accepts this as a
|
||||
# connection-level reset and discards an UDP session owned by
|
||||
# a previous server process. Serialize it with Hello so this
|
||||
# reset can never overtake a newly negotiated session.
|
||||
await self.send_message(
|
||||
self.reply_topic,
|
||||
json.dumps({"type": "goodbye"}),
|
||||
)
|
||||
logger.info("已通知重连MQTT设备清理旧UDP会话: {}", self.client_id)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.warning("通知重连MQTT设备清理旧会话失败: {}", e)
|
||||
|
||||
def mark_business_session_ready(self, session_id: str = None) -> None:
|
||||
"""Acknowledge readiness for the currently pending logical Hello."""
|
||||
event = self._hello_business_ready_event
|
||||
pending_session_id = self._hello_business_session_id
|
||||
if event is None:
|
||||
return
|
||||
if session_id is not None and pending_session_id != session_id:
|
||||
return
|
||||
event.set()
|
||||
|
||||
async def send_hello_reply(self, audio_params: Dict[str, Any], version: int = 3):
|
||||
"""发送hello回复（可在未收到设备hello时调用）"""
|
||||
hello_reply = self._prepare_hello_reply(audio_params, version)
|
||||
await self.send_message(self.reply_topic, json.dumps(hello_reply))
|
||||
|
||||
def _prepare_hello_reply(self, audio_params: Dict[str, Any], version: int = 3):
|
||||
"""Create and install one UDP session without publishing it yet."""
|
||||
import os
|
||||
|
||||
self.session_id = str(uuid.uuid4())
|
||||
self._last_goodbye_session_id = None
|
||||
udp_session_id = self.mqtt_server.bind_udp_session(
|
||||
self, self.udp_handler
|
||||
)
|
||||
nonce = self._generate_udp_header(
|
||||
0, 0, 0, connection_id=udp_session_id
|
||||
)
|
||||
self.udp_config = {
|
||||
'key': os.urandom(16),
|
||||
'encryption': 'aes-128-ctr',
|
||||
'server': self.mqtt_server.public_endpoint,
|
||||
'port': self.mqtt_server.udp_port,
|
||||
'nonce': nonce,
|
||||
'local_sequence': 0,
|
||||
'remote_sequence': 0
|
||||
}
|
||||
if self.udp_handler:
|
||||
self.udp_handler.configure_encryption(self.udp_config)
|
||||
|
||||
configured_audio_params = (
|
||||
getattr(self.mqtt_server, 'config', {})
|
||||
.get('xiaozhi', {})
|
||||
.get('audio_params', {})
|
||||
)
|
||||
hello_reply = {
|
||||
'type': 'hello',
|
||||
'version': version,
|
||||
'session_id': self.session_id,
|
||||
'transport': 'udp',
|
||||
'udp': {
|
||||
'server': self.udp_config['server'],
|
||||
'port': self.udp_config['port'],
|
||||
'encryption': self.udp_config['encryption'],
|
||||
'key': self.udp_config['key'].hex(),
|
||||
'nonce': nonce.hex()
|
||||
},
|
||||
'audio_params': configured_audio_params or audio_params or {}
|
||||
}
|
||||
|
||||
return hello_reply
|
||||
|
||||
def _generate_udp_header(
|
||||
self, length: int, timestamp: int, sequence: int,
|
||||
connection_id: int = None
|
||||
) -> bytes:
|
||||
header = bytearray(16)
|
||||
header[0] = 1 # type
|
||||
header[2:4] = length.to_bytes(2, 'big')
|
||||
udp_connection_id = connection_id or self.connection_id
|
||||
header[4:8] = udp_connection_id.to_bytes(4, 'big')
|
||||
header[8:12] = timestamp.to_bytes(4, 'big')
|
||||
header[12:16] = sequence.to_bytes(4, 'big')
|
||||
return bytes(header)
|
||||
|
||||
async def _handle_subscribe(self, subscribe_data: Dict[str, Any]):
|
||||
"""处理SUBSCRIBE消息"""
|
||||
try:
|
||||
topic = subscribe_data['topic']
|
||||
packet_id = subscribe_data['packetId']
|
||||
|
||||
logger.debug(f"客户端订阅主题: {topic}")
|
||||
|
||||
# 发送订阅确认
|
||||
await self.protocol.send_suback(packet_id, 0)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"处理SUBSCRIBE消息失败: {e}")
|
||||
|
||||
async def _handle_disconnect(self):
|
||||
"""处理DISCONNECT消息"""
|
||||
logger.info(f"客户端主动断开连接: {self.client_id}")
|
||||
await self.close()
|
||||
|
||||
async def _handle_close(self):
|
||||
"""处理连接关闭"""
|
||||
logger.info(f"MQTT连接关闭: {self.client_id}")
|
||||
await self.close()
|
||||
|
||||
async def _handle_error(self, error):
|
||||
"""处理连接错误"""
|
||||
logger.error(f"MQTT连接错误: {self.client_id}, error: {error}")
|
||||
await self.close()
|
||||
|
||||
async def _keep_alive_check(self):
|
||||
"""心跳检查任务"""
|
||||
try:
|
||||
while self.is_connected_flag and not self._closed:
|
||||
await asyncio.sleep(self.keep_alive_interval / 1000 / 2) # 检查间隔为心跳间隔的一半
|
||||
|
||||
current_time = time.monotonic()
|
||||
if current_time - self.last_activity > self.keep_alive_interval / 1000 * 1.5:
|
||||
logger.info(f"MQTT客户端心跳超时: {self.client_id}")
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self.notify_device_idle(),
|
||||
timeout=self.goodbye_timeout,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"MQTT心跳超时发送goodbye失败，继续关闭连接: {}", e
|
||||
)
|
||||
finally:
|
||||
await self.close()
|
||||
break
|
||||
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.error(f"心跳检查任务出错: {e}")
|
||||
|
||||
def set_message_callback(self, callback: Callable[[str, str], None]):
|
||||
"""设置消息接收回调"""
|
||||
self.message_callback = callback
|
||||
|
||||
async def send_message(self, topic: str, payload: str):
|
||||
"""发送MQTT消息"""
|
||||
if self._closed or not self.is_connected_flag:
|
||||
raise RuntimeError("MQTT connection is closed")
|
||||
|
||||
try:
|
||||
await self.protocol.send_publish(topic, payload, qos=0)
|
||||
logger.debug(f"发送MQTT消息: topic={topic}, payload={payload}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"发送MQTT消息失败: {e}")
|
||||
raise
|
||||
|
||||
async def notify_device_idle(self, session_id: str = None) -> bool:
|
||||
"""Send one session-scoped goodbye before the physical MQTT close."""
|
||||
async with self._goodbye_lock:
|
||||
target_session_id = session_id or self.session_id
|
||||
if (
|
||||
self._closed
|
||||
or not self.is_connected_flag
|
||||
or not self.udp_config
|
||||
or not self.reply_topic
|
||||
or not target_session_id
|
||||
or self._last_goodbye_session_id == target_session_id
|
||||
):
|
||||
return False
|
||||
|
||||
await self.send_message(
|
||||
self.reply_topic,
|
||||
json.dumps(
|
||||
{"type": "goodbye", "session_id": target_session_id}
|
||||
),
|
||||
)
|
||||
self._last_goodbye_session_id = target_session_id
|
||||
return True
|
||||
|
||||
def is_connected(self) -> bool:
|
||||
"""检查连接状态"""
|
||||
return self.is_connected_flag and not self._closed
|
||||
|
||||
async def close(self):
|
||||
"""关闭连接"""
|
||||
if not self._closed:
|
||||
self._closed = True
|
||||
self.is_connected_flag = False
|
||||
self.connect_processed_event.set()
|
||||
self.business_ready_event.set()
|
||||
if self._hello_business_ready_event is not None:
|
||||
self._hello_business_ready_event.set()
|
||||
# Run cleanup in a dedicated task. Shielding it lets a cancelled
|
||||
# first caller leave without falsely completing the close barrier;
|
||||
# later callers can still await the same cleanup owner.
|
||||
self._close_initiator = asyncio.current_task()
|
||||
self._close_task = asyncio.create_task(self._close_impl())
|
||||
|
||||
close_task = self._close_task
|
||||
if close_task is None:
|
||||
return
|
||||
|
||||
current_task = asyncio.current_task()
|
||||
dependency_tasks = {
|
||||
close_task,
|
||||
getattr(self.protocol, "_processing_task", None),
|
||||
getattr(self.protocol, "_dispatch_task", None),
|
||||
}
|
||||
if self.business_task is not self._close_initiator:
|
||||
dependency_tasks.add(self.business_task)
|
||||
if self.keep_alive_task is not self._close_initiator:
|
||||
dependency_tasks.add(self.keep_alive_task)
|
||||
# The dedicated closer can be waiting for these tasks. Let them unwind
|
||||
# instead of creating a reverse wait cycle.
|
||||
if current_task in dependency_tasks:
|
||||
return
|
||||
await asyncio.shield(close_task)
|
||||
|
||||
async def _close_impl(self):
|
||||
"""Own and complete physical cleanup independently of caller lifetime."""
|
||||
try:
|
||||
current_task = asyncio.current_task()
|
||||
|
||||
# The socket/protocol task can time out while ConnectionService is
|
||||
# blocked in private config or component initialization. Cancel the
|
||||
# owning server task so its finally block releases the SessionContext
|
||||
# and any partially initialized runtime instead of leaking per retry.
|
||||
if (
|
||||
self.business_task
|
||||
and self.business_task is not current_task
|
||||
and self.business_task is not self._close_initiator
|
||||
and not self.business_task.done()
|
||||
):
|
||||
self.business_task.cancel()
|
||||
# The business owner's finally block calls back into server
|
||||
# cleanup. Waiting for it here would create a close cycle:
|
||||
# close -> business finally -> transport.close -> close.
|
||||
# Give cancellation one loop turn, then let it unwind
|
||||
# independently while the physical socket is released.
|
||||
await asyncio.sleep(0)
|
||||
if not self.business_task.done():
|
||||
tracker = getattr(
|
||||
self.mqtt_server, "track_draining_task", None
|
||||
)
|
||||
if callable(tracker):
|
||||
tracker(self.business_task, "MQTT业务任务")
|
||||
else:
|
||||
self.business_task.add_done_callback(
|
||||
lambda task: self._consume_background_task(
|
||||
task, "MQTT业务任务"
|
||||
)
|
||||
)
|
||||
|
||||
# 取消心跳检查任务
|
||||
if (
|
||||
self.keep_alive_task
|
||||
and self.keep_alive_task is not current_task
|
||||
and self.keep_alive_task is not self._close_initiator
|
||||
and not self.keep_alive_task.done()
|
||||
):
|
||||
self.keep_alive_task.cancel()
|
||||
await self._wait_cancelled_task(
|
||||
self.keep_alive_task, "MQTT心跳任务"
|
||||
)
|
||||
|
||||
if (
|
||||
self._startup_recovery_task
|
||||
and self._startup_recovery_task is not current_task
|
||||
and not self._startup_recovery_task.done()
|
||||
):
|
||||
self._startup_recovery_task.cancel()
|
||||
await self._wait_cancelled_task(
|
||||
self._startup_recovery_task, "MQTT会话恢复任务"
|
||||
)
|
||||
|
||||
# 关闭协议处理器
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self.protocol.close(),
|
||||
timeout=self.close_timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning("关闭MQTT协议处理器超时，强制中止socket")
|
||||
abort = getattr(self.protocol, "abort", None)
|
||||
if callable(abort):
|
||||
abort()
|
||||
except (asyncio.CancelledError, Exception) as e:
|
||||
logger.error(f"关闭MQTT协议处理器失败: {e}")
|
||||
|
||||
# Publish disconnect only after the physical protocol close barrier.
|
||||
try:
|
||||
await self.mqtt_server.on_client_disconnected(self)
|
||||
except (asyncio.CancelledError, Exception) as e:
|
||||
logger.error(f"通知服务器连接关闭失败: {e}")
|
||||
|
||||
logger.info(f"MQTT连接已关闭: {self.client_id}")
|
||||
finally:
|
||||
self._close_complete.set()
|
||||
|
||||
async def _wait_cancelled_task(self, task: asyncio.Task, label: str) -> None:
|
||||
done, _ = await asyncio.wait({task}, timeout=self.close_timeout)
|
||||
if task not in done:
|
||||
logger.warning(
|
||||
"{}取消后{}秒仍未退出，继续释放物理连接",
|
||||
label,
|
||||
self.close_timeout,
|
||||
)
|
||||
return
|
||||
try:
|
||||
task.result()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception as exc:
|
||||
logger.error(f"{label}退出失败: {exc}")
|
||||
|
||||
@staticmethod
|
||||
def _consume_background_task(task: asyncio.Task, label: str) -> None:
|
||||
try:
|
||||
task.result()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception as exc:
|
||||
logger.error(f"{label}退出失败: {exc}")
|
||||
@@ -0,0 +1,651 @@
|
||||
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=None,
|
||||
reader=None,
|
||||
writer=None,
|
||||
max_payload_size=8192,
|
||||
event_queue_size=128,
|
||||
close_timeout=2,
|
||||
):
|
||||
self.socket = socket
|
||||
self.reader = reader
|
||||
self.writer = writer
|
||||
self.buffer = b''
|
||||
self.event_handlers = {}
|
||||
self.is_connected = False
|
||||
self.keep_alive_interval = 0
|
||||
self.last_activity = 0
|
||||
self.max_payload_size = int(max_payload_size or 0)
|
||||
self.close_timeout = max(0.1, float(close_timeout or 2))
|
||||
self._closed = False
|
||||
self._application_queue = asyncio.Queue(
|
||||
maxsize=max(1, int(event_queue_size or 128))
|
||||
)
|
||||
|
||||
# Application publishes stay ordered, while PINGREQ remains on the read
|
||||
# loop so a slow Hello/runtime refresh cannot starve MQTT keepalive.
|
||||
self._dispatch_task = asyncio.create_task(
|
||||
self._dispatch_application_messages()
|
||||
)
|
||||
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 emit_async(self, event: str, *args, **kwargs):
|
||||
"""Emit protocol events in packet order."""
|
||||
handler = self.event_handlers.get(event)
|
||||
if not handler:
|
||||
return None
|
||||
result = handler(*args, **kwargs)
|
||||
if asyncio.iscoroutine(result):
|
||||
return await result
|
||||
return result
|
||||
|
||||
async def _process_messages(self):
|
||||
"""处理消息的主循环"""
|
||||
try:
|
||||
while not self._closed:
|
||||
# 从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:
|
||||
if not self._closed:
|
||||
if self._dispatch_task.done():
|
||||
await self.emit_async('close')
|
||||
else:
|
||||
# Preserve parsed QoS0 publishes before a normal peer EOF.
|
||||
# The Hello/runtime barrier is separately time-bounded.
|
||||
await self._application_queue.put({'type': 'peer_close'})
|
||||
|
||||
async def _dispatch_application_messages(self):
|
||||
"""Dispatch non-heartbeat packets sequentially outside the read loop."""
|
||||
try:
|
||||
while not self._closed:
|
||||
message = await self._application_queue.get()
|
||||
try:
|
||||
await self._dispatch_application_message(message)
|
||||
finally:
|
||||
self._application_queue.task_done()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.error(f"MQTT应用消息处理循环出错: {e}")
|
||||
self.emit('error', e)
|
||||
|
||||
async def _dispatch_application_message(self, message: Dict[str, Any]):
|
||||
message_type = message.get('type')
|
||||
if message_type == 'publish':
|
||||
await self.emit_async('publish', message)
|
||||
elif message_type == 'disconnect':
|
||||
await self.emit_async('disconnect')
|
||||
self.is_connected = False
|
||||
elif message_type == 'peer_close':
|
||||
await self.emit_async('close')
|
||||
else:
|
||||
raise ValueError(f"不支持的MQTT应用消息类型: {message_type}")
|
||||
|
||||
def _enqueue_application_message(self, message: Dict[str, Any]) -> None:
|
||||
try:
|
||||
self._application_queue.put_nowait(message)
|
||||
except asyncio.QueueFull as exc:
|
||||
raise ValueError("MQTT应用消息队列已满") from exc
|
||||
|
||||
async def _read_socket(self) -> bytes:
|
||||
"""从socket读取数据"""
|
||||
try:
|
||||
if self.reader is not None:
|
||||
return await self.reader.read(4096)
|
||||
# 使用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)
|
||||
client_packet_types = {
|
||||
PacketType.CONNECT,
|
||||
PacketType.PUBLISH,
|
||||
PacketType.SUBSCRIBE,
|
||||
PacketType.PINGREQ,
|
||||
PacketType.DISCONNECT,
|
||||
}
|
||||
if packet_type not in client_packet_types:
|
||||
raise ValueError(f"不支持的客户端MQTT消息类型: {packet_type}")
|
||||
fixed_flags = first_byte & 0x0F
|
||||
expected_flags = {
|
||||
PacketType.CONNECT: 0,
|
||||
PacketType.SUBSCRIBE: 2,
|
||||
PacketType.PINGREQ: 0,
|
||||
PacketType.DISCONNECT: 0,
|
||||
}
|
||||
if (
|
||||
packet_type in expected_flags
|
||||
and fixed_flags != expected_flags[packet_type]
|
||||
):
|
||||
raise ValueError(
|
||||
f"MQTT packet type {packet_type} has invalid fixed-header flags "
|
||||
f"0x{fixed_flags:x}"
|
||||
)
|
||||
if packet_type == PacketType.PUBLISH and ((first_byte >> 1) & 0x03) == 3:
|
||||
raise ValueError("MQTT PUBLISH QoS 3 is invalid")
|
||||
|
||||
# 解析剩余长度
|
||||
remaining_length, bytes_read = self._decode_remaining_length()
|
||||
if remaining_length == -1:
|
||||
return 0, {} # 长度解析失败，等待更多数据
|
||||
max_payload_size = getattr(self, "max_payload_size", 0)
|
||||
if max_payload_size > 0 and remaining_length > max_payload_size:
|
||||
raise ValueError(
|
||||
f"MQTT remaining length {remaining_length} exceeds limit {max_payload_size}"
|
||||
)
|
||||
|
||||
# 计算完整消息长度
|
||||
total_length = 1 + bytes_read + remaining_length
|
||||
|
||||
if len(self.buffer) < total_length:
|
||||
return 0, {} # 消息不完整
|
||||
|
||||
if (
|
||||
packet_type in (PacketType.PINGREQ, PacketType.DISCONNECT)
|
||||
and remaining_length != 0
|
||||
):
|
||||
raise ValueError(
|
||||
f"MQTT packet type {packet_type} requires remaining length 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:
|
||||
if bytes_read + 1 >= len(self.buffer):
|
||||
return -1, 0
|
||||
digit = self.buffer[bytes_read + 1]
|
||||
bytes_read += 1
|
||||
|
||||
value += (digit & 127) * multiplier
|
||||
multiplier *= 128
|
||||
|
||||
if (digit & 128) == 0:
|
||||
return value, bytes_read
|
||||
if bytes_read == 4:
|
||||
raise ValueError("MQTT remaining length字段超过4字节")
|
||||
|
||||
raise ValueError("MQTT remaining length字段无效")
|
||||
|
||||
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
|
||||
|
||||
def read_bytes(binary=False):
|
||||
nonlocal pos
|
||||
if pos + 2 > len(message_data):
|
||||
raise ValueError("MQTT CONNECT字符串长度字段不完整")
|
||||
value_length = int.from_bytes(message_data[pos:pos + 2], 'big')
|
||||
pos += 2
|
||||
if pos + value_length > len(message_data):
|
||||
raise ValueError("MQTT CONNECT字符串内容不完整")
|
||||
value = message_data[pos:pos + value_length]
|
||||
pos += value_length
|
||||
return value if binary else value.decode('utf-8')
|
||||
|
||||
protocol = read_bytes()
|
||||
|
||||
# 协议级别
|
||||
if pos + 4 > len(message_data):
|
||||
raise ValueError("MQTT CONNECT可变头部不完整")
|
||||
protocol_level = message_data[pos]
|
||||
pos += 1
|
||||
if protocol != 'MQTT' or protocol_level != 4:
|
||||
raise ValueError(
|
||||
f"不支持的MQTT协议: {protocol}/{protocol_level}"
|
||||
)
|
||||
|
||||
# 连接标志
|
||||
connect_flags = message_data[pos]
|
||||
if connect_flags & 0x01:
|
||||
raise ValueError("MQTT CONNECT保留标志必须为0")
|
||||
has_username = (connect_flags & 0x80) != 0
|
||||
has_password = (connect_flags & 0x40) != 0
|
||||
will_retain = (connect_flags & 0x20) != 0
|
||||
will_qos = (connect_flags >> 3) & 0x03
|
||||
has_will = (connect_flags & 0x04) != 0
|
||||
clean_session = (connect_flags & 0x02) != 0
|
||||
if has_password and not has_username:
|
||||
raise ValueError("MQTT CONNECT密码标志要求用户名标志")
|
||||
if will_qos == 3:
|
||||
raise ValueError("MQTT CONNECT Will QoS 3无效")
|
||||
if not has_will and (will_retain or will_qos):
|
||||
raise ValueError("MQTT CONNECT未启用Will但设置了Will标志")
|
||||
pos += 1
|
||||
|
||||
# 保持连接时间
|
||||
keep_alive = int.from_bytes(message_data[pos:pos+2], 'big')
|
||||
pos += 2
|
||||
|
||||
client_id = read_bytes()
|
||||
if not client_id and not clean_session:
|
||||
raise ValueError("MQTT CONNECT空clientId必须启用clean session")
|
||||
|
||||
if has_will:
|
||||
read_bytes() # Will topic
|
||||
read_bytes(binary=True) # Will payload
|
||||
|
||||
# 用户名（如果存在）
|
||||
username = ''
|
||||
if has_username:
|
||||
username = read_bytes()
|
||||
|
||||
# 密码（如果存在）
|
||||
password = ''
|
||||
if has_password:
|
||||
password = read_bytes(binary=True).decode('utf-8')
|
||||
|
||||
if pos != len(message_data):
|
||||
raise ValueError("MQTT CONNECT包含未解析的尾部数据")
|
||||
|
||||
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
|
||||
|
||||
# 主题长度
|
||||
if pos + 2 > len(message_data):
|
||||
raise ValueError("MQTT PUBLISH缺少主题长度")
|
||||
topic_length = int.from_bytes(message_data[pos:pos+2], 'big')
|
||||
pos += 2
|
||||
if topic_length == 0 or pos + topic_length > len(message_data):
|
||||
raise ValueError("MQTT PUBLISH主题为空或不完整")
|
||||
|
||||
# 主题
|
||||
topic = message_data[pos:pos+topic_length].decode('utf-8')
|
||||
pos += topic_length
|
||||
if "\x00" in topic or "+" in topic or "#" in topic:
|
||||
raise ValueError("MQTT PUBLISH主题名称无效")
|
||||
|
||||
# 消息ID（QoS > 0时存在）
|
||||
packet_id = None
|
||||
if qos > 0:
|
||||
if pos + 2 > len(message_data):
|
||||
raise ValueError("MQTT PUBLISH缺少packetId")
|
||||
packet_id = int.from_bytes(message_data[pos:pos+2], 'big')
|
||||
pos += 2
|
||||
if packet_id == 0:
|
||||
raise ValueError("MQTT PUBLISH packetId不能为0")
|
||||
|
||||
# 有效载荷
|
||||
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
|
||||
if pos + 2 > len(message_data):
|
||||
raise ValueError("MQTT SUBSCRIBE缺少packetId")
|
||||
packet_id = int.from_bytes(message_data[pos:pos+2], 'big')
|
||||
pos += 2
|
||||
if packet_id == 0:
|
||||
raise ValueError("MQTT SUBSCRIBE packetId不能为0")
|
||||
|
||||
# 主题长度
|
||||
if pos + 2 > len(message_data):
|
||||
raise ValueError("MQTT SUBSCRIBE缺少主题长度")
|
||||
topic_length = int.from_bytes(message_data[pos:pos+2], 'big')
|
||||
pos += 2
|
||||
if topic_length == 0 or pos + topic_length > len(message_data):
|
||||
raise ValueError("MQTT SUBSCRIBE主题为空或不完整")
|
||||
|
||||
# 主题
|
||||
topic = message_data[pos:pos+topic_length].decode('utf-8')
|
||||
pos += topic_length
|
||||
if "\x00" in topic:
|
||||
raise ValueError("MQTT SUBSCRIBE主题过滤器无效")
|
||||
|
||||
# QoS
|
||||
if pos >= len(message_data):
|
||||
raise ValueError("MQTT SUBSCRIBE缺少请求QoS")
|
||||
qos = message_data[pos]
|
||||
pos += 1
|
||||
if qos > 2:
|
||||
raise ValueError("MQTT SUBSCRIBE请求QoS无效")
|
||||
if pos != len(message_data):
|
||||
raise ValueError("MQTT SUBSCRIBE当前仅支持单个主题过滤器")
|
||||
|
||||
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':
|
||||
if self.is_connected:
|
||||
raise ValueError("MQTT连接只能发送一次CONNECT")
|
||||
self.keep_alive_interval = message.get('keepAlive', 0)
|
||||
accepted = await self.emit_async('connect', message)
|
||||
self.is_connected = accepted is not False
|
||||
if self.is_connected:
|
||||
self.emit('activity')
|
||||
return
|
||||
|
||||
if not self.is_connected:
|
||||
raise ValueError("MQTT客户端必须先发送CONNECT")
|
||||
|
||||
if message_type == 'publish':
|
||||
self.emit('activity')
|
||||
self._enqueue_application_message(message)
|
||||
elif message_type == 'subscribe':
|
||||
self.emit('activity')
|
||||
await self.emit_async('subscribe', message)
|
||||
elif message_type == 'pingreq':
|
||||
self.emit('activity')
|
||||
await self.send_pingresp()
|
||||
elif message_type == 'disconnect':
|
||||
self.emit('activity')
|
||||
self._enqueue_application_message(message)
|
||||
else:
|
||||
raise ValueError(f"不支持的MQTT消息类型: {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:
|
||||
if self.writer is not None:
|
||||
self.writer.write(bytes(packet))
|
||||
await self.writer.drain()
|
||||
else:
|
||||
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):
|
||||
"""关闭协议处理器"""
|
||||
self._closed = True
|
||||
current_task = asyncio.current_task()
|
||||
if (
|
||||
hasattr(self, '_processing_task')
|
||||
and self._processing_task is not current_task
|
||||
and not self._processing_task.done()
|
||||
):
|
||||
self._processing_task.cancel()
|
||||
try:
|
||||
await self._processing_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
if (
|
||||
hasattr(self, '_dispatch_task')
|
||||
and self._dispatch_task is not current_task
|
||||
and not self._dispatch_task.done()
|
||||
):
|
||||
self._dispatch_task.cancel()
|
||||
try:
|
||||
await self._dispatch_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
try:
|
||||
if self.writer is not None:
|
||||
self.writer.close()
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self.writer.wait_closed(),
|
||||
timeout=self.close_timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
"等待MQTT socket关闭超时，强制中止transport"
|
||||
)
|
||||
self.abort()
|
||||
except Exception:
|
||||
pass
|
||||
elif self.socket:
|
||||
self.socket.close()
|
||||
except Exception as e:
|
||||
logger.error(f"关闭socket失败: {e}")
|
||||
|
||||
def abort(self):
|
||||
"""Force-close the underlying transport when graceful close stalls."""
|
||||
if self.writer is not None:
|
||||
transport = getattr(self.writer, "transport", None)
|
||||
if transport is not None:
|
||||
transport.abort()
|
||||
return
|
||||
if self.socket:
|
||||
self.socket.close()
|
||||
Reference in New Issue
Block a user