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

esphome/components/api/api_server.cpp

@@ -29,6 +29,29 @@ static const char *const TAG = "api";
 // APIServer
 APIServer *global_api_server = nullptr;  // NOLINT(cppcoreguidelines-avoid-non-const-global-variables)
 
+#ifdef USE_API_NOISE
+static constexpr uint32_t NOISE_PSK_PREF_HASH = 88491486UL;
+#endif
+
+#if defined(USE_API_NOISE) && defined(USE_OTA_ENCRYPTION_PROVISIONED)
+bool load_saved_noise_psk(noise::psk_t &out) {
+  SavedNoisePsk saved;
+#ifdef USE_PREFERENCE_KEY_LOOKUP
+  const bool loaded =
+      global_preferences->load_from_key(NOISE_PSK_PREF_HASH, reinterpret_cast<uint8_t *>(&saved), sizeof(saved));
+#else
+  // Slot backends need the reservation walk; it only lands on the record when the reservations before
+  // it match a normal boot, otherwise the type checked checksum fails the load
+  const bool loaded = global_preferences->make_preference<SavedNoisePsk>(NOISE_PSK_PREF_HASH, true).load(&saved);
+#endif
+  // The all-zeros record means no key
+  if (!loaded || noise::NoiseContext::is_all_zeros(saved.psk))
+    return false;
+  out = saved.psk;
+  return true;
+}
+#endif
+
 APIServer::APIServer() { global_api_server = this; }
 
 void APIServer::socket_failed_(const LogString *msg) {
@@ -43,8 +66,7 @@ void APIServer::setup() {
 #ifdef USE_API_NOISE
   // Always reserve the slot: flash preferences are positional on esp8266, so
   // a yaml key build must keep the layout of a runtime key build
-  uint32_t hash = 88491486UL;
-  this->noise_pref_ = global_preferences->make_preference<SavedNoisePsk>(hash, true);
+  this->noise_pref_ = global_preferences->make_preference<SavedNoisePsk>(NOISE_PSK_PREF_HASH, true);
 #ifndef USE_API_NOISE_PSK_FROM_YAML
   // A cleared record loads fine but holds no key
   if (this->load_and_apply_noise_psk_() && this->noise_ctx_.has_psk()) {

esphome/components/api/api_server.h

@@ -43,6 +43,11 @@ struct SavedNoisePsk {
   noise::psk_t psk;
 } PACKED;  // NOLINT
 #endif
+#if defined(USE_API_NOISE) && defined(USE_OTA_ENCRYPTION_PROVISIONED)
+/// One-shot read of the provisioned key for a boot without an api server (safe mode); false when
+/// there is no key
+bool load_saved_noise_psk(noise::psk_t &out);
+#endif
 
 class APIServer final : public Component,
                         public Controller

esphome/components/esphome/ota/__init__.py

@@ -316,20 +316,16 @@ async def to_code(config: ConfigType) -> None:
     # One key per device: an api encryption block supplies it (static or
     # runtime) and offers; the ota block only adds the requirement
     api_conf = CORE.config.get(CONF_API) or {}
-    encryption_conf = config.get(CONF_ENCRYPTION)
-    own_key = None
-    if encryption_conf is not None and static_encryption_key(api_conf) is None:
-        own_key = encryption_conf[CONF_KEY]
-    if own_key is not None:
+    if key := static_encryption_key(config) or static_encryption_key(api_conf):
+        # Build time key: the ota keeps its own pointer so safe mode, which
+        # has no api server, still has it
         cg.add_define("USE_OTA_ENCRYPTION")
-        cg.add(var.set_noise_psk(new_psk_progmem(config[CONF_ID], own_key)))
+        cg.add(var.set_noise_psk(new_psk_progmem(config[CONF_ID], key)))
     elif CONF_ENCRYPTION in api_conf:
+        # Runtime key: found in the api server, or in preferences in safe mode
         cg.add_define("USE_OTA_ENCRYPTION")
-        cg.add_define("USE_OTA_ENCRYPTION_FROM_API")
-        if static_encryption_key(api_conf) is None:
-            # The key arrives at runtime, so the offer has to look for it
-            cg.add_define("USE_OTA_ENCRYPTION_PROVISIONED")
-    if encryption_conf is not None:
+        cg.add_define("USE_OTA_ENCRYPTION_PROVISIONED")
+    if CONF_ENCRYPTION in config:
         cg.add_define("USE_OTA_ENCRYPTION_REQUIRED")
 
     # Build flag so lwip_fast_select.c (a .c file that can't include defines.h) sees it.

esphome/components/esphome/ota/ota_esphome.cpp

@@ -1,5 +1,5 @@
 #include "ota_esphome.h"
-#ifdef USE_OTA_ENCRYPTION_FROM_API
+#ifdef USE_OTA_ENCRYPTION_PROVISIONED
 #include "esphome/components/api/api_server.h"
 #endif
 #ifdef USE_OTA
@@ -32,11 +32,13 @@ static const char *const TAG = "esphome.ota";
 
 #ifdef USE_OTA_ENCRYPTION
 const noise::NoiseContext &ESPHomeOTAComponent::noise_context_() const {
-#ifdef USE_OTA_ENCRYPTION_FROM_API
-  return api::global_api_server->get_noise_ctx();
-#else
-  return this->noise_ctx_;
+#ifdef USE_OTA_ENCRYPTION_PROVISIONED
+  // The api server holds the live key; safe mode never constructs it, and then
+  // noise_ctx_ holds the saved key setup() found, if any
+  if (api::global_api_server != nullptr)
+    return api::global_api_server->get_noise_ctx();
 #endif
+  return this->noise_ctx_;
 }
 #endif
 static constexpr uint16_t OTA_BLOCK_SIZE = 8192;
@@ -58,6 +60,16 @@ extern "C" void esphome_wake_ota_component_any_context() {
 }
 
 void ESPHomeOTAComponent::setup() {
+#ifdef USE_OTA_ENCRYPTION_PROVISIONED
+  // Safe mode never constructs the api server, so read the key it saved
+  noise::psk_t psk;
+  if (api::global_api_server == nullptr && api::load_saved_noise_psk(psk)) {
+    this->saved_psk_ = RAMAllocator<noise::psk_t>().make_unique(psk);
+    if (this->saved_psk_ != nullptr) {
+      this->noise_ctx_.set_psk(this->saved_psk_->data());
+    }
+  }
+#endif
   this->server_ = socket::socket_ip_loop_monitored(SOCK_STREAM, 0).release();  // monitored for incoming connections
   if (this->server_ == nullptr) {
     this->server_failed_(LOG_STR("creation"));

esphome/components/esphome/ota/ota_esphome.h

@@ -44,7 +44,7 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
   }
 #endif  // USE_OTA_PASSWORD
 
-#if defined(USE_OTA_ENCRYPTION) && !defined(USE_OTA_ENCRYPTION_FROM_API)
+#ifdef USE_OTA_ENCRYPTION
   /// psk points at 32 bytes that live in flash for the life of the program
   void set_noise_psk(const uint8_t *psk) { this->noise_ctx_.set_psk(psk); }
 #endif
@@ -86,7 +86,8 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
     bool writing{false};    // a produced handshake frame is still being flushed
     uint8_t frame_buf[noise::FRAME_HEADER_SIZE + 1 + noise::MAX_HANDSHAKE_SIZE];
   };
-  // The api server's live context when the api has encryption, else our own
+  // The api server's live context when it exists, otherwise our own (a build
+  // time key, or the saved key loaded in safe mode)
   const noise::NoiseContext &noise_context_() const;
   bool noise_start_session_(uint8_t server_feature_flags);
   bool handle_noise_handshake_();
@@ -148,8 +149,10 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
   RAMUniquePtr<uint8_t[]> auth_buf_;
 #endif  // USE_OTA_PASSWORD
 #ifdef USE_OTA_ENCRYPTION
-#ifndef USE_OTA_ENCRYPTION_FROM_API
   noise::NoiseContext noise_ctx_;
+#ifdef USE_OTA_ENCRYPTION_PROVISIONED
+  // Backs noise_ctx_ in safe mode, where no api server holds the saved key
+  RAMUniquePtr<noise::psk_t> saved_psk_;
 #endif
   RAMUniquePtr<NoiseSession> noise_;
 #endif  // USE_OTA_ENCRYPTION

esphome/components/noise/__init__.py

@@ -5,11 +5,12 @@ from typing import Any
 import esphome.codegen as cg
 import esphome.config_validation as cv
 from esphome.const import CONF_ENCRYPTION, CONF_KEY
-from esphome.core import ID
+from esphome.core import CORE, ID
 from esphome.cpp_generator import MockObj
 from esphome.types import ConfigType
 
 CODEOWNERS = ["@esphome/core"]
+DOMAIN = "noise"
 
 noise_ns = cg.esphome_ns.namespace("noise")
 
@@ -70,11 +71,16 @@ def static_encryption_key(conf: ConfigType) -> str | None:
 
 def new_psk_progmem(parent_id: ID, key: str) -> MockObj:
     """Emit the decoded key as a PROGMEM array; the component keeps a pointer
-    so the key never occupies RAM."""
-    return cg.progmem_array(
-        ID(f"{parent_id.id}_psk", is_declaration=True, type=cg.uint8),
-        list(decode_encryption_key(key)),
-    )
+    so the key never occupies RAM. Components sharing one key (api and ota)
+    share the array."""
+    decoded = decode_encryption_key(key)
+    arrays: dict[bytes, MockObj] = CORE.data.setdefault(DOMAIN, {})
+    if (array := arrays.get(decoded)) is None:
+        array = arrays[decoded] = cg.progmem_array(
+            ID(f"{parent_id.id}_psk", is_declaration=True, type=cg.uint8),
+            list(decoded),
+        )
+    return array
 
 
 def encryption_schema(config: ConfigType | None) -> ConfigType:

esphome/core/defines.h

@@ -283,7 +283,6 @@
 #define USE_RUNTIME_STATS
 #define USE_OTA
 #define USE_OTA_ENCRYPTION
-#define USE_OTA_ENCRYPTION_FROM_API
 #define USE_OTA_ENCRYPTION_PROVISIONED
 #define USE_OTA_ENCRYPTION_REQUIRED
 #define USE_OTA_PASSWORD

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

@@ -15,6 +15,12 @@ import os
 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.
@@ -42,16 +48,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()