Fix will message encoding/decoding (bytes instead of string)

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:
     """

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)

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')