umqtt.simple: Encode subscribe Remaining Length as VBI.

micropython/umqtt.simple/test_umqtt_simple.py

@@ -0,0 +1,76 @@
+import io
+import sys
+
+
+class Socket:
+    def __init__(self, read_data=b""):
+        self._write_buffer = io.BytesIO()
+        self._read_buffer = io.BytesIO(read_data)
+
+    def write(self, buf, length=None):
+        if length is None:
+            length = len(buf)
+        self._write_buffer.write(buf[:length])
+
+    def read(self, n):
+        return self._read_buffer.read(n)
+
+    def setblocking(self, blocking):
+        pass
+
+    def close(self):
+        pass
+
+
+sys.path.insert(0, "micropython/umqtt.simple")
+# ruff: noqa: E402
+from umqtt.simple import MQTTClient
+
+
+def make_client(read_data):
+    c = MQTTClient(b"cid", "127.0.0.1")
+    c.sock = Socket(read_data)
+    c.set_callback(lambda topic, msg: None)
+    return c
+
+
+def test_subscribe_short_topic():
+    # Remaining Length = 5 + 4 = 9 -> single VBI byte 0x09, pid=1
+    c = make_client(b"\x90\x03\x00\x01\x00")
+    c.subscribe(b"abcd", qos=0)
+    out = c.sock._write_buffer.getvalue()
+    assert out[:2] == b"\x82\x09", out
+    assert out[2:4] == b"\x00\x01", out
+    assert out[4:6] == b"\x00\x04", out
+    assert out[6:10] == b"abcd", out
+    assert out[10:11] == b"\x00", out
+
+
+def test_subscribe_long_topic():
+    # Remaining Length = 5 + 123 = 128 -> VBI 0x80 0x01
+    topic = b"a" * 123
+    c = make_client(b"\x90\x03\x00\x01\x00")
+    c.subscribe(topic, qos=0)
+    out = c.sock._write_buffer.getvalue()
+    assert out[:3] == b"\x82\x80\x01", out
+    assert out[3:5] == b"\x00\x01", out
+    assert out[5:7] == b"\x00\x7b", out
+    assert out[7:130] == topic, out
+    assert out[130:131] == b"\x00", out
+
+
+def test_unsubscribe_long_topic():
+    # Remaining Length = 4 + 123 = 127 -> still one VBI byte; use 124 for 128
+    topic = b"a" * 124
+    c = make_client(b"\xb0\x02\x00\x01")
+    c.unsubscribe(topic)
+    out = c.sock._write_buffer.getvalue()
+    assert out[:3] == b"\xa2\x80\x01", out
+    assert out[3:5] == b"\x00\x01", out
+    assert out[5:7] == b"\x00\x7c", out
+    assert out[7:131] == topic, out
+
+
+test_subscribe_short_topic()
+test_subscribe_long_topic()
+test_unsubscribe_long_topic()

micropython/umqtt.simple/umqtt/simple.py

@@ -7,6 +7,16 @@ class MQTTException(Exception):
     pass
 
 
+def _encode_len(pkt, sz):
+    i = 1
+    while sz > 0x7F:
+        pkt[i] = (sz & 0x7F) | 0x80
+        sz >>= 7
+        i += 1
+    pkt[i] = sz
+    return i
+
+
 class MQTTClient:
     def __init__(
         self,
@@ -93,12 +103,7 @@ def connect(self, clean_session=True, timeout=None):
             msg[6] |= 0x4 | (self.lw_qos & 0x1) << 3 | (self.lw_qos & 0x2) << 3
             msg[6] |= self.lw_retain << 5
 
-        i = 1
-        while sz > 0x7F:
-            premsg[i] = (sz & 0x7F) | 0x80
-            sz >>= 7
-            i += 1
-        premsg[i] = sz
+        i = _encode_len(premsg, sz)
 
         self.sock.write(premsg, i + 2)
         self.sock.write(msg)
@@ -130,12 +135,7 @@ def publish(self, topic, msg, retain=False, qos=0):
         if qos > 0:
             sz += 2
         assert sz < 2097152
-        i = 1
-        while sz > 0x7F:
-            pkt[i] = (sz & 0x7F) | 0x80
-            sz >>= 7
-            i += 1
-        pkt[i] = sz
+        i = _encode_len(pkt, sz)
         # print(hex(len(pkt)), hexlify(pkt, ":"))
         self.sock.write(pkt, i + 1)
         self._send_str(topic)
@@ -158,37 +158,33 @@ def publish(self, topic, msg, retain=False, qos=0):
         elif qos == 2:
             assert 0
 
-    def subscribe(self, topic, qos=0):
-        assert self.cb is not None, "Subscribe callback is not set"
-        pkt = bytearray(b"\x82\0\0\0")
+    def _send_subunsub(self, topic, typ, ack_op, ack_n, qos=0):
+        pkt = bytearray(4)
+        pkt[0] = typ
         self.pid += 1
-        struct.pack_into("!BH", pkt, 1, 2 + 2 + len(topic) + 1, self.pid)
-        # print(hex(len(pkt)), hexlify(pkt, ":"))
-        self.sock.write(pkt)
+        pid = self.pid
+        i = _encode_len(pkt, 4 + len(topic) + (ack_n > 3))
+        self.sock.write(pkt, i + 1)
+        struct.pack_into("!H", pkt, 0, pid)
+        self.sock.write(pkt, 2)
         self._send_str(topic)
-        self.sock.write(qos.to_bytes(1, "little"))
+        if ack_n > 3:
+            self.sock.write(bytes((qos,)))
         while 1:
             op = self.wait_msg()
-            if op == 0x90:
-                resp = self.sock.read(4)
-                # print(resp)
-                assert resp[1] == pkt[2] and resp[2] == pkt[3]
-                if resp[3] == 0x80:
+            if op == ack_op:
+                resp = self.sock.read(ack_n)
+                assert (resp[1] << 8 | resp[2]) == pid
+                if ack_n > 3 and resp[3] == 0x80:
                     raise MQTTException(resp[3])
                 return
 
+    def subscribe(self, topic, qos=0):
+        assert self.cb is not None, "Subscribe callback is not set"
+        self._send_subunsub(topic, 0x82, 0x90, 4, qos)
+
     def unsubscribe(self, topic):
-        pkt = bytearray(b"\xa2\0\0\0")
-        self.pid += 1
-        struct.pack_into("!BH", pkt, 1, 2 + 2 + len(topic), self.pid)
-        self.sock.write(pkt)
-        self._send_str(topic)
-        while 1:
-            op = self.wait_msg()
-            if op == 0xB0:
-                resp = self.sock.read(3)
-                assert resp[1] == pkt[2] and resp[2] == pkt[3]
-                return
+        self._send_subunsub(topic, 0xA2, 0xB0, 3)
 
     # Wait for a single incoming MQTT message and process it.
     # Subscribed messages are delivered to a callback previously