From b83784a088018b8580a9649dfae08ac026a91e82 Mon Sep 17 00:00:00 2001 From: Nicolas Jouanin Date: Wed, 15 Jul 2015 23:01:40 +0200 Subject: [PATCH] Fix will message encoding/decoding (bytes instead of string) --- hbmqtt/codecs.py | 16 ++++++++++++++++ hbmqtt/mqtt/connect.py | 4 ++-- tests/mqtt/test_connect.py | 4 ++-- 3 files changed, 20 insertions(+), 4 deletions(-) diff --git a/hbmqtt/codecs.py b/hbmqtt/codecs.py index 97838cb..33faca5 100644 --- a/hbmqtt/codecs.py +++ b/hbmqtt/codecs.py @@ -69,11 +69,27 @@ def decode_string(reader) -> bytes: byte_str = yield from read_or_raise(reader, str_length) return byte_str.decode(encoding='utf-8') +@asyncio.coroutine +def decode_data_with_length(reader) -> bytes: + """ + Read data from a reader. Data is prefixed with 2 bytes length + :param reader: Stream reader + :return: bytes read from stream (without length) + """ + length_bytes = yield from read_or_raise(reader, 2) + bytes_length = bytes_to_int(length_bytes) + data = yield from read_or_raise(reader, bytes_length) + return data + def encode_string(string: str) -> bytes: data = string.encode(encoding='utf-8') data_length = len(data) return int_to_bytes(data_length, 2) + data +def encode_data_with_length(data: bytes) -> bytes: + data_length = len(data) + return int_to_bytes(data_length, 2) + data + @asyncio.coroutine def decode_packet_id(reader) -> int: """ diff --git a/hbmqtt/mqtt/connect.py b/hbmqtt/mqtt/connect.py index 49902d0..c7ce7bd 100644 --- a/hbmqtt/mqtt/connect.py +++ b/hbmqtt/mqtt/connect.py @@ -156,7 +156,7 @@ class ConnectPayload(MQTTPayload): if variable_header.will_flag: try: payload.will_topic = yield from decode_string(reader) - payload.will_message = yield from decode_string(reader) + payload.will_message = yield from decode_data_with_length(reader) except NoDataException: payload.will_topic = None payload.will_message = None @@ -182,7 +182,7 @@ class ConnectPayload(MQTTPayload): # Will topic / message if variable_header.will_flag: out += encode_string(self.will_topic) - out += encode_string(self.will_message) + out += encode_data_with_length(self.will_message) # username if variable_header.username_flag: out += encode_string(self.username) diff --git a/tests/mqtt/test_connect.py b/tests/mqtt/test_connect.py index 9fa18cd..94c4fec 100644 --- a/tests/mqtt/test_connect.py +++ b/tests/mqtt/test_connect.py @@ -28,7 +28,7 @@ class ConnectPacketTest(unittest.TestCase): self.assertFalse(message.variable_header.reserved_flag) self.assertEqual(message.payload.client_id, '0123456789') self.assertEqual(message.payload.will_topic, 'WillTopic') - self.assertEqual(message.payload.will_message, 'WillMessage') + self.assertEqual(message.payload.will_message, b'WillMessage') self.assertEqual(message.payload.username, 'user') self.assertEqual(message.payload.password, 'password') @@ -101,7 +101,7 @@ class ConnectPacketTest(unittest.TestCase): def test_encode(self): header = MQTTFixedHeader(PacketType.CONNECT, 0x00, 0) variable_header = ConnectVariableHeader(0xce, 0, 'MQTT', 4) - payload = ConnectPayload('0123456789', 'WillTopic', 'WillMessage', 'user', 'password') + payload = ConnectPayload('0123456789', 'WillTopic', b'WillMessage', 'user', 'password') message = ConnectPacket(header, variable_header, payload) encoded = message.to_bytes() self.assertEqual(encoded, b'\x10\x3e\x00\x04MQTT\x04\xce\x00\x00\x00\x0a0123456789\x00\x09WillTopic\x00\x0bWillMessage\x00\x04user\x00\x08password')