Proper child table restoration (#686)

puddly · web-flow · commit 9bb0cf78b33f · 2025-08-21T09:56:44.000-04:00
* Restore child data more reliably to NVRAM

* Remove NVRAM token writing

* Revert `NV3ChildTableEntry`

* Re-add NVRAM stuff

* Ensure we reset after writing to NVRAM

* Clean things up

* Tests

* Ensure things work without NV3 interface
diff --git a/bellows/ezsp/v10/__init__.py b/bellows/ezsp/v10/__init__.py
@@ -6,6 +6,7 @@
 import voluptuous
 
 import bellows.config
+from bellows.exception import InvalidCommandError
 import bellows.types as t
 
 from . import commands, config
@@ -39,3 +40,47 @@ async def write_child_data(self, children: dict[t.EUI64, t.NWK]) -> None:
                     timeout_remaining=0,
                 ),
             )
+
+            # There is unfortunately an SDK bug with `setChildData`: some sort of
+            # internal flag in the NVRAM child table is not correctly set (0x00). For
+            # working coordinators, it holds the value 0x80. We need to carefully tweak
+            # this value to ensure restoration works 100%.
+            try:
+                rsp = await self.getTokenData(
+                    token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE, index=index
+                )
+                if t.sl_Status.from_ember_status(rsp.status) != t.sl_Status.OK:
+                    LOGGER.warning(
+                        "Failed to read NVRAM child info for %d: %r", index, rsp
+                    )
+                    continue
+            except (InvalidCommandError, AttributeError):
+                LOGGER.debug("NV3 interface not available, skipping")
+                continue
+
+            # We need to be careful and ensure that the value in NVRAM matches our
+            # expected format (other than the flag byte)
+            expected_entry = t.NV3ChildTableEntry(eui64=eui64, id=nwk, flags=0x80)
+
+            if rsp.value != expected_entry.replace(flags=rsp.value[-1]).serialize():
+                LOGGER.warning(
+                    "Unexpected NVRAM child info for %d: %r, expected %r",
+                    index,
+                    rsp.value,
+                    expected_entry.serialize(),
+                )
+                continue
+
+            # Once we have it fully parsed, write in the correct value (if necessary).
+            # The reason we do this roundabout read/write dance is because we can't be
+            # sure the format in NVRAM will be static.
+            entry, remaining = t.NV3ChildTableEntry.deserialize(rsp.value)
+            assert not remaining
+
+            if entry.flags != 0x80:
+                (status,) = await self.setTokenData(
+                    token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE,
+                    index=index,
+                    token_data=expected_entry.serialize(),
+                )
+                assert t.sl_Status.from_ember_status(status) == t.sl_Status.OK
diff --git a/bellows/types/struct.py b/bellows/types/struct.py
@@ -717,6 +717,14 @@ class NV3StackNetworkManagementToken(EzspStruct):
     padding: basic.uint8_t
 
 
+class NV3ChildTableEntry(EzspStruct):
+    """NV3 node table entry token value."""
+
+    eui64: named.EUI64
+    id: named.NWK
+    flags: basic.uint8_t
+
+
 class SlRxPacketInfo(EzspStruct):
     """Received packet information.
 
diff --git a/bellows/zigbee/application.py b/bellows/zigbee/application.py
@@ -413,14 +413,6 @@ async def write_network_info(
 
         await ezsp.write_link_keys(network_info.key_table)
 
-        children_with_nwk_addresses = {
-            eui64: network_info.nwk_addresses[eui64]
-            for eui64 in network_info.children
-            if eui64 in network_info.nwk_addresses
-        }
-
-        await ezsp.write_child_data(children_with_nwk_addresses)
-
         # Set the network settings
         parameters = t.EmberNetworkParameters()
         parameters.panId = t.EmberPanId(network_info.pan_id)
@@ -439,6 +431,17 @@ async def write_network_info(
         if network_info.nwk_update_id != 0:
             await ezsp.write_nwk_update_id(network_info.nwk_update_id)
 
+        children_with_nwk_addresses = {
+            eui64: network_info.nwk_addresses[eui64]
+            for eui64 in network_info.children
+            if eui64 in network_info.nwk_addresses
+        }
+
+        # Resetting the adapter after writing the child table is important! Otherwise,
+        # NVRAM will not be fully re-read, causing issues.
+        await ezsp.write_child_data(children_with_nwk_addresses)
+        await self._reset()
+
         await self._ensure_network_running()
 
     async def reset_network_info(self):
diff --git a/tests/test_ezsp_v10.py b/tests/test_ezsp_v10.py
@@ -2,6 +2,8 @@
 
 import pytest
 
+from bellows.exception import InvalidCommandError
+from bellows.ezsp.v9.commands import GetTokenDataRsp
 import bellows.ezsp.v10
 import bellows.types as t
 
@@ -48,12 +50,32 @@ async def test_pre_permit(ezsp_f):
 async def test_write_child_data(ezsp_f) -> None:
     ezsp_f.setChildData.return_value = [t.EmberStatus.SUCCESS]
 
-    await ezsp_f.write_child_data(
-        {
-            t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57"): 0xC06B,
-            t.EUI64.convert("00:18:4b:00:1c:a1:b8:46"): 0x1234,
-        }
-    )
+    # Mock NVRAM entries with correct flags
+    eui64_1 = t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57")
+    eui64_2 = t.EUI64.convert("00:18:4b:00:1c:a1:b8:46")
+
+    def mock_get_token_data(token, index):
+        if index == 0:
+            return GetTokenDataRsp(
+                status=t.EmberStatus.SUCCESS,
+                value=t.NV3ChildTableEntry(
+                    eui64=eui64_1, id=0xC06B, flags=0x80
+                ).serialize(),
+            )
+        elif index == 1:
+            return GetTokenDataRsp(
+                status=t.EmberStatus.SUCCESS,
+                value=t.NV3ChildTableEntry(
+                    eui64=eui64_2, id=0x1234, flags=0x80
+                ).serialize(),
+            )
+        else:
+            return GetTokenDataRsp(status=t.EmberStatus.ERR_FATAL, value=b"")
+
+    ezsp_f.getTokenData = AsyncMock(side_effect=mock_get_token_data)
+    ezsp_f.setTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
+
+    await ezsp_f.write_child_data({eui64_1: 0xC06B, eui64_2: 0x1234})
 
     assert ezsp_f.setChildData.mock_calls == [
         call(
@@ -81,3 +103,212 @@ async def test_write_child_data(ezsp_f) -> None:
             ),
         ),
     ]
+
+    assert ezsp_f.getTokenData.mock_calls == [
+        call(token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE, index=0),
+        call(token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE, index=1),
+    ]
+    assert ezsp_f.setTokenData.mock_calls == []
+
+
+async def test_write_child_data_nvram_read_failure(ezsp_f) -> None:
+    """Test write_child_data when NVRAM read fails for some entries."""
+    ezsp_f.setChildData.return_value = [t.EmberStatus.SUCCESS]
+
+    def mock_get_token_data(token, index):
+        if index == 0:
+            return GetTokenDataRsp(status=t.EmberStatus.ERR_FATAL, value=b"")
+        elif index == 1:
+            return GetTokenDataRsp(
+                status=t.EmberStatus.SUCCESS,
+                value=t.NV3ChildTableEntry(
+                    eui64=t.EUI64.convert("00:18:4b:00:1c:a1:b8:46"),
+                    id=0x1234,
+                    flags=0x80,
+                ).serialize(),
+            )
+        else:
+            return GetTokenDataRsp(status=t.EmberStatus.ERR_FATAL, value=b"")
+
+    ezsp_f.getTokenData = AsyncMock(side_effect=mock_get_token_data)
+    ezsp_f.setTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
+
+    await ezsp_f.write_child_data(
+        {
+            t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57"): 0xC06B,
+            t.EUI64.convert("00:18:4b:00:1c:a1:b8:46"): 0x1234,
+        }
+    )
+
+    assert ezsp_f.getTokenData.mock_calls == [
+        call(token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE, index=0),
+        call(token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE, index=1),
+    ]
+    assert ezsp_f.setTokenData.mock_calls == []
+
+
+async def test_write_child_data_nvram_format_mismatch(ezsp_f) -> None:
+    """Test write_child_data when NVRAM format doesn't match expected format."""
+    ezsp_f.setChildData.return_value = [t.EmberStatus.SUCCESS]
+
+    ezsp_f.getTokenData = AsyncMock(
+        return_value=GetTokenDataRsp(
+            status=t.EmberStatus.SUCCESS,
+            value=b"\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b",  # Wrong format
+        )
+    )
+    ezsp_f.setTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
+
+    await ezsp_f.write_child_data({t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57"): 0xC06B})
+
+    assert ezsp_f.getTokenData.mock_calls == [
+        call(token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE, index=0)
+    ]
+    # No setTokenData due to format mismatch
+    assert ezsp_f.setTokenData.mock_calls == []
+
+
+async def test_write_child_data_nvram_flags_correction_needed(ezsp_f) -> None:
+    """Test write_child_data when NVRAM flags need correction."""
+    ezsp_f.setChildData.return_value = [t.EmberStatus.SUCCESS]
+
+    ezsp_f.getTokenData = AsyncMock(
+        return_value=GetTokenDataRsp(
+            status=t.EmberStatus.SUCCESS,
+            value=t.NV3ChildTableEntry(
+                eui64=t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57"), id=0xC06B, flags=0x00
+            ).serialize(),
+        )
+    )
+    ezsp_f.setTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
+
+    await ezsp_f.write_child_data({t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57"): 0xC06B})
+
+    assert ezsp_f.getTokenData.mock_calls == [
+        call(token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE, index=0)
+    ]
+
+    assert ezsp_f.setTokenData.mock_calls == [
+        call(
+            token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE,
+            index=0,
+            token_data=t.NV3ChildTableEntry(
+                eui64=t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57"), id=0xC06B, flags=0x80
+            ).serialize(),
+        )
+    ]
+
+
+async def test_write_child_data_nvram_flags_already_correct(ezsp_f) -> None:
+    """Test write_child_data when NVRAM flags are already correct."""
+    ezsp_f.setChildData.return_value = [t.EmberStatus.SUCCESS]
+
+    ezsp_f.getTokenData = AsyncMock(
+        return_value=GetTokenDataRsp(
+            status=t.EmberStatus.SUCCESS,
+            value=t.NV3ChildTableEntry(
+                eui64=t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57"), id=0xC06B, flags=0x80
+            ).serialize(),
+        )
+    )
+    ezsp_f.setTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
+
+    await ezsp_f.write_child_data({t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57"): 0xC06B})
+
+    assert ezsp_f.getTokenData.mock_calls == [
+        call(token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE, index=0)
+    ]
+    assert ezsp_f.setTokenData.mock_calls == []
+
+
+async def test_write_child_data_nvram_set_token_failure(ezsp_f) -> None:
+    """Test write_child_data when setTokenData fails during flag correction."""
+    ezsp_f.setChildData.return_value = [t.EmberStatus.SUCCESS]
+
+    ezsp_f.getTokenData = AsyncMock(
+        return_value=GetTokenDataRsp(
+            status=t.EmberStatus.SUCCESS,
+            value=t.NV3ChildTableEntry(
+                eui64=t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57"), id=0xC06B, flags=0x00
+            ).serialize(),
+        )
+    )
+    ezsp_f.setTokenData = AsyncMock(return_value=[t.EmberStatus.ERR_FATAL])
+
+    with pytest.raises(AssertionError):
+        await ezsp_f.write_child_data(
+            {t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57"): 0xC06B}
+        )
+
+
+async def test_write_child_data_multiple_entries_mixed_scenarios(ezsp_f) -> None:
+    """Test write_child_data with multiple entries covering various scenarios."""
+    ezsp_f.setChildData.return_value = [t.EmberStatus.SUCCESS]
+
+    def mock_get_token_data(token, index):
+        if index == 0:  # Read failure
+            return GetTokenDataRsp(status=t.EmberStatus.ERR_FATAL, value=b"")
+        elif index == 1:  # Needs correction
+            return GetTokenDataRsp(
+                status=t.EmberStatus.SUCCESS,
+                value=t.NV3ChildTableEntry(
+                    eui64=t.EUI64.convert("00:18:4b:00:1c:a1:b8:46"),
+                    id=0x1234,
+                    flags=0x00,
+                ).serialize(),
+            )
+        elif index == 2:  # Already correct
+            return GetTokenDataRsp(
+                status=t.EmberStatus.SUCCESS,
+                value=t.NV3ChildTableEntry(
+                    eui64=t.EUI64.convert("00:22:33:44:55:66:77:88"),
+                    id=0x5678,
+                    flags=0x80,
+                ).serialize(),
+            )
+        return GetTokenDataRsp(status=t.EmberStatus.ERR_FATAL, value=b"")
+
+    ezsp_f.getTokenData = AsyncMock(side_effect=mock_get_token_data)
+    ezsp_f.setTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
+
+    await ezsp_f.write_child_data(
+        {
+            t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57"): 0xC06B,
+            t.EUI64.convert("00:18:4b:00:1c:a1:b8:46"): 0x1234,
+            t.EUI64.convert("00:22:33:44:55:66:77:88"): 0x5678,
+        }
+    )
+
+    assert ezsp_f.getTokenData.mock_calls == [
+        call(token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE, index=0),
+        call(token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE, index=1),
+        call(token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE, index=2),
+    ]
+
+    # Only entry 2 needs flag correction
+    assert ezsp_f.setTokenData.mock_calls == [
+        call(
+            token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE,
+            index=1,
+            token_data=t.NV3ChildTableEntry(
+                eui64=t.EUI64.convert("00:18:4b:00:1c:a1:b8:46"), id=0x1234, flags=0x80
+            ).serialize(),
+        )
+    ]
+
+
+async def test_write_child_data_nv3_interface_unavailable(ezsp_f) -> None:
+    """Test write_child_data when NV3 interface is not available."""
+    ezsp_f.setChildData.return_value = [t.EmberStatus.SUCCESS]
+    ezsp_f.getTokenData = AsyncMock(
+        side_effect=InvalidCommandError("NV3 not available")
+    )
+    ezsp_f.setTokenData = AsyncMock(return_value=[t.EmberStatus.SUCCESS])
+
+    # Should complete without raising an exception
+    await ezsp_f.write_child_data({t.EUI64.convert("00:0b:57:ff:fe:2b:d4:57"): 0xC06B})
+
+    assert ezsp_f.getTokenData.mock_calls == [
+        call(token=t.NV3KeyId.NVM3KEY_STACK_CHILD_TABLE, index=0)
+    ]
+    assert ezsp_f.setTokenData.mock_calls == []
