def server_mock(reader, writer):
try:
packet = PublishPacket.build("/topic", b"test_data", rand_packet_id(), False, QOS_2, False)
yield from packet.to_stream(writer)
pubrec = yield from PubrecPacket.from_stream(reader)
self.assertIsNotNone(pubrec)
self.assertEqual(packet.packet_id, pubrec.packet_id)
self.assertIn(packet.packet_id, self.handler._pubrel_waiters)
pubrel = PubrelPacket.build(packet.packet_id)
yield from pubrel.to_stream(writer)
pubcomp = yield from PubcompPacket.from_stream(reader)
self.assertIsNotNone(pubcomp)
self.assertEqual(packet.packet_id, pubcomp.packet_id)
except Exception as ae:
future.set_exception(ae)
def server_mock(reader, writer):
try:
packet = PublishPacket.build('/topic', b'test_data', rand_packet_id(), False, QOS_2, False)
yield from packet.to_stream(writer)
pubrec = yield from PubrecPacket.from_stream(reader)
self.assertIsNotNone(pubrec)
self.assertEqual(packet.packet_id, pubrec.packet_id)
self.assertIn(packet.packet_id, self.handler._pubrel_waiters)
pubrel = PubrelPacket.build(packet.packet_id)
yield from pubrel.to_stream(writer)
pubcomp = yield from PubcompPacket.from_stream(reader)
self.assertIsNotNone(pubcomp)
self.assertEqual(packet.packet_id, pubcomp.packet_id)
except Exception as ae:
future.set_exception(ae)
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)
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.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 test_to_bytes(self):
variable_header = PacketIdVariableHeader(10)
publish = PubcompPacket(variable_header=variable_header)
out = publish.to_bytes()
self.assertEqual(out, b'\x70\x02\x00\x0a')
def test_from_stream(self):
data = b'\x70\x02\x00\x0a'
stream = BufferReader(data)
message = self.loop.run_until_complete(PubcompPacket.from_stream(stream))
self.assertEqual(message.variable_header.packet_id, 10)
def test_to_bytes(self):
variable_header = PacketIdVariableHeader(10)
publish = PubcompPacket(variable_header=variable_header)
out = publish.to_bytes()
self.assertEqual(out, b'\x70\x02\x00\x0a')
def test_from_stream(self):
data = b'\x70\x02\x00\x0a'
stream = BufferReader(data)
message = self.loop.run_until_complete(PubcompPacket.from_stream(stream))
self.assertEqual(message.variable_header.packet_id, 10)
def test_from_stream(self):
data = b'\x70\x02\x00\x0a'
stream = asyncio.StreamReader(loop=self.loop)
stream.feed_data(data)
message = self.loop.run_until_complete(PubcompPacket.from_stream(stream))
self.assertEqual(message.variable_header.packet_id, 10)
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")
