async def server_mock(stream): try: packet = await PublishPacket.from_stream(stream) self.assertEqual(packet.topic_name, '/topic') self.assertEqual(packet.qos, QOS_2) self.assertIsNotNone(packet.packet_id) self.assertIn(packet.packet_id, self.session.inflight_out) self.assertIn(packet.packet_id, self.handler._pubrec_waiters) pubrec = PubrecPacket.build(packet.packet_id) await pubrec.to_stream(stream) await PubrelPacket.from_stream(stream) self.assertIn(packet.packet_id, self.handler._pubcomp_waiters) pubcomp = PubcompPacket.build(packet.packet_id) await pubcomp.to_stream(stream) except Exception as ae: raise
def server_mock(reader, writer): try: packet = yield from PublishPacket.from_stream(reader) self.assertEqual(packet.topic_name, '/topic') self.assertEqual(packet.qos, QOS_2) self.assertIsNotNone(packet.packet_id) self.assertIn(packet.packet_id, self.session.inflight_out) self.assertIn(packet.packet_id, self.handler._pubrec_waiters) pubrec = PubrecPacket.build(packet.packet_id) yield from pubrec.to_stream(writer) pubrel = yield from PubrelPacket.from_stream(reader) self.assertIn(packet.packet_id, self.handler._pubcomp_waiters) pubcomp = PubcompPacket.build(packet.packet_id) yield from pubcomp.to_stream(writer) except Exception as ae: future.set_exception(ae)
def server_mock(reader, writer): try: packet = yield from PublishPacket.from_stream(reader) self.assertEquals(packet.topic_name, "/topic") self.assertEquals(packet.qos, QOS_2) self.assertIsNotNone(packet.packet_id) self.assertIn(packet.packet_id, self.session.inflight_out) self.assertIn(packet.packet_id, self.handler._pubrec_waiters) pubrec = PubrecPacket.build(packet.packet_id) yield from pubrec.to_stream(writer) pubrel = yield from PubrelPacket.from_stream(reader) self.assertIn(packet.packet_id, self.handler._pubcomp_waiters) pubcomp = PubcompPacket.build(packet.packet_id) yield from pubcomp.to_stream(writer) except Exception as ae: future.set_exception(ae)
def _handle_qos2_message_flow(self, app_message): """ Handle QOS_2 application message acknowledgment For incoming messages, this method stores the message, sends PUBREC, waits for PUBREL, initiate delivery and send PUBCOMP For outgoing messages, this methods sends PUBLISH, waits for PUBREC, discards messages and wait for PUBCOMP :param app_message: :return: """ assert app_message.qos == QOS_2 if app_message.direction == OUTGOING: if app_message.pubrel_packet and app_message.pubcomp_packet: raise HBMQTTException( "Message '%d' has already been acknowledged" % app_message.packet_id) if not app_message.pubrel_packet: # Store message if app_message.publish_packet is not None: # This is a retry flow, no need to store just check the message exists in session if app_message.packet_id not in self.session.inflight_out: raise HBMQTTException( "Unknown inflight message '%d' in session" % app_message.packet_id) publish_packet = app_message.build_publish_packet(dup=True) else: # Store message in session self.session.inflight_out[ app_message.packet_id] = app_message publish_packet = app_message.build_publish_packet() # Send PUBLISH packet yield from self._send_packet(publish_packet) app_message.publish_packet = publish_packet # Wait PUBREC if app_message.packet_id in self._pubrec_waiters: # PUBREC waiter already exists for this packet ID message = "Can't add PUBREC waiter, a waiter already exists for message Id '%s'" \ % app_message.packet_id self.logger.warning(message) raise HBMQTTException(message) waiter = asyncio.Future(loop=self._loop) self._pubrec_waiters[app_message.packet_id] = waiter yield from waiter del self._pubrec_waiters[app_message.packet_id] app_message.pubrec_packet = waiter.result() if not app_message.pubcomp_packet: # Send pubrel app_message.pubrel_packet = PubrelPacket.build( app_message.packet_id) yield from self._send_packet(app_message.pubrel_packet) # Wait for PUBCOMP waiter = asyncio.Future(loop=self._loop) self._pubcomp_waiters[app_message.packet_id] = waiter yield from waiter del self._pubcomp_waiters[app_message.packet_id] app_message.pubcomp_packet = waiter.result() # Discard inflight message del self.session.inflight_out[app_message.packet_id] elif app_message.direction == INCOMING: self.session.inflight_in[app_message.packet_id] = app_message # Send pubrec pubrec_packet = PubrecPacket.build(app_message.packet_id) yield from self._send_packet(pubrec_packet) app_message.pubrec_packet = pubrec_packet # Wait PUBREL if app_message.packet_id in self._pubrel_waiters and not self._pubrel_waiters[ app_message.packet_id].done(): # PUBREL waiter already exists for this packet ID message = "A waiter already exists for message Id '%s', canceling it" \ % app_message.packet_id self.logger.warning(message) self._pubrel_waiters[app_message.packet_id].cancel() try: waiter = asyncio.Future(loop=self._loop) self._pubrel_waiters[app_message.packet_id] = waiter yield from waiter del self._pubrel_waiters[app_message.packet_id] app_message.pubrel_packet = waiter.result() # Initiate delivery and discard message yield from self.session.delivered_message_queue.put( app_message) del self.session.inflight_in[app_message.packet_id] # Send pubcomp pubcomp_packet = PubcompPacket.build(app_message.packet_id) yield from self._send_packet(pubcomp_packet) app_message.pubcomp_packet = pubcomp_packet except asyncio.CancelledError: self.logger.debug("Message flow cancelled")
def _handle_qos2_message_flow(self, app_message): """ Handle QOS_2 application message acknowledgment For incoming messages, this method stores the message, sends PUBREC, waits for PUBREL, initiate delivery and send PUBCOMP For outgoing messages, this methods sends PUBLISH, waits for PUBREC, discards messages and wait for PUBCOMP :param app_message: :return: """ assert app_message.qos == QOS_2 if app_message.direction == OUTGOING: if app_message.pubrel_packet and app_message.pubcomp_packet: raise HBMQTTException("Message '%d' has already been acknowledged" % app_message.packet_id) if not app_message.pubrel_packet: # Store message if app_message.publish_packet is not None: # This is a retry flow, no need to store just check the message exists in session if app_message.packet_id not in self.session.inflight_out: raise HBMQTTException("Unknown inflight message '%d' in session" % app_message.packet_id) publish_packet = app_message.build_publish_packet(dup=True) else: # Store message in session self.session.inflight_out[app_message.packet_id] = app_message publish_packet = app_message.build_publish_packet() # Send PUBLISH packet yield from self._send_packet(publish_packet) app_message.publish_packet = publish_packet # Wait PUBREC if app_message.packet_id in self._pubrec_waiters: # PUBREC waiter already exists for this packet ID message = "Can't add PUBREC waiter, a waiter already exists for message Id '%s'" \ % app_message.packet_id self.logger.warning(message) raise HBMQTTException(message) waiter = asyncio.Future(loop=self._loop) self._pubrec_waiters[app_message.packet_id] = waiter yield from waiter del self._pubrec_waiters[app_message.packet_id] app_message.pubrec_packet = waiter.result() if not app_message.pubcomp_packet: # Send pubrel app_message.pubrel_packet = PubrelPacket.build(app_message.packet_id) yield from self._send_packet(app_message.pubrel_packet) # Wait for PUBCOMP waiter = asyncio.Future(loop=self._loop) self._pubcomp_waiters[app_message.packet_id] = waiter yield from waiter del self._pubcomp_waiters[app_message.packet_id] app_message.pubcomp_packet = waiter.result() # Discard inflight message del self.session.inflight_out[app_message.packet_id] elif app_message.direction == INCOMING: self.session.inflight_in[app_message.packet_id] = app_message # Send pubrec pubrec_packet = PubrecPacket.build(app_message.packet_id) yield from self._send_packet(pubrec_packet) app_message.pubrec_packet = pubrec_packet # Wait PUBREL if app_message.packet_id in self._pubrel_waiters and not self._pubrel_waiters[app_message.packet_id].done(): # PUBREL waiter already exists for this packet ID message = "Can't add PUBREL waiter, a waiter already exists for message Id '%s'" \ % app_message.packet_id self.logger.warning(message) raise HBMQTTException(message) waiter = asyncio.Future(loop=self._loop) self._pubrel_waiters[app_message.packet_id] = waiter yield from waiter del self._pubrel_waiters[app_message.packet_id] app_message.pubrel_packet = waiter.result() # Initiate delivery and discard message yield from self.session.delivered_message_queue.put(app_message) del self.session.inflight_in[app_message.packet_id] # Send pubcomp pubcomp_packet = PubcompPacket.build(app_message.packet_id) yield from self._send_packet(pubcomp_packet) app_message.pubcomp_packet = pubcomp_packet
def handle_publish(self, publish_packet: PublishPacket): incoming_message = None packet_id = publish_packet.variable_header.packet_id qos = publish_packet.qos if qos == 0: if publish_packet.dup_flag: self.logger.warn("[MQTT-3.3.1-2] DUP flag must set to 0 for QOS 0 message. Message ignored: %s" % repr(publish_packet)) else: # Assign packet_id as it's needed internally packet_id = self.session.next_packet_id publish_packet.variable_header.packet_id = packet_id incoming_message = IncomingInFlightMessage(publish_packet, qos, self.session.publish_retry_delay, self._loop) incoming_message.received_publish() self.session.incoming_msg[packet_id] = incoming_message yield from self.session.delivered_message_queue.put(packet_id) else: # Check if publish is a retry if packet_id in self.session.incoming_msg: incoming_message = self.session.incoming_msg[packet_id] else: incoming_message = IncomingInFlightMessage(publish_packet, qos, self.session.publish_retry_delay, self._loop) self.session.incoming_msg[packet_id] = incoming_message incoming_message.publish() if qos == 1: # Initiate delivery yield from self.session.delivered_message_queue.put(packet_id) ack = yield from incoming_message.wait_acknowledge() if ack: # Send PUBACK puback = PubackPacket.build(packet_id) yield from self.outgoing_queue.put(puback) #Discard message del self.session.incoming_msg[packet_id] self.logger.debug("Discarded incoming message %d" % packet_id) else: raise HBMQTTException("Something wrong, ack is False") if qos == 2: # Send PUBREC pubrec = PubrecPacket.build(packet_id) yield from self.outgoing_queue.put(pubrec) incoming_message.sent_pubrec() # Wait for pubrel ack = yield from incoming_message.wait_pubrel() if ack: # Initiate delivery yield from self.session.delivered_message_queue.put(packet_id) else: raise HBMQTTException("Something wrong, ack is False") ack = yield from incoming_message.wait_acknowledge() if ack: # Send PUBCOMP pubcomp = PubcompPacket.build(packet_id) yield from self.outgoing_queue.put(pubcomp) incoming_message.sent_pubcomp() #Discard message del self.session.incoming_msg[packet_id] self.logger.debug("Discarded incoming message %d" % packet_id) else: raise HBMQTTException("Something wrong, ack is False")
def handle_publish(self, publish_packet: PublishPacket): incoming_message = None packet_id = publish_packet.variable_header.packet_id qos = publish_packet.qos if qos == 0: if publish_packet.dup_flag: self.logger.warn( "[MQTT-3.3.1-2] DUP flag must set to 0 for QOS 0 message. Message ignored: %s" % repr(publish_packet)) else: # Assign packet_id as it's needed internally packet_id = self.session.next_packet_id publish_packet.variable_header.packet_id = packet_id incoming_message = IncomingInFlightMessage( publish_packet, qos, self.session.publish_retry_delay, self._loop) incoming_message.received_publish() self.session.incoming_msg[packet_id] = incoming_message yield from self.session.delivered_message_queue.put(packet_id) else: # Check if publish is a retry if packet_id in self.session.incoming_msg: incoming_message = self.session.incoming_msg[packet_id] else: incoming_message = IncomingInFlightMessage( publish_packet, qos, self.session.publish_retry_delay, self._loop) self.session.incoming_msg[packet_id] = incoming_message incoming_message.publish() if qos == 1: # Initiate delivery yield from self.session.delivered_message_queue.put(packet_id) ack = yield from incoming_message.wait_acknowledge() if ack: # Send PUBACK puback = PubackPacket.build(packet_id) yield from self.outgoing_queue.put(puback) #Discard message del self.session.incoming_msg[packet_id] self.logger.debug("Discarded incoming message %d" % packet_id) else: raise HBMQTTException("Something wrong, ack is False") if qos == 2: # Send PUBREC pubrec = PubrecPacket.build(packet_id) yield from self.outgoing_queue.put(pubrec) incoming_message.sent_pubrec() # Wait for pubrel ack = yield from incoming_message.wait_pubrel() if ack: # Initiate delivery yield from self.session.delivered_message_queue.put( packet_id) else: raise HBMQTTException("Something wrong, ack is False") ack = yield from incoming_message.wait_acknowledge() if ack: # Send PUBCOMP pubcomp = PubcompPacket.build(packet_id) yield from self.outgoing_queue.put(pubcomp) incoming_message.sent_pubcomp() #Discard message del self.session.incoming_msg[packet_id] self.logger.debug("Discarded incoming message %d" % packet_id) else: raise HBMQTTException("Something wrong, ack is False")