Check firmware CRC before writing

Author: puddly

tests/test_tools_flash.py

@@ -1,9 +1,11 @@
+import pytest
+
 import random
 import zigpy_znp.types as t
 import zigpy_znp.commands as c
 
 from zigpy_znp.tools.flash_read import main as flash_read
-from zigpy_znp.tools.flash_write import main as flash_write
+from zigpy_znp.tools.flash_write import get_firmware_crcs, main as flash_write
 
 from test_api import pytest_mark_asyncio_timeout  # noqa: F401
 from test_application import znp_server  # noqa: F401
@@ -12,7 +14,12 @@
 
 random.seed(12345)
 FAKE_IMAGE_SIZE = 2 ** 10
-FAKE_FLASH = random.getrandbits(FAKE_IMAGE_SIZE * 8).to_bytes(FAKE_IMAGE_SIZE, "little")
+FAKE_FLASH = bytearray(
+    random.getrandbits(FAKE_IMAGE_SIZE * 8).to_bytes(FAKE_IMAGE_SIZE, "little")
+)
+FAKE_FLASH[c.ubl.IMAGE_CRC_OFFSET : c.ubl.IMAGE_CRC_OFFSET + 2] = get_firmware_crcs(
+    FAKE_FLASH
+)[1].to_bytes(2, "little")
 random.seed()
 
 
@@ -89,3 +96,26 @@ def write_flash(req):
 
     # They should be identical
     assert backup_file.read_bytes() == FAKE_FLASH
+
+
+@pytest_mark_asyncio_timeout(seconds=5)
+async def test_flash_write_bad_crc(
+    openable_serial_znp_server, tmp_path, mocker  # noqa: F811
+):
+    # It takes too long otherwise
+    mocker.patch("zigpy_znp.commands.ubl.IMAGE_SIZE", FAKE_IMAGE_SIZE)
+
+    # Flip the bits in one byte, the CRC should fail
+    BAD_FIRMWARE = bytearray(len(FAKE_FLASH))
+    BAD_FIRMWARE[FAKE_IMAGE_SIZE - 1] = BAD_FIRMWARE[FAKE_IMAGE_SIZE - 1] ^ 0xFF
+
+    # No communication will happen because the CRC will be invalid
+    firmware_file = tmp_path / "bad-firmware.bin"
+    firmware_file.write_bytes(BAD_FIRMWARE)
+
+    with pytest.raises(ValueError) as e:
+        await flash_write(
+            [openable_serial_znp_server._port_path, "-i", str(firmware_file)]
+        )
+
+    assert "Firmware CRC is incorrect" in str(e)

zigpy_znp/tools/flash_write.py

@@ -1,4 +1,5 @@
 import sys
+import typing
 import asyncio
 import logging
 import argparse
@@ -17,13 +18,59 @@
 LOGGER = logging.getLogger(__name__)
 
 
+def compute_crc16(data: bytes) -> int:
+    poly = 0x1021
+    crc = 0x0000
+
+    for byte in data:
+        for _ in range(8):
+            msb = 1 if (crc & 0x8000) else 0
+
+            crc <<= 1
+            crc &= 0xFFFF
+
+            if byte & 0x80:
+                crc |= 0x0001
+
+            if msb:
+                crc ^= poly
+
+            byte <<= 1
+
+    return crc
+
+
+def get_firmware_crcs(firmware: bytes) -> typing.Tuple[int, int]:
+    # There is room for *two* CRCs in the firmware file: the expected and the computed
+    firmware_wihout_crcs = (
+        firmware[: c.ubl.IMAGE_CRC_OFFSET]
+        + firmware[c.ubl.IMAGE_CRC_OFFSET + 4 :]
+        + b"\x00\x00"
+    )
+
+    # We only use the first one. The second one is written by the bootloader into flash.
+    real_crc = int.from_bytes(
+        firmware[c.ubl.IMAGE_CRC_OFFSET : c.ubl.IMAGE_CRC_OFFSET + 2], "little"
+    )
+
+    return real_crc, compute_crc16(firmware_wihout_crcs)
+
+
 async def write_firmware(firmware: bytes, radio_path: str):
     if len(firmware) != c.ubl.IMAGE_SIZE:
         raise ValueError(
             f"Firmware is the wrong size."
             f" Expected {c.ubl.IMAGE_SIZE}, got {len(firmware)}"
         )
 
+    expected_crc, computed_crc = get_firmware_crcs(firmware)
+
+    if expected_crc != computed_crc:
+        raise ValueError(
+            f"Firmware CRC is incorrect. "
+            f"Expected 0x{expected_crc:04X}, got 0x{computed_crc:04X}"
+        )
+
     znp = ZNP(CONFIG_SCHEMA({"device": {"path": radio_path}}))
 
     # The bootloader handshake must be the very first command
@@ -63,7 +110,6 @@ async def write_firmware(firmware: bytes, radio_path: str):
         assert write_rsp.Status == c.ubl.BootloaderStatus.SUCCESS
 
     # Now we have to read it all back
-    # TODO: figure out how the CRC is computed!
     for offset in range(0, c.ubl.IMAGE_SIZE, buffer_size):
         address = offset // c.ubl.FLASH_WORD_SIZE
         LOGGER.info(
@@ -79,7 +125,7 @@ async def write_firmware(firmware: bytes, radio_path: str):
         assert read_rsp.FlashWordAddr == address
         assert read_rsp.Data == firmware[offset : offset + buffer_size]
 
-    # This seems to cause the firmware to compute and verify the CRC
+    # This seems to cause the bootloader to compute and verify the CRC
     enable_rsp = await znp.request_callback_rsp(
         request=c.UBL.EnableReq.Req(), callback=c.UBL.EnableRsp.Callback(partial=True),
     )