diff --git a/pyomnilogic_local/api/protocol.py b/pyomnilogic_local/api/protocol.py index 69e9869..d92bf31 100644 --- a/pyomnilogic_local/api/protocol.py +++ b/pyomnilogic_local/api/protocol.py @@ -258,7 +258,9 @@ async def _reassemble_multipart(self, lead_msg: OmniLogicMessage) -> tuple[bytes received_block_count += 1 _LOGGER.debug("received block %d/%d", received_block_count, lead.msg_block_count) - return payload_data, compressed + # MsgSize identifies the response boundary. Do not strip null bytes from + # the payload: they may be valid bytes in a compressed stream's trailer. + return payload_data[: lead.msg_size], compressed def _parse_lead_message(self, msg: OmniLogicMessage) -> LeadMessage: """Parse the XML payload of an MSP_LEADMESSAGE into a LeadMessage model. @@ -284,7 +286,10 @@ def _decode_payload(self, data: bytes, compressed: bool) -> str: Decoded UTF-8 string with leading/trailing null bytes stripped. """ if compressed: - data = zlib.decompress(data.rstrip(b"\x00")) + # zlib ignores bytes after the end of a complete stream, so keeping + # the wire payload intact preserves a valid trailer that may end in + # a null byte while still tolerating protocol padding. + data = zlib.decompress(data) return data.decode("utf-8").strip("\x00") # ------------------------------------------------------------------------- diff --git a/tests/test_protocol.py b/tests/test_protocol.py index 1b97a26..93fb3ab 100644 --- a/tests/test_protocol.py +++ b/tests/test_protocol.py @@ -594,6 +594,46 @@ async def test_receive_file_decompresses_data() -> None: assert result == original.decode("utf-8") +@pytest.mark.asyncio +async def test_receive_file_decompresses_data_with_protocol_padding() -> None: + """Test that valid compressed trailers are not removed with protocol padding.""" + protocol = OmniLogicProtocol() + + original = b"" + compressed_data = b"" + for value in range(256): + candidate = b"" + b"a" * value + compressed_candidate = zlib.compress(candidate) + if compressed_candidate.endswith(b"\x00"): + original = candidate + compressed_data = compressed_candidate + break + else: + pytest.fail("Could not create a zlib stream ending in a null byte") + + leadmsg_payload = ( + 'LeadMessage' + '1003' + f'{len(compressed_data)}' + '10' + "" + ) + leadmsg = OmniLogicMessage(100, MessageType.MSP_LEADMESSAGE, payload=leadmsg_payload) + leadmsg.compressed = True + + block = OmniLogicMessage(101, MessageType.MSP_BLOCKMESSAGE) + block.payload = b"\x00" * 8 + compressed_data + b"\x00" * 3 + + await protocol._receive_queue.put(leadmsg) + await protocol._receive_queue.put(block) + + raw_data, compressed = await protocol._receive_response() + + assert raw_data == compressed_data + assert compressed is True + assert protocol._decode_payload(raw_data, compressed) == original.decode() + + @pytest.mark.asyncio async def test_receive_file_decompression_error() -> None: """Test that _decode_payload raises OmniMessageFormatError for invalid compressed data."""