Merge pull request #92 from puddly/puddly/nvram-migrations

Fix broken empty address manager entries after restore on affected devices

Author: puddly

tests/application/test_nvram_migration.py

@@ -0,0 +1,114 @@
+import pytest
+
+import zigpy_znp.const as const
+import zigpy_znp.types as t
+import zigpy_znp.commands as c
+from zigpy_znp.types.nvids import ExNvIds, OsalNvIds
+
+from ..conftest import FORMED_DEVICES, FormedZStack3CC2531
+
+pytestmark = [pytest.mark.asyncio]
+
+
+@pytest.mark.parametrize("device", FORMED_DEVICES)
+async def test_addrmgr_empty_entries(make_connected_znp, device):
+    znp, znp_server = await make_connected_znp(server_cls=device)
+
+    if znp.version >= 3.30:
+        entries = [
+            entry
+            async for entry in znp.nvram.read_table(
+                item_id=ExNvIds.ADDRMGR,
+                item_type=t.AddrMgrEntry,
+            )
+        ]
+    else:
+        entries = list(
+            await znp.nvram.osal_read(
+                OsalNvIds.ADDRMGR, item_type=t.AddressManagerTable
+            )
+        )
+
+    num_empty = 0
+
+    for entry in entries:
+        if entry.extAddr != t.EUI64.convert("FF:FF:FF:FF:FF:FF:FF:FF"):
+            continue
+
+        num_empty += 1
+
+        if znp.version >= 3.30:
+            assert entry == const.EMPTY_ADDR_MGR_ENTRY_ZSTACK3
+        else:
+            assert entry == const.EMPTY_ADDR_MGR_ENTRY_ZSTACK1
+
+    assert num_empty > 0
+
+
+@pytest.mark.parametrize("device", [FormedZStack3CC2531])
+async def test_addrmgr_rewrite_fix(device, make_application, mocker):
+    # Keep track of reads
+    addrmgr_reads = []
+
+    correct_entry = t.AddrMgrEntry(
+        type=t.AddrMgrUserType.Default,
+        nwkAddr=0xFFFF,
+        extAddr=t.EUI64.convert("FF:FF:FF:FF:FF:FF:FF:FF"),
+    )
+
+    bad_entry = t.AddrMgrEntry(
+        type=t.AddrMgrUserType(0xFF),
+        nwkAddr=0xFFFF,
+        extAddr=t.EUI64.convert("FF:FF:FF:FF:FF:FF:FF:FF"),
+    )
+
+    app, znp_server = make_application(server_cls=device)
+    znp_server.callback_for_response(
+        c.SYS.OSALNVReadExt.Req(Id=OsalNvIds.ADDRMGR, Offset=0), addrmgr_reads.append
+    )
+
+    nvram = znp_server._nvram[ExNvIds.LEGACY]
+    old_addrmgr, _ = t.AddressManagerTable.deserialize(nvram[OsalNvIds.ADDRMGR])
+
+    # Ensure the table looks the way we expect
+    assert old_addrmgr.count(correct_entry) == 58
+    assert old_addrmgr.count(bad_entry) == 0
+
+    assert nvram[OsalNvIds.ADDRMGR] == b"".join([e.serialize() for e in old_addrmgr])
+
+    # Purposefully corrupt the empty entries
+    nvram[OsalNvIds.ADDRMGR] = b"".join(
+        [(bad_entry if e == correct_entry else e).serialize() for e in old_addrmgr]
+    )
+    assert old_addrmgr != nvram[OsalNvIds.ADDRMGR]
+
+    assert len(addrmgr_reads) == 0
+    await app.startup()
+    await app.shutdown()
+    assert len(addrmgr_reads) == 2
+
+    # Bad entries have been fixed
+    new_addrmgr, _ = t.AddressManagerTable.deserialize(nvram[OsalNvIds.ADDRMGR])
+    assert new_addrmgr == old_addrmgr
+
+    # Migration has been created
+    assert t.uint8_t.deserialize(nvram[OsalNvIds.ZIGPY_ZNP_MIGRATION_ID])[0] >= 1
+
+    # Will not be read again
+    assert len(addrmgr_reads) == 2
+    await app.startup()
+    await app.shutdown()
+    assert len(addrmgr_reads) == 2
+
+    # Will be migrated again if the migration NVID is deleted
+    del nvram[OsalNvIds.ZIGPY_ZNP_MIGRATION_ID]
+
+    old_addrmgr2 = nvram[OsalNvIds.ADDRMGR]
+
+    assert len(addrmgr_reads) == 2
+    await app.startup()
+    await app.shutdown()
+    assert len(addrmgr_reads) == 3
+
+    # But nothing will change
+    assert nvram[OsalNvIds.ADDRMGR] == old_addrmgr2

tests/conftest.py

@@ -978,7 +978,7 @@ def _create_network_nvram(self):
         super()._create_network_nvram()
         self._nvram[ExNvIds.LEGACY][OsalNvIds.APS_LINK_KEY_TABLE] = b"\xFF" * 20
         self._nvram[ExNvIds.ADDRMGR] = {
-            addr: self.nvram_serialize(const.EMPTY_ADDR_MGR_ENTRY)
+            addr: self.nvram_serialize(const.EMPTY_ADDR_MGR_ENTRY_ZSTACK3)
             for addr in range(0x0000, 0x0100 + 1)
         }
 
@@ -1094,7 +1094,7 @@ def _create_network_nvram(self):
         super()._create_network_nvram()
         self._nvram[ExNvIds.LEGACY][OsalNvIds.APS_LINK_KEY_TABLE] = b"\xFF" * 17
         self._nvram[ExNvIds.LEGACY][OsalNvIds.ADDRMGR] = 124 * self.nvram_serialize(
-            const.EMPTY_ADDR_MGR_ENTRY
+            const.EMPTY_ADDR_MGR_ENTRY_ZSTACK1
         )
 
     def create_nib(self, _=None):

zigpy_znp/api.py

@@ -191,9 +191,10 @@ async def write_network_info(
 
         from zigpy_znp.znp import security
 
-        # Delete any existing HAS_CONFIGURED_ZSTACK* NV items. One (or both) may fail.
+        # Delete any existing NV items that store formation state
         await self.nvram.osal_delete(OsalNvIds.HAS_CONFIGURED_ZSTACK1)
         await self.nvram.osal_delete(OsalNvIds.HAS_CONFIGURED_ZSTACK3)
+        await self.nvram.osal_delete(OsalNvIds.ZIGPY_ZNP_MIGRATION_ID)
         await self.nvram.osal_delete(OsalNvIds.BDBNODEISONANETWORK)
 
         # Instruct Z-Stack to reset everything on the next boot

zigpy_znp/const.py

@@ -12,7 +12,13 @@
 )
 ZSTACK_CONFIGURE_SUCCESS = t.uint8_t(0x55)
 
-EMPTY_ADDR_MGR_ENTRY = t.AddrMgrEntry(
+EMPTY_ADDR_MGR_ENTRY_ZSTACK1 = t.AddrMgrEntry(
+    type=t.AddrMgrUserType.Default,
+    nwkAddr=0xFFFF,
+    extAddr=t.EUI64.convert("FF:FF:FF:FF:FF:FF:FF:FF"),
+)
+
+EMPTY_ADDR_MGR_ENTRY_ZSTACK3 = t.AddrMgrEntry(
     type=t.AddrMgrUserType(0xFF),
     nwkAddr=0xFFFF,
     extAddr=t.EUI64.convert("FF:FF:FF:FF:FF:FF:FF:FF"),

zigpy_znp/types/nvids.py

@@ -40,6 +40,10 @@ class OsalNvIds(BaseNvIds):
     HAS_CONFIGURED_ZSTACK1 = 0x0F00
     HAS_CONFIGURED_ZSTACK3 = 0x0060
 
+    # Although the docs say "IDs reserved for applications range from 0x0401 to 0x0FFF",
+    # no OSAL NVID beyond 0x03FF is writable with the MT interface when using Z-Stack 3.
+    ZIGPY_ZNP_MIGRATION_ID = 0x005F
+
     # OSAL NV item IDs
     EXTADDR = 0x0001
     BOOTCOUNTER = 0x0002

zigpy_znp/zigbee/application.py

@@ -30,6 +30,7 @@
 import zigpy_znp.config as conf
 import zigpy_znp.commands as c
 from zigpy_znp.api import ZNP
+from zigpy_znp.znp import security
 from zigpy_znp.utils import combine_concurrent_calls
 from zigpy_znp.exceptions import CommandNotRecognized, InvalidCommandResponse
 from zigpy_znp.types.nvids import OsalNvIds
@@ -186,6 +187,9 @@ async def _startup(self, auto_form=False, force_form=False, read_only=False):
         self._znp = znp
         self._znp.set_application(self)
 
+        if not read_only and not force_form:
+            await self._migrate_nvram()
+
         self._bind_callbacks()
 
         # Next, read out the NVRAM item that Zigbee2MQTT writes when it has configured
@@ -928,6 +932,58 @@ async def _set_led_mode(self, *, led: t.uint8_t, mode: c.util.LEDMode) -> None:
         except (asyncio.TimeoutError, CommandNotRecognized):
             LOGGER.info("This build of Z-Stack does not appear to support LED control")
 
+    async def _migrate_nvram(self):
+        """
+        Migrates NVRAM entries using the `ZIGPY_ZNP_MIGRATION_ID` NVRAM item.
+        """
+
+        try:
+            migration_id = await self._znp.nvram.osal_read(
+                OsalNvIds.ZIGPY_ZNP_MIGRATION_ID, item_type=t.uint8_t
+            )
+        except KeyError:
+            migration_id = 0
+
+        initial_migration_id = migration_id
+
+        # Migration 1: empty `ADDRMGR` entries are version-dependent and were improperly
+        #              written for CC253x devices.
+        #
+        #              This migration is stateless and can safely be run more than once:
+        #              the only downside is that startup times increase by 10s on newer
+        #              coordinators, which is why the migration ID is persisted.
+        if migration_id < 1:
+            try:
+                entries = await security.read_addr_manager_entries(self._znp)
+            except KeyError:
+                pass
+            else:
+                fixed_entries = []
+
+                for entry in entries:
+                    if entry.extAddr != t.EUI64.convert("FF:FF:FF:FF:FF:FF:FF:FF"):
+                        fixed_entries.append(entry)
+                    elif self._znp.version == 3.30:
+                        fixed_entries.append(const.EMPTY_ADDR_MGR_ENTRY_ZSTACK3)
+                    else:
+                        fixed_entries.append(const.EMPTY_ADDR_MGR_ENTRY_ZSTACK1)
+
+                if entries != fixed_entries:
+                    LOGGER.warning(
+                        "Repairing %d invalid empty address manager entries (total %d)",
+                        sum(i != j for i, j in zip(entries, fixed_entries)),
+                        len(entries),
+                    )
+                    await security.write_addr_manager_entries(self._znp, fixed_entries)
+
+            migration_id = 1
+
+        if initial_migration_id != migration_id:
+            await self._znp.nvram.osal_write(
+                OsalNvIds.ZIGPY_ZNP_MIGRATION_ID, t.uint8_t(migration_id), create=True
+            )
+            await self._znp.reset()
+
     async def _write_stack_settings(self, *, reset_if_changed: bool) -> None:
         """
         Writes network-independent Z-Stack settings to NVRAM.

zigpy_znp/znp/security.py

@@ -153,7 +153,7 @@ async def write_tc_frame_counter(
         )
 
 
-async def read_addr_mgr_entries(znp: ZNP) -> typing.Sequence[t.AddrMgrEntry]:
+async def read_addr_manager_entries(znp: ZNP) -> typing.Sequence[t.AddrMgrEntry]:
     if znp.version >= 3.30:
         entries = [
             entry
@@ -259,7 +259,7 @@ async def read_devices(znp: ZNP) -> typing.Sequence[StoredDevice]:
     if znp.version > 1.2:
         tclk_seed = await znp.nvram.osal_read(OsalNvIds.TCLK_SEED, item_type=t.KeyData)
 
-    addr_mgr = await read_addr_mgr_entries(znp)
+    addr_mgr = await read_addr_manager_entries(znp)
     devices = {}
 
     for entry in addr_mgr:
@@ -317,30 +317,13 @@ async def read_devices(znp: ZNP) -> typing.Sequence[StoredDevice]:
 
 
 async def write_addr_manager_entries(
-    znp: ZNP, devices: typing.Sequence[StoredDevice]
+    znp: ZNP, entries: typing.Sequence[t.AddrMgrEntry]
 ) -> None:
-    entries = []
-
-    for dev in devices:
-        entry = t.AddrMgrEntry(
-            type=t.AddrMgrUserType.Default,
-            nwkAddr=dev.node_info.nwk,
-            extAddr=dev.node_info.ieee,
-        )
-
-        if dev.key is not None:
-            entry.type |= t.AddrMgrUserType.Security
-
-        if dev.is_child:
-            entry.type |= t.AddrMgrUserType.Assoc
-
-        entries.append(entry)
-
     if znp.version >= 3.30:
         await znp.nvram.write_table(
             item_id=ExNvIds.ADDRMGR,
             values=entries,
-            fill_value=const.EMPTY_ADDR_MGR_ENTRY,
+            fill_value=const.EMPTY_ADDR_MGR_ENTRY_ZSTACK3,
         )
         return
 
@@ -349,7 +332,11 @@ async def write_addr_manager_entries(
     old_entries = await znp.nvram.osal_read(
         OsalNvIds.ADDRMGR, item_type=t.AddressManagerTable
     )
-    new_entries = len(old_entries) * [const.EMPTY_ADDR_MGR_ENTRY]
+
+    if znp.version >= 3.30:
+        new_entries = len(old_entries) * [const.EMPTY_ADDR_MGR_ENTRY_ZSTACK3]
+    else:
+        new_entries = len(old_entries) * [const.EMPTY_ADDR_MGR_ENTRY_ZSTACK1]
 
     # Purposefully throw an `IndexError` if we are trying to write too many entries
     for index, entry in enumerate(entries):
@@ -449,8 +436,25 @@ async def write_devices(
         if len(new_link_key_table_value) > len(old_link_key_table):
             raise RuntimeError("New link key table is larger than the current one")
 
+    addr_mgr_entries = []
+
+    for dev in devices:
+        entry = t.AddrMgrEntry(
+            type=t.AddrMgrUserType.Default,
+            nwkAddr=dev.node_info.nwk,
+            extAddr=dev.node_info.ieee,
+        )
+
+        if dev.key is not None:
+            entry.type |= t.AddrMgrUserType.Security
+
+        if dev.is_child:
+            entry.type |= t.AddrMgrUserType.Assoc
+
+        addr_mgr_entries.append(entry)
+
     # Postpone writes until all of the table entries have been created
-    await write_addr_manager_entries(znp, devices)
+    await write_addr_manager_entries(znp, addr_mgr_entries)
 
     if old_link_key_table is None:
         return