diff --git a/roborock/protocols/v1_protocol.py b/roborock/protocols/v1_protocol.py index 733f48a1..7e5c857e 100644 --- a/roborock/protocols/v1_protocol.py +++ b/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 diff --git a/roborock/roborock_message.py b/roborock/roborock_message.py index 9fab9357..b044b8ac 100644 --- a/roborock/roborock_message.py +++ b/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 diff --git a/tests/protocols/test_v1_protocol.py b/tests/protocols/test_v1_protocol.py index c4bcb35c..21ee1de3 100644 --- a/tests/protocols/test_v1_protocol.py +++ b/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"), [