Coverage for zenmqtt/mqtt/protocol/core.py: 96%
191 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 21:55 +0000
1import asyncio
2from asyncio import Task
3from logging import getLogger
4from typing import Optional, Sequence, assert_never
6from zenmqtt.connection import MQTTConnection
7from zenmqtt.exceptions import ConnectionLostError, ProtocolError
8from zenmqtt.instrumentation.metrics import MetricsCollector, guard_metrics
9from zenmqtt.mqtt.auth import AuthResult
10from zenmqtt.mqtt.codec import IncomingPacket, decode
11from zenmqtt.mqtt.connect import (
12 SESSION_PRESENT_FLAG,
13 ConnectionResult,
14 ConnectProperties,
15 DisconnectProperties,
16 ServerDisconnect,
17 WillMessage,
18 pack_connect_packet,
19 pack_disconnect_packet,
20)
21from zenmqtt.mqtt.limits import ServerLimits
22from zenmqtt.mqtt.ping import PingRespResult
23from zenmqtt.mqtt.protocol.commands import Commands
24from zenmqtt.mqtt.protocol.context import ProtocolContext
25from zenmqtt.mqtt.protocol.incoming import IncomingFlow
26from zenmqtt.mqtt.protocol.keepalive import KeepAlive
27from zenmqtt.mqtt.protocol.outgoing import OutgoingFlow
28from zenmqtt.mqtt.protocol.stream import PacketReader
29from zenmqtt.mqtt.publish import (
30 MQTTMessage,
31 PubAckResult,
32 PubCompResult,
33 PublishAcknowledgement,
34 PublishProperties,
35 PubRecResult,
36 PubRelResult,
37)
38from zenmqtt.mqtt.reason_codes import FAILURE_REASON_CODE, IMPLEMENTATION_SPECIFIC_ERROR
39from zenmqtt.mqtt.subscribe import (
40 SubscribeResult,
41 SubscriptionProperties,
42 SubscriptionRequest,
43 UnsubscribeProperties,
44 UnsubscribeResult,
45)
46from zenmqtt.session import MQTTSession
48logger = getLogger(__name__)
51class MQTTProtocol:
52 """
53 MQTT 5 client protocol: the connection lifecycle and the read loop, which
54 passes incoming packets to the parts of the protocol: outgoing and
55 incoming PUBLISH flows, SUBSCRIBE/UNSUBSCRIBE and keep alive.
56 """
58 def __init__(
59 self,
60 session: MQTTSession,
61 metrics: Optional[MetricsCollector] = None,
62 ) -> None:
63 """
64 A lost connection doesn't end anything: QoS 1/2 publish calls keep
65 waiting for the acknowledgement of the message re-sent on the next
66 connect, and next_message() keeps waiting. The owner calls close()
67 when nothing will connect anymore.
68 """
69 self._session = session
70 self._context = ProtocolContext(guard_metrics(metrics))
72 self._incoming = IncomingFlow(self._context, session)
73 self._outgoing = OutgoingFlow(self._context, session)
74 self._commands = Commands(self._context, session)
75 self._keep_alive = KeepAlive(self._context)
77 self._read_loop_task: Optional[Task[None]] = None
78 self._packet_reader: Optional[PacketReader] = None
79 self._connection_future: Optional[asyncio.Future[ConnectionResult]] = None
81 # True if the connection is closed by disconnect call
82 self._disconnecting = False
84 # DISCONNECT of the server for the current/last connection
85 self.server_disconnect: Optional[ServerDisconnect] = None
87 # "Authentication Method" of the last CONNECT
88 self._authentication_method: Optional[str] = None
90 @property
91 def is_connected(self) -> bool:
92 """True between successful CONNACK and loss of the connection."""
93 return self._context.connected
95 def set_connection(self, connection: MQTTConnection) -> None:
96 self._context.set_connection(connection)
97 self._connection_future = asyncio.get_running_loop().create_future()
98 self._disconnecting = False
99 self.server_disconnect = None
101 self._packet_reader = PacketReader(connection)
102 self._read_loop_task = asyncio.create_task(
103 self._read_loop(connection, self._packet_reader),
104 name="mqtt-protocol-read-loop",
105 )
107 async def connect(
108 self,
109 client_id: str,
110 username: Optional[str],
111 password: Optional[str],
112 clean_start: bool = False,
113 keepalive: int = 0,
114 properties: Optional[ConnectProperties] = None,
115 will: Optional[WillMessage] = None,
116 ) -> ConnectionResult:
117 """
118 :param keepalive: seconds between control packets sent by the client,
119 0 disables keep alive; "Server Keep Alive" from CONNACK overrides it
120 """
121 connection = self._context.connection
122 assert connection
123 assert self._connection_future
125 properties = properties or {}
127 self._authentication_method = properties.get("authentication_method")
128 self._incoming.on_new_connection(properties)
130 if self._packet_reader: 130 ↛ 135line 130 didn't jump to line 135 because the condition on line 130 was always true
131 self._packet_reader.maximum_packet_size = properties.get(
132 "maximum_packet_size"
133 )
135 if clean_start:
136 await self._session.reset()
138 login_packet = pack_connect_packet(
139 client_id,
140 username,
141 password,
142 clean_start=clean_start,
143 keepalive=keepalive,
144 properties=properties,
145 will=will,
146 )
148 logger.debug("mqtt_protocol.send_connect_packet")
149 sent_at = self._context.now()
150 await self._context.write(connection, login_packet)
152 connection_result = await self._connection_future
154 self._context.metrics.on_connect(
155 connection_result.reason_code, self._context.now() - sent_at
156 )
158 if connection_result.reason_code >= FAILURE_REASON_CODE:
159 return connection_result
161 if not connection_result.flags & SESSION_PRESENT_FLAG:
162 await self._discard_session()
164 self._context.set_server_limits(
165 ServerLimits.from_connack(connection_result.properties)
166 )
168 await self._outgoing.on_connected(connection_result)
170 self._context.mark_connected()
172 self._keep_alive.start(
173 connection_result.properties.get("server_keep_alive", keepalive)
174 )
176 return connection_result
178 async def disconnect(
179 self, reason_code: int, properties: Optional[DisconnectProperties] = None
180 ) -> None:
181 connection = self._context.connection
182 assert connection
184 self._disconnecting = True
186 if not connection.is_closing():
187 logger.debug("mqtt_protocol.send_disconnect_packet")
188 await self._context.write(
189 connection, pack_disconnect_packet(reason_code, properties or {})
190 )
192 await connection.disconnect()
194 if read_loop_task := self._read_loop_task:
195 await read_loop_task
196 self._read_loop_task = None
198 async def wait_outgoing_completed(self) -> None:
199 """
200 Waits until no outgoing QoS 1/2 message is in flight; see
201 MQTTClient.disconnect(timeout=...).
202 """
203 await self._outgoing.wait_completed()
205 async def wait_closed(self) -> None:
206 """Waits until the current connection is closed and handled."""
207 if read_loop_task := self._read_loop_task: 207 ↛ exitline 207 didn't return from function 'wait_closed' because the condition on line 207 was always true
208 await asyncio.shield(read_loop_task)
210 async def close(self, exc: Exception) -> None:
211 """
212 Nothing will connect anymore: fails publish calls which wait for
213 acknowledgements and ends the messages.
214 """
215 self._outgoing.close(exc)
216 self._incoming.close()
218 async def publish(
219 self,
220 topic: str,
221 payload: bytes,
222 qos: int = 0,
223 retain: bool = False,
224 properties: Optional[PublishProperties] = None,
225 ) -> Optional[PublishAcknowledgement]:
226 return await self._outgoing.publish(topic, payload, qos, retain, properties)
228 async def next_message(self) -> Optional[MQTTMessage]:
229 """
230 The next incoming message, None after close(); QoS 1/2 messages must
231 be acknowledged by ack(). Meant for one reader, see MQTTClient.messages.
232 """
233 return await self._incoming.next_message()
235 async def ack(self, message: MQTTMessage, reason_code: int = 0) -> None:
236 await self._incoming.ack(message, reason_code)
238 async def subscribe(
239 self,
240 topics: Sequence[SubscriptionRequest],
241 properties: Optional[SubscriptionProperties] = None,
242 ) -> SubscribeResult:
243 return await self._commands.subscribe(topics, properties)
245 async def unsubscribe(
246 self, topics: Sequence[str], properties: Optional[UnsubscribeProperties] = None
247 ) -> UnsubscribeResult:
248 return await self._commands.unsubscribe(topics, properties)
250 async def ping(self) -> None:
251 await self._keep_alive.ping()
253 async def _discard_session(self) -> None:
254 """The server has no session, so the client must discard its own."""
255 if pending := await self._session.get_pending_outgoing_messages():
256 logger.warning(
257 "mqtt_protocol.session_not_present.discard_pending_messages "
258 "count:%s, set session_expiry_interval > 0 to keep them",
259 len(pending),
260 )
262 await self._session.reset()
264 self._outgoing.on_session_lost()
266 async def _handle_connack(self, connection_result: ConnectionResult) -> None:
267 assert self._connection_future
269 logger.debug("mqtt_protocol.handle_connack_packet packet:%s", connection_result)
271 if not self._connection_future.done():
272 self._connection_future.set_result(connection_result)
274 async def _handle_disconnect(self, disconnect_packet: ServerDisconnect) -> None:
275 connection = self._context.connection
276 assert connection
278 if disconnect_packet.reason_code >= FAILURE_REASON_CODE: 278 ↛ 283line 278 didn't jump to line 283 because the condition on line 278 was always true
279 logger.warning(
280 "mqtt_protocol.handle_disconnect_packet packet:%s", disconnect_packet
281 )
282 else:
283 logger.debug(
284 "mqtt_protocol.handle_disconnect_packet packet:%s", disconnect_packet
285 )
287 self.server_disconnect = disconnect_packet
289 # pending futures are failed by the read loop once the stream ends
290 await connection.disconnect()
292 async def _handle_auth(self, packet: AuthResult) -> None:
293 # the server sends AUTH only if CONNECT had an "Authentication Method"
294 # (MQTT 5, 4.12)
295 if self._authentication_method is None:
296 raise ProtocolError("AUTH without authentication method in CONNECT")
298 # closes the connection with "Implementation specific error"
299 raise NotImplementedError("Enhanced authentication isn't supported")
301 async def _read_loop(
302 self, connection: MQTTConnection, packet_reader: PacketReader
303 ) -> None:
304 try:
305 while packet := await packet_reader.read_packet():
306 header, reader, size = packet
308 logger.debug(
309 "mqtt_protocol.read_loop.new_packet_received header:%s", header
310 )
312 self._context.metrics.on_packet_received(header.packet_type, size)
314 try:
315 await self._dispatch(decode(header, reader))
316 except ProtocolError:
317 raise
318 except Exception as exc:
319 # e.g. the session storage failed: the state of the flow
320 # is unknown, so the connection is closed; after reconnect
321 # the session and the server re-send what isn't completed
322 logger.error(
323 "mqtt_protocol.handle_incoming_packet.error", exc_info=exc
324 )
325 await self._abort(connection, IMPLEMENTATION_SPECIFIC_ERROR)
326 return
327 except ProtocolError as exc:
328 logger.error("mqtt_protocol.protocol_error", exc_info=exc)
329 await self._abort(connection, exc.reason_code)
330 except Exception as exc:
331 logger.error("mqtt_protocol.read_loop.error", exc_info=exc)
332 finally:
333 await self._handle_connection_lost(connection)
335 async def _dispatch(self, packet: IncomingPacket) -> None:
336 match packet:
337 case MQTTMessage():
338 await self._incoming.handle_publish(packet)
339 case PubAckResult():
340 await self._outgoing.handle_puback(packet)
341 case PubRecResult():
342 await self._outgoing.handle_pubrec(packet)
343 case PubRelResult():
344 await self._incoming.handle_pubrel(packet)
345 case PubCompResult():
346 await self._outgoing.handle_pubcomp(packet)
347 case SubscribeResult():
348 await self._commands.handle_suback(packet)
349 case UnsubscribeResult():
350 await self._commands.handle_unsuback(packet)
351 case PingRespResult():
352 await self._keep_alive.handle_pingresp(packet)
353 case ConnectionResult():
354 await self._handle_connack(packet)
355 case ServerDisconnect():
356 await self._handle_disconnect(packet)
357 case AuthResult(): 357 ↛ 359line 357 didn't jump to line 359 because the pattern on line 357 always matched
358 await self._handle_auth(packet)
359 case _:
360 assert_never(packet)
362 async def _abort(self, connection: MQTTConnection, reason_code: int) -> None:
363 await self._context.write(connection, pack_disconnect_packet(reason_code, {}))
364 await connection.disconnect()
366 async def _handle_connection_lost(self, connection: MQTTConnection) -> None:
367 logger.debug("mqtt_protocol.handle_connection_lost")
369 # everything up to the first await is synchronous, so no new
370 # futures can be registered after the parts handled the loss
371 if self._context.connected:
372 self._context.metrics.on_connection_closed(lost=not self._disconnecting)
374 self._context.mark_disconnected()
375 exc = ConnectionLostError(self.server_disconnect)
377 if self._connection_future and not self._connection_future.done(): 377 ↛ 378line 377 didn't jump to line 378 because the condition on line 377 was never true
378 self._connection_future.set_exception(exc)
380 self._keep_alive.on_connection_lost(exc)
381 self._outgoing.on_connection_lost(exc)
382 command_packet_identifiers = self._commands.on_connection_lost(exc)
383 self._incoming.on_connection_lost()
385 # publish packets stay in the session, commands aren't re-sent
386 for packet_identifier in command_packet_identifiers:
387 await self._session.release_packet_identifier(packet_identifier)
389 if not connection.is_closing():
390 await connection.disconnect()