Fix codec engine

hbmqtt/messages/connack.py

@@ -2,8 +2,9 @@
 #
 # See the file license.txt for copying permission.
 import asyncio
-from hbmqtt.messages.packet import MQTTPacket, MQTTFixedHeader, PacketType, MQTTVariableHeader
+from hbmqtt.messages.packet import MQTTPacket, MQTTFixedHeader, PacketType, MQTTVariableHeader, MQTTPayload
 from hbmqtt.codecs import int_to_bytes, read_or_raise
+from hbmqtt.errors import HBMQTTException
 from enum import Enum
 
 
@@ -41,10 +42,21 @@ class ConnackVariableHeader(MQTTVariableHeader):
 
         return out
 
+    def __repr__(self):
+        return 'ConnackVariableHeader(session_parent={0}, return_code={1})'.format(hex(self.session_parent), hex(self.return_code))
+
 
 class ConnackPacket(MQTTPacket):
-    def __init__(self, vh: ConnackVariableHeader):
-        header = MQTTFixedHeader(PacketType.CONNACK, 0x00)
+    VARIABLE_HEADER = ConnackVariableHeader
+    PAYLOAD = None
+
+    def __init__(self, fixed: MQTTFixedHeader, variable_header: ConnackVariableHeader=None, payload: MQTTPayload=None):
+        if fixed is None:
+            header = MQTTFixedHeader(PacketType.CONNACK, 0x00)
+        else:
+            if fixed.packet_type is not PacketType.CONNACK:
+                raise HBMQTTException("Invalid fixed packet type %s for ConnackPacket init" % fixed.packet_type)
+            header = fixed
         super().__init__(header)
-        self.variable_header = vh
+        self.variable_header = variable_header
         self.payload = None

hbmqtt/messages/connect.py

@@ -4,7 +4,7 @@
 import asyncio
 from hbmqtt.messages.packet import MQTTPacket, MQTTFixedHeader, PacketType, MQTTVariableHeader, MQTTPayload
 from hbmqtt.codecs import *
-from hbmqtt.errors import MQTTException, CodecException
+from hbmqtt.errors import MQTTException, CodecException, HBMQTTException
 from hbmqtt.session import Session
 
 
@@ -191,8 +191,16 @@ class ConnectPayload(MQTTPayload):
 
 
 class ConnectPacket(MQTTPacket):
-    def __init__(self, vh: ConnectVariableHeader, payload: ConnectPayload):
-        header = MQTTFixedHeader(PacketType.CONNECT, 0x00)
+    VARIABLE_HEADER = ConnectVariableHeader
+    PAYLOAD = ConnectPayload
+
+    def __init__(self, fixed: MQTTFixedHeader, vh: ConnectVariableHeader, payload: ConnectPayload):
+        if fixed is None:
+            header = MQTTFixedHeader(PacketType.CONNECT, 0x00)
+        else:
+            if fixed.packet_type is not PacketType.CONNECT:
+                raise HBMQTTException("Invalid fixed packet type %s for ConnectPacket init" % fixed.packet_type)
+            header = fixed
         super().__init__(header)
         self.variable_header = vh
         self.payload = payload
@@ -226,5 +234,6 @@ class ConnectPacket(MQTTPacket):
         else:
             vh.will_flag = False
 
-        packet = cls(vh, payload)
+        header = MQTTFixedHeader(PacketType.CONNECT, 0x00)
+        packet = cls(header, vh, payload)
         return packet

hbmqtt/messages/packet.py

@@ -119,6 +119,8 @@ class MQTTFixedHeader:
         remain_length = yield from decode_remaining_length()
         return cls(msg_type, flags, remain_length)
 
+    def __repr__(self):
+        return 'MQTTFixedHeader(type={0}, length={1}, flags={2})'.format(self.packet_type, self.remaining_length, hex(self.flags))
 
 class MQTTVariableHeader:
     def __init__(self):
@@ -162,6 +164,9 @@ class MQTTPayload:
 
 
 class MQTTPacket:
+    VARIABLE_HEADER = None
+    PAYLOAD = None
+
     def __init__(self, fixed: MQTTFixedHeader, variable_header: MQTTVariableHeader=None, payload: MQTTPayload=None):
         self.fixed_header = fixed
         self.variable_header = variable_header
@@ -178,7 +183,7 @@ class MQTTPacket:
         else:
             variable_header_bytes = b''
         if self.payload:
-            payload_bytes = self.payload.to_bytes()
+            payload_bytes = self.payload.to_bytes(self.fixed_header, self.variable_header)
         else:
             payload_bytes = b''
 
@@ -191,7 +196,16 @@ class MQTTPacket:
     @asyncio.coroutine
     def from_stream(cls, reader: asyncio.StreamReader):
         fixed_header = yield from MQTTFixedHeader.from_stream(reader)
-        variable_header = yield from MQTTVariableHeader.from_stream(reader, fixed_header)
-        payload = yield from MQTTPayload.from_stream(reader, fixed_header, variable_header)
+        if cls.VARIABLE_HEADER:
+            variable_header = yield from cls.VARIABLE_HEADER.from_stream(reader, fixed_header)
+        else:
+            variable_header = None
+        if cls.PAYLOAD:
+            payload = yield from cls.PAYLOAD.from_stream(reader, fixed_header, variable_header)
+        else:
+            payload= None
 
         return cls(fixed_header, variable_header, payload)
+
+    def __repr__(self):
+        return 'MQPacket(fixed={0!r}, variable={1!r}, payload={2!r})'.format(self.fixed_header, self.variable_header, self.payload)