[esphome] Keep the OTA encryption key without the api server so safe mode uploads work (#19349)

tests/component_tests/ota/test_esphome_ota.py

@@ -476,43 +476,36 @@ def test_static_encryption_key() -> None:
     ("yaml_name", "defines_present", "defines_absent"),
     [
         # An api key alone compiles the transport in without requiring it;
-        # the device uses the api server's key, not a copy
+        # the ota keeps its own pointer to the key so safe mode, which never
+        # constructs the api server, can still use it
         (
             "api_key_offer",
-            {"USE_OTA_ENCRYPTION", "USE_OTA_ENCRYPTION_FROM_API"},
+            {"USE_OTA_ENCRYPTION"},
             {"USE_OTA_ENCRYPTION_REQUIRED", "USE_OTA_ENCRYPTION_PROVISIONED"},
         ),
         # A password still guards plaintext uploads on an offering device
         (
             "api_key_offer_password",
-            {"USE_OTA_ENCRYPTION", "USE_OTA_ENCRYPTION_FROM_API", "USE_OTA_PASSWORD"},
+            {"USE_OTA_ENCRYPTION", "USE_OTA_PASSWORD"},
             {"USE_OTA_ENCRYPTION_REQUIRED", "USE_OTA_ENCRYPTION_PROVISIONED"},
         ),
         # The ota encryption block is what makes the device refuse plaintext
         (
             "encryption_required",
-            {
-                "USE_OTA_ENCRYPTION",
-                "USE_OTA_ENCRYPTION_REQUIRED",
-                "USE_OTA_ENCRYPTION_FROM_API",
-            },
+            {"USE_OTA_ENCRYPTION", "USE_OTA_ENCRYPTION_REQUIRED"},
             {"USE_OTA_ENCRYPTION_PROVISIONED"},
         ),
         # Without api encryption the ota key is the device's own
         (
             "own_key",
             {"USE_OTA_ENCRYPTION", "USE_OTA_ENCRYPTION_REQUIRED"},
-            {"USE_OTA_ENCRYPTION_FROM_API", "USE_OTA_ENCRYPTION_PROVISIONED"},
+            {"USE_OTA_ENCRYPTION_PROVISIONED"},
         ),
         # A key provisioned at runtime lives in the api server; the device
         # offers with it once provisioned and never requires it
         (
             "runtime_api_key",
-            {
-                "USE_OTA_ENCRYPTION",
-                "USE_OTA_ENCRYPTION_FROM_API",
-                "USE_OTA_ENCRYPTION_PROVISIONED",
-            },
+            {"USE_OTA_ENCRYPTION", "USE_OTA_ENCRYPTION_PROVISIONED"},
             {"USE_OTA_ENCRYPTION_REQUIRED"},
         ),
         # No api encryption at all keeps the noise glue out of the build
@@ -522,7 +515,6 @@ def test_static_encryption_key() -> None:
             {
                 "USE_OTA_ENCRYPTION",
                 "USE_OTA_ENCRYPTION_REQUIRED",
-                "USE_OTA_ENCRYPTION_FROM_API",
                 "USE_OTA_ENCRYPTION_PROVISIONED",
             },
         ),
@@ -541,8 +533,10 @@ def test_encryption_offer_codegen(
     assert defines_present <= defines
     assert not (defines_absent & defines)
     encrypted = "USE_OTA_ENCRYPTION" in defines_present
-    own_key = encrypted and "USE_OTA_ENCRYPTION_FROM_API" not in defines_present
+    own_key = encrypted and "USE_OTA_ENCRYPTION_PROVISIONED" not in defines_present
     assert ("esphome_esphomeotacomponent_id->set_noise_psk(" in main_cpp) is own_key
+    # The api shares the ota's array instead of emitting the same key twice
+    assert main_cpp.count("_psk[] PROGMEM") == (1 if own_key else 0)
     assert ("set_auth_password(" in main_cpp) is ("USE_OTA_PASSWORD" in defines_present)
     # The noise transport source compiles only when the define is set
     assert FILTER_SOURCE_FILES() == ([] if encrypted else ["ota_esphome_noise.cpp"])

tests/integration/fixtures/host_ota_encrypted_safe_mode.yaml

@@ -0,0 +1,13 @@
+esphome:
+  name: host-ota-test
+host:
+api:
+  encryption:
+    key: "AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8="
+ota:
+  - platform: esphome
+    port: __OTA_PORT__
+    encryption:
+safe_mode:
+logger:
+  level: DEBUG

tests/integration/fixtures/host_ota_provisioned_api_key_safe_mode.yaml

@@ -0,0 +1,11 @@
+esphome:
+  name: host-ota-test
+host:
+api:
+  encryption:
+ota:
+  - platform: esphome
+    port: __OTA_PORT__
+safe_mode:
+logger:
+  level: DEBUG

tests/integration/host_prefs.py

@@ -14,6 +14,12 @@ from __future__ import annotations
 from pathlib import Path
 import struct
 
+_ENTRY = struct.Struct("<IB")  # key, data length
+# Must match esphome::safe_mode::RTC_KEY in safe_mode.h
+_SAFE_MODE_RTC_KEY = 233825507
+# Must match esphome::safe_mode::SafeModeComponent::ENTER_SAFE_MODE_MAGIC
+_ENTER_SAFE_MODE_MAGIC = 0x5AFE5AFE
+
 
 def host_prefs_path(device_name: str) -> Path:
     """Return the on-disk prefs file path for a host-platform device."""
@@ -34,16 +40,33 @@ def write_host_prefs(device_name: str, entries: dict[int, bytes]) -> Path:
     for key, data in entries.items():
         if len(data) > 255:
             raise ValueError(f"Preference data too long: {len(data)} bytes (max 255)")
-        payload += struct.pack("<IB", key, len(data)) + data
+        payload += _ENTRY.pack(key, len(data)) + data
     path = host_prefs_path(device_name)
     path.parent.mkdir(parents=True, exist_ok=True)
     path.write_bytes(payload)
     return path
 
 
-def write_host_pref(device_name: str, key: int, data: bytes) -> Path:
-    """Write a single preference entry, replacing the file's contents.
+def read_host_prefs(device_name: str) -> dict[int, bytes]:
+    """Read the preference entries of a host-platform device; empty when
+    the file does not exist."""
+    path = host_prefs_path(device_name)
+    if not path.exists():
+        return {}
+    payload = path.read_bytes()
+    entries: dict[int, bytes] = {}
+    pos = 0
+    while pos < len(payload):
+        key, length = _ENTRY.unpack_from(payload, pos)
+        pos += _ENTRY.size
+        entries[key] = payload[pos : pos + length]
+        pos += length
+    return entries
 
-    Returns the path that was written.
-    """
-    return write_host_prefs(device_name, {key: data})
+
+def force_safe_mode(device_name: str) -> None:
+    """Make the next boot of a host-platform device enter safe mode; other
+    saved preferences are kept."""
+    entries = read_host_prefs(device_name)
+    entries[_SAFE_MODE_RTC_KEY] = struct.pack("<I", _ENTER_SAFE_MODE_MAGIC)
+    write_host_prefs(device_name, entries)

tests/integration/test_host_ota.py

@@ -29,6 +29,7 @@ from .const import (
     PROVISIONING_PSK,
     ZERO_PSK,
 )
+from .host_prefs import force_safe_mode
 from .types import APIClientConnectedFactory, CompileFunction, ConfigWriter
 
 DEVICE_NAME = "host-ota-test"
@@ -166,6 +167,15 @@ class _Device:
         assert self.proc.returncode is None, "process died on rejected OTA"
 
 
+async def _provision_key(
+    dev: _Device, api_client_connected: APIClientConnectedFactory
+) -> None:
+    """Provision PROVISIONING_PSK over the api and wait for it to activate."""
+    async with api_client_connected(port=dev.api_port, noise_psk=ZERO_PSK) as client:
+        assert await client.noise_encryption_set_key(PROVISIONING_PSK) is True
+    await asyncio.sleep(KEY_ACTIVATION_DELAY)
+
+
 @pytest.mark.asyncio
 async def test_host_ota_self_update(
     yaml_config: str,
@@ -227,6 +237,31 @@ async def test_host_ota_encrypted(
         await dev.ota(None, API_KEY, "encrypted OTA reported failure")
 
 
+@pytest.mark.asyncio
+async def test_host_ota_encrypted_safe_mode(
+    yaml_config: str,
+    write_yaml_config: ConfigWriter,
+    compile_esphome: CompileFunction,
+    reserved_tcp_port: tuple[int, socket.socket],
+) -> None:
+    """Safe mode never constructs the api server, so an encrypted OTA with the
+    api key has to run on the ota component's own copy of that key."""
+    pytest.importorskip("aioesphomeapi.noise")
+    dev = _Device(
+        *await _build(
+            yaml_config, write_yaml_config, compile_esphome, reserved_tcp_port
+        )
+    )
+    # The api port never opens in safe mode, so wait for the log line instead
+    force_safe_mode(DEVICE_NAME)
+    async with run_binary(dev.binary_path, line_callback=dev.on_log) as (proc, lines):
+        dev.proc = proc
+        await _wait_for_line(lines, "SAFE MODE IS ACTIVE", PORT_WAIT_TIMEOUT)
+        await _wait_for_port(LOCALHOST, dev.ota_port, PORT_WAIT_TIMEOUT)
+        # The safe mode boot clears the counter, so the re-exec boots normally
+        await dev.ota(None, API_KEY, "encrypted OTA in safe mode reported failure")
+
+
 @pytest.mark.asyncio
 async def test_host_ota_api_key_offer_with_password(
     yaml_config: str,
@@ -305,11 +340,7 @@ async def test_host_ota_provisioned_api_key(
             None, None, "plaintext upload to an unprovisioned device must succeed"
         )
 
-        async with api_client_connected(
-            port=dev.api_port, noise_psk=ZERO_PSK
-        ) as client:
-            assert await client.noise_encryption_set_key(PROVISIONING_PSK) is True
-        await asyncio.sleep(KEY_ACTIVATION_DELAY)
+        await _provision_key(dev, api_client_connected)
 
         key = PROVISIONING_PSK.decode()
         await dev.ota(
@@ -319,6 +350,45 @@ async def test_host_ota_provisioned_api_key(
         await dev.ota(None, None, "plaintext must stay accepted on an offering device")
 
 
+@pytest.mark.asyncio
+async def test_host_ota_provisioned_api_key_safe_mode(
+    yaml_config: str,
+    write_yaml_config: ConfigWriter,
+    compile_esphome: CompileFunction,
+    reserved_tcp_port: tuple[int, socket.socket],
+    api_client_connected: APIClientConnectedFactory,
+) -> None:
+    """Safe mode never constructs the api server, so the OTA has to load the
+    provisioned key from preferences itself to keep encrypting there."""
+    pytest.importorskip("aioesphomeapi.noise")
+    dev = _Device(
+        *await _build(
+            yaml_config, write_yaml_config, compile_esphome, reserved_tcp_port
+        )
+    )
+    async with run_binary(dev.binary_path, line_callback=dev.on_log) as (proc, _lines):
+        dev.proc = proc
+        await _wait_for_port(LOCALHOST, dev.api_port, PORT_WAIT_TIMEOUT)
+        await _provision_key(dev, api_client_connected)
+
+    # The saved key is already on disk; a host reboot outside an OTA just
+    # exits, so safe mode takes a second start
+    force_safe_mode(DEVICE_NAME)
+    key = PROVISIONING_PSK.decode()
+    async with run_binary(dev.binary_path, line_callback=dev.on_log) as (proc, lines):
+        dev.proc = proc
+        await _wait_for_line(lines, "SAFE MODE IS ACTIVE", PORT_WAIT_TIMEOUT)
+        await _wait_for_port(LOCALHOST, dev.ota_port, PORT_WAIT_TIMEOUT)
+        await dev.ota(
+            None,
+            key,
+            "encrypted upload with the provisioned key must succeed in safe mode",
+        )
+        # The re-exec boots normally and the api reads the same record
+        async with api_client_connected(port=dev.api_port, noise_psk=key):
+            pass
+
+
 @pytest.mark.asyncio
 async def test_host_ota_rejects_garbage(
     yaml_config: str,

tests/integration/test_safe_mode_loop_runs.py

@@ -25,19 +25,13 @@ from __future__ import annotations
 
 import asyncio
 import re
-import struct
 
 import pytest
 
 from .conftest import run_binary
-from .host_prefs import clear_host_prefs, write_host_pref
+from .host_prefs import clear_host_prefs, force_safe_mode
 from .types import CompileFunction, ConfigWriter
 
-# Must match esphome::safe_mode::RTC_KEY in safe_mode.h
-SAFE_MODE_RTC_KEY = 233825507
-# Must match esphome::safe_mode::SafeModeComponent::ENTER_SAFE_MODE_MAGIC
-ENTER_SAFE_MODE_MAGIC = 0x5AFE5AFE
-
 DEVICE_NAME = "safe-mode-loop-runs"
 THREAD_LOG_MARKER = "looping component ran in safe mode"
 
@@ -56,9 +50,7 @@ async def test_safe_mode_loop_runs(
 
     # Compile finished successfully; pre-populate prefs so the *next* run
     # enters safe mode immediately.
-    write_host_pref(
-        DEVICE_NAME, SAFE_MODE_RTC_KEY, struct.pack("<I", ENTER_SAFE_MODE_MAGIC)
-    )
+    force_safe_mode(DEVICE_NAME)
 
     try:
         loop = asyncio.get_running_loop()