diff --git a/roborock/devices/rpc/b01_q7_channel.py b/roborock/devices/rpc/b01_q7_channel.py index a26611d22..d91f10f2d 100644 --- a/roborock/devices/rpc/b01_q7_channel.py +++ b/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) diff --git a/roborock/protocols/b01_q7_protocol.py b/roborock/protocols/b01_q7_protocol.py index 1eac3dcb7..e9fb32b07 100644 --- a/roborock/protocols/b01_q7_protocol.py +++ b/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 diff --git a/tests/devices/rpc/test_b01_q7_channel.py b/tests/devices/rpc/test_b01_q7_channel.py index 087291d9c..206c7947b 100644 --- a/tests/devices/rpc/test_b01_q7_channel.py +++ b/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()