add test case for long list of topics

this uncovered a bug in SUBACK processing

Author: vladak

adafruit_minimqtt/adafruit_minimqtt.py

@@ -625,7 +625,7 @@ def _connect(
             var_header[7] |= 0x4 | (self._lw_qos & 0x1) << 3 | (self._lw_qos & 0x2) << 3
             var_header[7] |= self._lw_retain << 5
 
-        self.encode_remaining_length(fixed_header, remaining_length)
+        self._encode_remaining_length(fixed_header, remaining_length)
         self.logger.debug("Sending CONNECT to broker...")
         self.logger.debug(f"Fixed Header: {fixed_header}")
         self.logger.debug(f"Variable Header: {var_header}")
@@ -663,10 +663,13 @@ def _connect(
                     )
 
     # pylint: disable=no-self-use
-    def encode_remaining_length(self, fixed_header: bytearray, remaining_length: int):
-        """
-        Encode Remaining Length [2.2.3]
-        """
+    def _encode_remaining_length(
+        self, fixed_header: bytearray, remaining_length: int
+    ) -> None:
+        """Encode Remaining Length [2.2.3]"""
+        if remaining_length > 268_435_455:
+            raise MMQTTException("invalid remaining length")
+
         # Remaining length calculation
         if remaining_length > 0x7F:
             while remaining_length > 0:
@@ -765,7 +768,7 @@ def publish(
             pub_hdr_var.append(self._pid >> 8)
             pub_hdr_var.append(self._pid & 0xFF)
 
-        self.encode_remaining_length(pub_hdr_fixed, remaining_length)
+        self._encode_remaining_length(pub_hdr_fixed, remaining_length)
 
         self.logger.debug(
             "Sending PUBLISH\nTopic: %s\nMsg: %s\
@@ -836,7 +839,7 @@ def subscribe(self, topic: Optional[Union[tuple, str, list]], qos: int = 0) -> N
         fixed_header = bytearray([MQTT_SUB])
         packet_length = 2 + (2 * len(topics)) + (1 * len(topics))
         packet_length += sum(len(topic.encode("utf-8")) for topic, qos in topics)
-        self.encode_remaining_length(fixed_header, remaining_length=packet_length)
+        self._encode_remaining_length(fixed_header, remaining_length=packet_length)
         self.logger.debug(f"Fixed Header: {fixed_header}")
         self._sock.send(fixed_header)
         self._pid = self._pid + 1 if self._pid < 0xFFFF else 1
@@ -864,13 +867,13 @@ def subscribe(self, topic: Optional[Union[tuple, str, list]], qos: int = 0) -> N
                     )
             else:
                 if op == 0x90:
-                    rc = self._sock_exact_recv(3)
-                    # Check packet identifier.
-                    assert rc[1] == var_header[0] and rc[2] == var_header[1]
-                    remaining_len = rc[0] - 2
+                    remaining_len = self._decode_remaining_length()
                     assert remaining_len > 0
-                    rc = self._sock_exact_recv(remaining_len)
-                    for i in range(0, remaining_len):
+                    rc = self._sock_exact_recv(2)
+                    # Check packet identifier.
+                    assert rc[0] == var_header[0] and rc[1] == var_header[1]
+                    rc = self._sock_exact_recv(remaining_len - 2)
+                    for i in range(0, remaining_len - 2):
                         if rc[i] not in [0, 1, 2]:
                             raise MMQTTException(
                                 f"SUBACK Failure for topic {topics[i][0]}: {hex(rc[i])}"
@@ -915,7 +918,7 @@ def unsubscribe(self, topic: Optional[Union[str, list]]) -> None:
         fixed_header = bytearray([MQTT_UNSUB])
         packet_length = 2 + (2 * len(topics))
         packet_length += sum(len(topic.encode("utf-8")) for topic in topics)
-        self.encode_remaining_length(fixed_header, remaining_length=packet_length)
+        self._encode_remaining_length(fixed_header, remaining_length=packet_length)
         self.logger.debug(f"Fixed Header: {fixed_header}")
         self._sock.send(fixed_header)
         self._pid = self._pid + 1 if self._pid < 0xFFFF else 1
@@ -1090,7 +1093,7 @@ def _wait_for_msg(self) -> Optional[int]:
             return pkt_type
 
         # Handle only the PUBLISH packet type from now on.
-        sz = self._recv_len()
+        sz = self._decode_remaining_length()
         # topic length MSB & LSB
         topic_len_buf = self._sock_exact_recv(2)
         topic_len = int((topic_len_buf[0] << 8) | topic_len_buf[1])
@@ -1123,11 +1126,13 @@ def _wait_for_msg(self) -> Optional[int]:
 
         return pkt_type
 
-    def _recv_len(self) -> int:
-        """Unpack MQTT message length."""
+    def _decode_remaining_length(self) -> int:
+        """Decode Remaining Length [2.2.3]"""
         n = 0
         sh = 0
         while True:
+            if sh > 28:
+                raise MMQTTException("invalid remaining length encoding")
             b = self._sock_exact_recv(1)[0]
             n |= (b & 0x7F) << sh
             if not b & 0x80:

tests/test_subscribe.py

@@ -49,6 +49,29 @@ def handle_subscribe(client, user_data, topic, qos):
             ]
         ),
     ),
+    # same as before but with tuple
+    (
+        ("foo/bar", 0),
+        bytearray([0x90, 0x03, 0x00, 0x01, 0x00]),  # SUBACK
+        bytearray(
+            [
+                0x82,  # fixed header
+                0x0C,  # remaining length
+                0x00,
+                0x01,  # message ID
+                0x00,
+                0x07,  # topic length
+                0x66,  # topic
+                0x6F,
+                0x6F,
+                0x2F,
+                0x62,
+                0x61,
+                0x72,
+                0x00,  # QoS
+            ]
+        ),
+    ),
     # remaining length is encoded as 2 bytes due to long topic name.
     (
         "f" + "o" * 257,
@@ -113,13 +136,52 @@ def handle_subscribe(client, user_data, topic, qos):
             ]
         ),
     ),
+    # use list of topics for more coverage. If the range was (1, 10000), that would be
+    # long enough to use 3 bytes for remaining length, however that would make the test
+    # run for many minutes even on modern systems, so 1001 is used instead.
+    # This results in 2 bytes for the remaining length.
+    (
+        [(f"foo/bar{x:04}", 0) for x in range(1, 1001)],
+        bytearray(
+            [
+                0x90,
+                0xEA,  # remaining length
+                0x07,
+                0x00,  # message ID
+                0x01,
+            ]
+            + [0x00] * 1000  # success for all topics
+        ),
+        bytearray(
+            [
+                0x82,  # fixed header
+                0xB2,  # remaining length
+                0x6D,
+                0x00,  # message ID
+                0x01,
+            ]
+            + sum(
+                [
+                    [0x00, 0x0B] + list(f"foo/bar{x:04}".encode("ascii")) + [0x00]
+                    for x in range(1, 1001)
+                ],
+                [],
+            )
+        ),
+    ),
 ]
 
 
 @pytest.mark.parametrize(
     "topic,to_send,exp_recv",
     testdata,
-    ids=["short_topic", "long_topic", "publish_first"],
+    ids=[
+        "short_topic",
+        "short_topic_tuple",
+        "long_topic",
+        "publish_first",
+        "topic_list_long",
+    ],
 )
 def test_subscribe(topic, to_send, exp_recv) -> None:
     """
@@ -157,5 +219,10 @@ def test_subscribe(topic, to_send, exp_recv) -> None:
     logger.info(f"subscribing to {topic}")
     mqtt_client.subscribe(topic)
 
-    assert topic in subscribed_topics
+    if isinstance(topic, str):
+        assert topic in subscribed_topics
+    elif isinstance(topic, list):
+        for topic_name, _ in topic:
+            assert topic_name in subscribed_topics
     assert mocket.sent == exp_recv
+    assert len(mocket._to_send) == 0