Fix broken empty address manager entries on startup on affected devices

Author: puddly

tests/application/test_startup.py

@@ -1,5 +1,6 @@
 import pytest
 
+import zigpy_znp.const as const
 import zigpy_znp.types as t
 import zigpy_znp.config as conf
 import zigpy_znp.commands as c
@@ -267,3 +268,31 @@ async def test_auto_form_necessary(device, make_application, mocker):
     assert nvram[OsalNvIds.ZDO_DIRECT_CB] == t.Bool(True).serialize()
 
     await app.shutdown()
+
+
+@pytest.mark.parametrize("device", [FormedZStack1CC2531, FormedZStack3CC2531])
+async def test_addrmgr_rewrite_fix(device, make_application, mocker):
+    app, znp_server = make_application(server_cls=device)
+
+    nvram = znp_server._nvram[ExNvIds.LEGACY]
+
+    assert (
+        nvram[OsalNvIds.ADDRMGR].count(const.EMPTY_ADDR_MGR_ENTRY_ZSTACK1.serialize())
+        > 2
+    )
+    nvram[OsalNvIds.ADDRMGR] = nvram[OsalNvIds.ADDRMGR].replace(
+        const.EMPTY_ADDR_MGR_ENTRY_ZSTACK1.serialize(),
+        const.EMPTY_ADDR_MGR_ENTRY_ZSTACK3.serialize(),
+        2,
+    )
+
+    assert const.EMPTY_ADDR_MGR_ENTRY_ZSTACK1.serialize() in nvram[OsalNvIds.ADDRMGR]
+    assert const.EMPTY_ADDR_MGR_ENTRY_ZSTACK3.serialize() in nvram[OsalNvIds.ADDRMGR]
+
+    await app.startup()
+    await app.shutdown()
+
+    assert const.EMPTY_ADDR_MGR_ENTRY_ZSTACK1.serialize() in nvram[OsalNvIds.ADDRMGR]
+    assert (
+        const.EMPTY_ADDR_MGR_ENTRY_ZSTACK3.serialize() not in nvram[OsalNvIds.ADDRMGR]
+    )

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/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

@@ -39,6 +39,7 @@ class OsalNvIds(BaseNvIds):
     # Introduced by zigbeer/zigbee-shepherd and now used by Zigbee2MQTT
     HAS_CONFIGURED_ZSTACK1 = 0x0F00
     HAS_CONFIGURED_ZSTACK3 = 0x0060
+    ZIGPY_ZNP_MIGRATION_ID = 0x0FEE
 
     # OSAL NV item IDs
     EXTADDR = 0x0001

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,53 @@ 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
+        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 not in (
+                        const.EMPTY_ADDR_MGR_ENTRY_ZSTACK1,
+                        const.EMPTY_ADDR_MGR_ENTRY_ZSTACK3,
+                    ):
+                        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("Fixing broken address manager entries")
+                    await security.write_addr_manager_entries(self._znp, fixed_entries)
+
+            migration_id = 1
+
+        await self._znp.nvram.osal_write(
+            OsalNvIds.ZIGPY_ZNP_MIGRATION_ID, t.uint8_t(migration_id), create=True
+        )
+
+        if initial_migration_id != migration_id:
+            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