fix(b01): reassemble fragmented B01/Q7 live map responses

roborock/devices/rpc/b01_q7_channel.py

@@ -18,6 +18,7 @@
     ParamsType,
     Q7RequestMessage,
     create_map_key,
+    create_map_response_assembler,
     decode_map_payload,
     decode_rpc_response,
     encode_mqtt_payload,
@@ -58,15 +59,12 @@ async def subscribe_map_pushes(self, callback: Callable[[bytes], None]) -> Calla
         ...
 
 
-def _matches_map_response(response_message: RoborockMessage, *, version: bytes | None) -> bytes | None:
-    """Return raw map payload bytes for matching MAP_RESPONSE messages."""
-    if (
-        response_message.protocol == RoborockMessageProtocol.MAP_RESPONSE
-        and response_message.payload
-        and response_message.version == version
-    ):
-        return response_message.payload
-    return None
+def _is_map_message(response_message: RoborockMessage) -> bool:
+    """Check whether a message is a B01 map response part or final part."""
+    return response_message.protocol in (
+        RoborockMessageProtocol.MAP_RESPONSE_PART,
+        RoborockMessageProtocol.MAP_RESPONSE,
+    )
 
 
 async def _send_command(
@@ -172,6 +170,19 @@ class B01Q7Channel(Channel, Q7RpcChannel, Q7MapRpcChannel):
     def __init__(self, mqtt_channel: MqttChannel, map_key: MapKey) -> None:
         self._mqtt_channel = mqtt_channel
         self._map_key = map_key
+        self._map_assembler = create_map_response_assembler()
+
+    def _match_map_response(self, response_message: RoborockMessage) -> bytes | None:
+        """Return the complete raw map payload for matching map messages.
+
+        Large maps arrive split into several messages; the assembler buffers
+        the parts and returns the joined payload once the frame is complete.
+        """
+        if not _is_map_message(response_message):
+            return None
+        if not response_message.payload or response_message.version != B01_VERSION:
+            return None
+        return self._map_assembler(response_message)
 
     @property
     def is_connected(self) -> bool:
@@ -205,7 +216,7 @@ async def send_map_command(
             raw_payload = await _send_command(
                 self._mqtt_channel,
                 request_message,
-                response_matcher=lambda response_message: _matches_map_response(response_message, version=B01_VERSION),
+                response_matcher=self._match_map_response,
             )
         except TimeoutError as ex:
             raise RoborockException(f"B01 map command timed out after {_TIMEOUT}s ({request_message})") from ex
@@ -217,10 +228,11 @@ async def subscribe_map_pushes(self, callback: Callable[[bytes], None]) -> Calla
 
         The device streams full SCMap frames on its own during cleaning; the
         callback receives the decoded (inflated) SCMap bytes for each frame.
+        Frames split across several messages are reassembled first.
         """
 
         def on_message(message: RoborockMessage) -> None:
-            if (raw_payload := _matches_map_response(message, version=B01_VERSION)) is None:
+            if (raw_payload := self._match_map_response(message)) is None:
                 return
             try:
                 decoded = decode_map_payload(raw_payload, map_key=self._map_key)

roborock/protocols/b01_q7_protocol.py

@@ -6,6 +6,7 @@
 import json
 import logging
 import zlib
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -132,3 +133,70 @@ def _decode_base64_payload(raw_payload: bytes) -> bytes:
         return base64.b64decode(padded, validate=True)
     except binascii.Error as err:
         raise RoborockException("Failed to decode B01 map payload") from err
+
+
+# Number of incomplete multi-part map responses kept while waiting for missing parts.
+_MAX_PENDING_MAP_RESPONSES = 4
+
+
+def create_map_response_assembler() -> Callable[[RoborockMessage], bytes | None]:
+    """Create an assembler for B01 map response messages.
+
+    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). Parts may arrive out of order, so they are buffered
+    until complete. A lone final part whose decoded length is not a multiple of
+    the AES block size is a trailing fragment of a split frame, not a complete
+    frame, so it is buffered too. This ports the V1 reassembly from #1004.
+    """
+    # (random, timestamp) -> seq -> (is last part, raw payload)
+    pending: dict[tuple[int, int], dict[int, tuple[bool, bytes]]] = {}
+
+    def _decoded_length(raw_payload: bytes) -> int | None:
+        try:
+            return len(_decode_base64_payload(raw_payload))
+        except RoborockException:
+            return None
+
+    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 _complete_payload(parts: dict[int, tuple[bool, bytes]]) -> bytes | None:
+        seqs = sorted(parts)
+        if not parts[seqs[-1]][0]:
+            return None
+        if seqs != list(range(seqs[0], seqs[-1] + 1)):
+            return None
+        joined = b"".join(parts[seq][1] for seq in seqs)
+        length = _decoded_length(joined)
+        if length is None or length % AES.block_size != 0:
+            return None
+        return joined
+
+    def assemble(message: RoborockMessage) -> bytes | None:
+        """Return the complete raw payload once all parts of a frame arrived."""
+        if not message.payload:
+            return None
+        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):
+            parts = pending.pop(key, {})
+            parts[message.seq] = (is_last, message.payload)
+            if (joined := _complete_payload(parts)) is not None:
+                return joined
+            _remember(key, parts)
+            return None
+        if is_last:
+            length = _decoded_length(message.payload)
+            if length is not None and length % AES.block_size != 0:
+                # Trailing fragment of a split frame whose earlier parts have
+                # not arrived yet; wait for them instead of failing the decode.
+                _remember(key, {message.seq: (True, message.payload)})
+                return None
+            return message.payload
+        return None
+
+    return assemble

tests/devices/rpc/test_b01_q7_channel.py

@@ -1,4 +1,5 @@
 import asyncio
+import base64
 import json
 import math
 import time
@@ -62,6 +63,35 @@ def build_map_response(self, payload: bytes) -> RoborockMessage:
             seq=self.seq,
         )
 
+    def build_map_part(
+        self,
+        payload: bytes,
+        protocol: RoborockMessageProtocol,
+        seq: int,
+        random: int,
+        timestamp: int,
+    ) -> RoborockMessage:
+        """Build a MAP_RESPONSE_PART (300) or MAP_RESPONSE (301) message."""
+        return RoborockMessage(
+            protocol=protocol,
+            payload=payload,
+            version=b"B01",
+            seq=seq,
+            random=random,
+            timestamp=timestamp,
+        )
+
+
+def _fragmented_map_payload() -> tuple[bytes, bytes, bytes]:
+    """A base64 map payload split into two fragments, like issue #1043.
+
+    The full payload decodes to 48 bytes (a multiple of the AES block size);
+    each fragment alone decodes to 24 bytes, so neither fragment is decodable
+    on its own.
+    """
+    full = base64.b64encode(bytes(range(48)))
+    return full, full[:32], full[32:]
+
 
 @pytest.fixture(name="fake_channel")
 def fake_channel_fixture() -> FakeChannel:
@@ -350,3 +380,167 @@ async def test_subscribe_map_pushes_skips_undecodable_frames(
     ):
         fake_channel.notify_subscribers(message_builder.build_map_response(b"raw-map-payload"))
     assert received == [b"inflated-payload"]
+
+
+# ----------------------------------------------------------------------------- fragmented map responses (#1043)
+
+MAP_RANDOM = 2716216127
+MAP_TIMESTAMP = 1794578237
+
+
+async def test_send_map_command_reassembles_fragmented_frames(
+    device: HomeDataDevice,
+    product: HomeDataProduct,
+    fake_channel: FakeChannel,
+    message_builder: B01MessageBuilder,
+) -> None:
+    """A map split into a 300 part and a 301 part is decoded once, joined."""
+    channel = create_b01_q7_channel(device, product, fake_channel)  # type: ignore[arg-type]
+    full, first_part, last_part = _fragmented_map_payload()
+
+    with patch(
+        "roborock.devices.rpc.b01_q7_channel.decode_map_payload",
+        return_value=b"inflated-payload",
+    ) as mock_decode:
+        task = asyncio.create_task(channel.send_map_command("service.upload_by_mapid", {"map_id": 123}))
+        await asyncio.sleep(0)
+
+        fake_channel.notify_subscribers(
+            message_builder.build_map_part(
+                first_part, RoborockMessageProtocol.MAP_RESPONSE_PART, 1, MAP_RANDOM, MAP_TIMESTAMP
+            )
+        )
+        await asyncio.sleep(0)
+        assert not task.done()  # the first part alone is not a complete frame
+        fake_channel.notify_subscribers(
+            message_builder.build_map_part(
+                last_part, RoborockMessageProtocol.MAP_RESPONSE, 2, MAP_RANDOM, MAP_TIMESTAMP
+            )
+        )
+        assert await task == b"inflated-payload"
+        mock_decode.assert_called_once()
+        assert mock_decode.call_args[0][0] == full
+
+
+async def test_send_map_command_reassembles_out_of_order_fragments(
+    device: HomeDataDevice,
+    product: HomeDataProduct,
+    fake_channel: FakeChannel,
+    message_builder: B01MessageBuilder,
+) -> None:
+    """The trailing 301 part arriving first is buffered until the 300 part arrives."""
+    channel = create_b01_q7_channel(device, product, fake_channel)  # type: ignore[arg-type]
+    full, first_part, last_part = _fragmented_map_payload()
+
+    with patch(
+        "roborock.devices.rpc.b01_q7_channel.decode_map_payload",
+        return_value=b"inflated-payload",
+    ) as mock_decode:
+        task = asyncio.create_task(channel.send_map_command("service.upload_by_mapid", {"map_id": 123}))
+        await asyncio.sleep(0)
+
+        fake_channel.notify_subscribers(
+            message_builder.build_map_part(
+                last_part, RoborockMessageProtocol.MAP_RESPONSE, 2, MAP_RANDOM, MAP_TIMESTAMP
+            )
+        )
+        await asyncio.sleep(0)
+        assert not task.done()  # the trailing fragment alone is not decodable
+        fake_channel.notify_subscribers(
+            message_builder.build_map_part(
+                first_part, RoborockMessageProtocol.MAP_RESPONSE_PART, 1, MAP_RANDOM, MAP_TIMESTAMP
+            )
+        )
+        assert await task == b"inflated-payload"
+        mock_decode.assert_called_once()
+        assert mock_decode.call_args[0][0] == full
+
+
+async def test_subscribe_map_pushes_reassembles_fragmented_frames(
+    device: HomeDataDevice,
+    product: HomeDataProduct,
+    fake_channel: FakeChannel,
+    message_builder: B01MessageBuilder,
+) -> None:
+    """Fragmented map pushes are delivered once, as one decoded frame."""
+    channel = create_b01_q7_channel(device, product, fake_channel)  # type: ignore[arg-type]
+    full, first_part, last_part = _fragmented_map_payload()
+
+    received: list[bytes] = []
+    await channel.subscribe_map_pushes(received.append)
+
+    with patch(
+        "roborock.devices.rpc.b01_q7_channel.decode_map_payload",
+        return_value=b"inflated-payload",
+    ) as mock_decode:
+        fake_channel.notify_subscribers(
+            message_builder.build_map_part(
+                first_part, RoborockMessageProtocol.MAP_RESPONSE_PART, 1, MAP_RANDOM, MAP_TIMESTAMP
+            )
+        )
+        assert received == []
+        fake_channel.notify_subscribers(
+            message_builder.build_map_part(
+                last_part, RoborockMessageProtocol.MAP_RESPONSE, 2, MAP_RANDOM, MAP_TIMESTAMP
+            )
+        )
+        assert received == [b"inflated-payload"]
+        mock_decode.assert_called_once()
+        assert mock_decode.call_args[0][0] == full
+
+
+async def test_subscribe_map_pushes_buffers_lone_trailing_fragment(
+    device: HomeDataDevice,
+    product: HomeDataProduct,
+    fake_channel: FakeChannel,
+    message_builder: B01MessageBuilder,
+) -> None:
+    """A lone 301 fragment is buffered, not decoded, until its 300 part arrives."""
+    channel = create_b01_q7_channel(device, product, fake_channel)  # type: ignore[arg-type]
+    _full, first_part, last_part = _fragmented_map_payload()
+
+    received: list[bytes] = []
+    await channel.subscribe_map_pushes(received.append)
+
+    with patch(
+        "roborock.devices.rpc.b01_q7_channel.decode_map_payload",
+        return_value=b"inflated-payload",
+    ):
+        fake_channel.notify_subscribers(
+            message_builder.build_map_part(
+                last_part, RoborockMessageProtocol.MAP_RESPONSE, 2, MAP_RANDOM, MAP_TIMESTAMP
+            )
+        )
+        assert received == []
+        fake_channel.notify_subscribers(
+            message_builder.build_map_part(
+                first_part, RoborockMessageProtocol.MAP_RESPONSE_PART, 1, MAP_RANDOM, MAP_TIMESTAMP
+            )
+        )
+        assert received == [b"inflated-payload"]
+
+
+async def test_subscribe_map_pushes_ignores_incomplete_frame(
+    device: HomeDataDevice,
+    product: HomeDataProduct,
+    fake_channel: FakeChannel,
+    message_builder: B01MessageBuilder,
+) -> None:
+    """A lone 300 part with no final part delivers nothing."""
+    channel = create_b01_q7_channel(device, product, fake_channel)  # type: ignore[arg-type]
+    _full, first_part, _last_part = _fragmented_map_payload()
+
+    received: list[bytes] = []
+    await channel.subscribe_map_pushes(received.append)
+
+    with patch(
+        "roborock.devices.rpc.b01_q7_channel.decode_map_payload",
+        return_value=b"inflated-payload",
+    ) as mock_decode:
+        fake_channel.notify_subscribers(
+            message_builder.build_map_part(
+                first_part, RoborockMessageProtocol.MAP_RESPONSE_PART, 1, MAP_RANDOM, MAP_TIMESTAMP
+            )
+        )
+        assert received == []
+        mock_decode.assert_not_called()