fix: reassemble V1 map responses split into multiple messages

roborock/protocols/v1_protocol.py

@@ -306,16 +306,29 @@ def _decode_blob_response(message: RoborockMessage) -> BlobResponse | None:
     return _decode_blob_response
 
 
+# Number of incomplete multi-part map responses kept while waiting for missing parts.
+_MAX_PENDING_MAP_RESPONSES = 4
+
+
 def create_map_response_decoder(security_data: SecurityData) -> Callable[[RoborockMessage], MapResponse | None]:
-    """Create a decoder for V1 map response messages."""
+    """Create a decoder for V1 map response messages.
 
-    def _decode_map_response(message: RoborockMessage) -> MapResponse | None:
-        """Decode a V1 map response message."""
-        if not message.payload or len(message.payload) < 24:
+    Large maps are split into multiple messages that share the same ``random``
+    and ``timestamp`` and are numbered by ``seq``. All parts except the last
+    one use protocol ``MAP_RESPONSE_PART`` (300), the last one uses
+    ``MAP_RESPONSE`` (301). Only the first part starts with the map header.
+    Parts may arrive out of order, so they are buffered until complete.
+    """
+    endpoint_bytes = security_data.endpoint.encode()
+    # (random, timestamp) -> seq -> (is last part, payload)
+    pending: dict[tuple[int, int], dict[int, tuple[bool, bytes]]] = {}
+
+    def _decode_payload(payload: bytes) -> MapResponse | None:
+        if len(payload) < 24:
             raise RoborockException("Invalid V1 map response format: missing payload")
-        header, body = message.payload[:24], message.payload[24:]
+        header, body = payload[:24], payload[24:]
         [endpoint, _, request_id, _] = struct.unpack("<8s8sH6s", header)
-        if not endpoint.decode().startswith(security_data.endpoint):
+        if not endpoint.startswith(endpoint_bytes):
             _LOGGER.debug("Received map response not requested by this device, ignoring.")
             return None
         try:
@@ -325,6 +338,39 @@ def _decode_map_response(message: RoborockMessage) -> MapResponse | None:
         decompressed = Utils.decompress(decrypted)
         return MapResponse(request_id=request_id, data=decompressed)
 
+    def _remember(key: tuple[int, int], parts: dict[int, tuple[bool, bytes]]) -> None:
+        pending[key] = parts
+        while len(pending) > _MAX_PENDING_MAP_RESPONSES:
+            del pending[next(iter(pending))]
+
+    def _add_part(key: tuple[int, int], seq: int, is_last: bool, payload: bytes) -> MapResponse | None:
+        parts = pending.pop(key, {})
+        parts[seq] = (is_last, payload)
+        seqs = sorted(parts)
+        complete = (
+            parts[seqs[-1]][0]
+            and seqs == list(range(seqs[0], seqs[-1] + 1))
+            and parts[seqs[0]][1].startswith(endpoint_bytes)
+        )
+        if not complete:
+            _remember(key, parts)
+            return None
+        return _decode_payload(b"".join(parts[s][1] for s in seqs))
+
+    def _decode_map_response(message: RoborockMessage) -> MapResponse | None:
+        """Decode a V1 map response message."""
+        if not message.payload:
+            raise RoborockException("Invalid V1 map response format: missing payload")
+        payload = message.payload
+        key = (message.random, message.timestamp)
+        is_last = message.protocol == RoborockMessageProtocol.MAP_RESPONSE
+        if message.protocol == RoborockMessageProtocol.MAP_RESPONSE_PART or (is_last and key in pending):
+            return _add_part(key, message.seq, is_last, payload)
+        if is_last and not payload.startswith(endpoint_bytes):
+            # Possibly the last part of a multi-part response that arrived first.
+            _remember(key, {message.seq: (True, payload)})
+        return _decode_payload(payload)
+
     return _decode_map_response
 
 

roborock/roborock_message.py

@@ -15,6 +15,7 @@ class RoborockMessageProtocol(RoborockEnum):
     GENERAL_RESPONSE = 5
     RPC_REQUEST = 101
     RPC_RESPONSE = 102
+    MAP_RESPONSE_PART = 300  # Non-final part of a map response split into multiple messages
     MAP_RESPONSE = 301
 
 

tests/protocols/test_v1_protocol.py

@@ -1,8 +1,10 @@
 """Tests for the v1 protocol message encoding and decoding."""
 
+import gzip
 import json
 import logging
 import pathlib
+import random
 from collections.abc import Generator
 from unittest.mock import patch
 
@@ -229,6 +231,67 @@ def test_create_map_response_decoder_invalid_payload():
         decoder(message)
 
 
+def _multipart_map_messages(test_data: bytes, split_at: int) -> tuple[RoborockMessage, RoborockMessage]:
+    """Build a map response split into two messages like large maps sent over MQTT."""
+    header = TEST_ENDPOINT_BYTES + b"\x00" * 8 + b"\xdc\xad" + b"\x00" * 6
+    payload = header + Utils.encrypt_cbc(gzip.compress(test_data), SECURITY_DATA.nonce)
+    first = RoborockMessage(
+        protocol=RoborockMessageProtocol.MAP_RESPONSE_PART,
+        payload=payload[:split_at],
+        seq=1,
+        version=b"1.0",
+        random=1248,
+        timestamp=1792857010,
+    )
+    last = RoborockMessage(
+        protocol=RoborockMessageProtocol.MAP_RESPONSE,
+        payload=payload[split_at:],
+        seq=2,
+        version=b"1.0",
+        random=1248,
+        timestamp=1792857010,
+    )
+    return first, last
+
+
+@pytest.mark.parametrize("last_first", [False, True], ids=["in_order", "out_of_order"])
+def test_create_map_response_decoder_multipart(last_first: bool):
+    """Test that a map response split into multiple messages is reassembled."""
+    test_data = random.Random(0).randbytes(4000)  # incompressible, so both parts carry data
+    first, last = _multipart_map_messages(test_data, split_at=1000)
+    messages = [last, first] if last_first else [first, last]
+
+    decoder = create_map_response_decoder(SECURITY_DATA)
+    assert decoder(messages[0]) is None
+    result = decoder(messages[1])
+    assert result is not None
+    assert result.request_id == 44508
+    assert result.data == test_data
+
+
+def test_create_map_response_decoder_multipart_incomplete():
+    """Test that a lone first part of a multi-part map response is buffered, not decoded."""
+    first, _ = _multipart_map_messages(b"some map\n" * 100, split_at=100)
+
+    decoder = create_map_response_decoder(SECURITY_DATA)
+    assert decoder(first) is None
+
+
+def test_create_map_response_decoder_non_utf8_header():
+    """Test that a map response part without header does not raise on non UTF-8 bytes."""
+    message = RoborockMessage(
+        protocol=RoborockMessageProtocol.MAP_RESPONSE,
+        payload=b'\x9bv9\x15\x8e;"R' + b"\xff" * 40,
+        seq=2,
+        version=b"1.0",
+        random=1248,
+        timestamp=1792857010,
+    )
+
+    decoder = create_map_response_decoder(SECURITY_DATA)
+    assert decoder(message) is None
+
+
 @pytest.mark.parametrize(
     ("payload", "expected_data", "expected_error"),
     [