Change flags management

hbmqtt/message.py

@@ -26,13 +26,11 @@ def get_message_type(byte):
     return MessageType(byte)
 
 class FixedHeader:
-    def __init__(self, msg_type, length, dup_flag=False, qos=0, retain_flag=False):
+    def __init__(self, msg_type, flags, length):
         if isinstance(msg_type, int):
             enum_type = msg_type
         else:
             enum_type = get_message_type(msg_type)
         self.message_type = enum_type
         self.remainingLength = length
-        self.dup_flag = dup_flag
-        self.qos = qos
-        self.retain_flag = retain_flag
+        self.flags = flags

hbmqtt/streams/fixed_header.py

@@ -7,7 +7,10 @@ from hbmqtt.utils import (
     hex_to_int,
 )
 from hbmqtt.message import FixedHeader, MessageType
-from hbmqtt.streams.errors import FixedHeaderException
+from hbmqtt.streams.errors import StreamException, NoDataException
+
+class FixedHeaderException(StreamException):
+    pass
 
 class FixedHeaderStream:
     def __init__(self):
@@ -15,12 +18,15 @@ class FixedHeaderStream:
 
     def decode(self, reader) -> FixedHeader:
         b1 = yield from reader.read(1)
+        if not b1:
+            raise NoDataException
         msg_type = FixedHeaderStream.get_message_type(b1)
         if msg_type is MessageType.RESERVED_0 or msg_type is MessageType.RESERVED_15:
             raise FixedHeaderException("Usage of control packet type %s is forbidden" % msg_type)
-        (dup_flag, qos, retain_flag) = FixedHeaderStream.get_flags(b1)
+        flags = FixedHeaderStream.get_flags(b1)
+
         remain_length = yield from self.decode_remaining_length(reader)
-        return FixedHeader(msg_type, remain_length, dup_flag, qos, retain_flag)
+        return FixedHeader(msg_type, flags, remain_length)
 
     @staticmethod
     def get_message_type(byte):
@@ -30,10 +36,7 @@ class FixedHeaderStream:
     @staticmethod
     def get_flags(data):
         byte = hex_to_int(data)
-        b3 = True if (byte & 0x08) >> 3 else False
-        b21 = (byte & 0x06) >> 1
-        b0 = True if (byte & 0x01) else False
-        return b3, b21, b0
+        return byte & 0x0f
 
     @asyncio.coroutine
     def decode_remaining_length(self, reader):
@@ -42,6 +45,8 @@ class FixedHeaderStream:
         length_bytes = b''
         while True:
             encoded_byte = yield from reader.read(1)
+            if not encoded_byte:
+                raise NoDataException
             length_bytes += encoded_byte
             int_byte = hex_to_int(encoded_byte)
             value += (int_byte & 0x7f) * multiplier

tests/streams/test_fixed_header.py

@@ -3,8 +3,7 @@
 # See the file license.txt for copying permission.
 import unittest
 import asyncio
-from hbmqtt.streams.fixed_header import FixedHeaderStream
-from hbmqtt.streams.errors import FixedHeaderException
+from hbmqtt.streams.fixed_header import FixedHeaderStream, FixedHeaderException
 from hbmqtt.message import MessageType
 
 class TestFixedHeader(unittest.TestCase):
@@ -16,10 +15,12 @@ class TestFixedHeader(unittest.TestCase):
         self.assertEqual(m_type, MessageType.CONNECT)
 
     def test_get_flags(self):
-        (dup_flag, qos, retain_flag) = FixedHeaderStream.get_flags(b'\x1f')
-        self.assertTrue(dup_flag)
-        self.assertEqual(qos, 3)
-        self.assertTrue(retain_flag)
+        flags = FixedHeaderStream.get_flags(b'\x1f')
+        self.assertTrue(flags & 0x08)
+        self.assertTrue(flags & 0x04)
+        self.assertTrue(flags & 0x02)
+        self.assertTrue(flags & 0x01)
+        self.assertFalse(flags & 0x10)
 
     def test_decode_remaining_length1(self):
         stream = asyncio.StreamReader(loop=self.loop)
@@ -62,9 +63,9 @@ class TestFixedHeader(unittest.TestCase):
         s = FixedHeaderStream()
         header = self.loop.run_until_complete(s.decode(stream))
         self.assertEqual(header.message_type, MessageType.CONNECT)
-        self.assertFalse(header.dup_flag)
-        self.assertEqual(header.qos, 0)
-        self.assertFalse(header.retain_flag)
+        self.assertFalse(header.flags & 0x08)
+        self.assertEqual((header.flags & 0x06) >> 1 , 0)
+        self.assertFalse(header.flags & 0x01)
 
     def test_decode_ko(self):
         stream = asyncio.StreamReader(loop=self.loop)