diff --git a/esphome/__main__.py b/esphome/__main__.py index 174f9c3c36..11081e271b 100644 --- a/esphome/__main__.py +++ b/esphome/__main__.py @@ -1092,6 +1092,15 @@ def upload_program( port_type = get_port_type(host) + # MQTT and MQTTIP are also OTA paths; MQTTIP gets resolved to a real IP later by + # _resolve_network_devices(). Only SERIAL and BOOTSEL are non-OTA upload paths. + if port_type in (PortType.SERIAL, PortType.BOOTSEL) and getattr( + args, "partition_table", False + ): + raise EsphomeError( + "The option --partition-table can only be used for Over The Air updates." + ) + if port_type == PortType.BOOTSEL: exit_code = upload_using_picotool(config) # Return None for device - BOOTSEL can't be used for logging, @@ -1132,12 +1141,80 @@ def upload_program( binary = CORE.firmware_bin ota_type = espota2.OTA_TYPE_UPDATE_APP + if getattr(args, "partition_table", False): + # Fail fast if the resolved ESPHome OTA config does not enable allow_partition_access. + # The device-side handshake also rejects this with "Device only supports app updates", + # but checking here surfaces the misconfiguration before opening a network connection. + if not ota_conf.get("allow_partition_access"): + raise EsphomeError( + "The option --partition-table requires 'allow_partition_access: true' on the " + "esphome OTA platform in the device's YAML configuration. Add it, recompile, " + "flash a build with the option enabled, and then retry --partition-table." + ) + binary = CORE.partition_table_bin + ota_type = espota2.OTA_TYPE_UPDATE_PARTITION_TABLE if getattr(args, "file", None) is not None: binary = Path(args.file) + if ota_type == espota2.OTA_TYPE_UPDATE_PARTITION_TABLE: + _validate_partition_table_binary(binary) + return espota2.run_ota(network_devices, remote_port, password, binary, ota_type) +# Layout of esp_partition_info_t on flash. Each entry is 32 bytes, leading with a +# 16-bit little-endian magic. ESP-IDF defines ESP_PARTITION_MAGIC = 0x50AA (stored as +# bytes 0xAA, 0x50) for partition entries and ESP_PARTITION_MAGIC_MD5 = 0xEBEB for the +# trailing checksum entry. Padding past the last entry is 0xFF. The full table is +# exactly ESP_PARTITION_TABLE_MAX_LEN bytes. +_PARTITION_TABLE_MAX_LEN = 0xC00 +_ESP_PARTITION_MAGIC = 0x50AA +_ESP_PARTITION_MAGIC_MD5 = 0xEBEB + + +def _validate_partition_table_binary(binary: Path) -> None: + """Validate that ``binary`` looks like an ESP32 partition table image. + + Catches common mistakes (wrong file, truncated build output, swapped --file path) + before opening a network connection so the failure mode is a clear local error + instead of a post-handshake device rejection. + """ + try: + data = binary.read_bytes() + except OSError as err: + raise EsphomeError( + f"Cannot read partition table file '{binary}': {err}" + ) from err + + if len(data) != _PARTITION_TABLE_MAX_LEN: + raise EsphomeError( + f"Partition table file '{binary}' has wrong size: expected " + f"{_PARTITION_TABLE_MAX_LEN} bytes, got {len(data)}. " + "Pass the partition table image (e.g. partitions.bin / partition-table.bin), " + "not the firmware image." + ) + + first_magic = data[0] | (data[1] << 8) + if first_magic != _ESP_PARTITION_MAGIC: + raise EsphomeError( + f"Partition table file '{binary}' does not start with the expected " + f"partition magic 0x{_ESP_PARTITION_MAGIC:04X} (got 0x{first_magic:04X}). " + "This file does not look like an ESP32 partition table." + ) + + # The MD5 checksum entry is required: without it the device-side + # esp_partition_table_verify will accept the table but the bootloader will + # refuse to boot from it. Scan the 32-byte entries for the MD5 magic. + if not any( + (data[off] | (data[off + 1] << 8)) == _ESP_PARTITION_MAGIC_MD5 + for off in range(0, _PARTITION_TABLE_MAX_LEN, 32) + ): + raise EsphomeError( + f"Partition table file '{binary}' is missing the MD5 checksum entry. " + "Regenerate the partition table with gen_esp32part.py or rebuild the project." + ) + + def show_logs(config: ConfigType, args: ArgsProtocol, devices: list[str]) -> int | None: try: module = importlib.import_module("esphome.components." + CORE.target_platform) @@ -1805,6 +1882,11 @@ def parse_args(argv): "--file", help="Manually specify the binary file to upload.", ) + parser_upload.add_argument( + "--partition-table", + help="Upload as partition table (OTA).", + action="store_true", + ) parser_logs = subparsers.add_parser( "logs", diff --git a/esphome/components/esphome/ota/__init__.py b/esphome/components/esphome/ota/__init__.py index bfa5ffb55c..ee3b7f0c20 100644 --- a/esphome/components/esphome/ota/__init__.py +++ b/esphome/components/esphome/ota/__init__.py @@ -16,11 +16,13 @@ from esphome.const import ( CONF_SAFE_MODE, CONF_VERSION, ) -from esphome.core import coroutine_with_priority +from esphome.core import CORE, coroutine_with_priority from esphome.coroutine import CoroPriority import esphome.final_validate as fv from esphome.types import ConfigType +CONF_ALLOW_PARTITION_ACCESS = "allow_partition_access" + _LOGGER = logging.getLogger(__name__) @@ -75,6 +77,10 @@ def ota_esphome_final_validate(config): merged_ota_esphome_configs_by_port[conf_port] = merge_config( merged_ota_esphome_configs_by_port[conf_port], ota_conf ) + if ota_conf.get(CONF_ALLOW_PARTITION_ACCESS) and not CORE.is_esp32: + raise cv.Invalid( + f"{CONF_ALLOW_PARTITION_ACCESS} is only supported on the esp32" + ) else: new_ota_conf.append(ota_conf) @@ -125,6 +131,7 @@ CONFIG_SCHEMA = cv.All( ln882x=8820, rtl87xx=8892, ): cv.port, + cv.Optional(CONF_ALLOW_PARTITION_ACCESS, default=False): cv.boolean, cv.Optional(CONF_PASSWORD): cv.string, cv.Optional(CONF_NUM_ATTEMPTS): cv.invalid( f"'{CONF_SAFE_MODE}' (and its related configuration variables) has moved from 'ota' to its own component. See https://esphome.io/components/safe_mode" @@ -159,6 +166,10 @@ async def to_code(config: ConfigType) -> None: if config[CONF_PASSWORD]: cg.add(var.set_auth_password(config[CONF_PASSWORD])) cg.add_define("USE_OTA_VERSION", config[CONF_VERSION]) + + if config.get(CONF_ALLOW_PARTITION_ACCESS): + cg.add_define("USE_OTA_PARTITIONS") + # Build flag so lwip_fast_select.c (a .c file that can't include defines.h) sees it. cg.add_build_flag("-DUSE_OTA_PLATFORM_ESPHOME") diff --git a/esphome/components/esphome/ota/ota_esphome.cpp b/esphome/components/esphome/ota/ota_esphome.cpp index 955b4dc96f..3ce3f2302d 100644 --- a/esphome/components/esphome/ota/ota_esphome.cpp +++ b/esphome/components/esphome/ota/ota_esphome.cpp @@ -87,6 +87,10 @@ void ESPHomeOTAComponent::setup() { // no wakes fire and loop() falls back to the self-disable safety net. esphome_fast_select_set_ota_listener_sock(esphome_lwip_get_sock(this->server_->get_fd())); #endif + +#ifdef USE_OTA_PARTITIONS + ota::get_running_app_position(this->running_app_offset_, this->running_app_size_); +#endif } void ESPHomeOTAComponent::dump_config() { @@ -100,6 +104,29 @@ void ESPHomeOTAComponent::dump_config() { ESP_LOGCONFIG(TAG, " Password configured"); } #endif +#ifdef USE_OTA_PARTITIONS + ESP_LOGCONFIG(TAG, + " Partition access allowed\n" + " Running app:\n" + " Partition address: 0x%X\n" + " Used size: %zu bytes (0x%X)", + this->running_app_offset_, this->running_app_size_, this->running_app_size_); + +#ifdef USE_ESP32 + ESP_LOGCONFIG(TAG, + " Partition table:\n" + " %-12s %-4s %-8s %-10s %-10s", + "Name", "Type", "Subtype", "Address", "Size"); + esp_partition_iterator_t it = esp_partition_find(ESP_PARTITION_TYPE_ANY, ESP_PARTITION_SUBTYPE_ANY, NULL); + while (it != NULL) { + const esp_partition_t *partition = esp_partition_get(it); + ESP_LOGCONFIG(TAG, " %-12s 0x%-2X 0x%-6X 0x%-8" PRIX32 " 0x%-8" PRIX32, partition->label, partition->type, + partition->subtype, partition->address, partition->size); + it = esp_partition_next(it); + } + esp_partition_iterator_release(it); +#endif +#endif } void ESPHomeOTAComponent::loop() { @@ -118,6 +145,7 @@ static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_COMPRESSION = 0x01; static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_SHA256_AUTH = 0x02; static constexpr uint8_t CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL = 0x04; static constexpr uint8_t SERVER_FEATURE_SUPPORTS_COMPRESSION = 0x01; +static constexpr uint8_t SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS = 0x02; void ESPHomeOTAComponent::handle_handshake_() { /// Handle the OTA handshake and authentication. @@ -215,6 +243,9 @@ void ESPHomeOTAComponent::handle_handshake_() { static_assert(HANDSHAKE_BUF_SIZE >= 2, "handshake_buf_ must hold the 2-byte extended-protocol feature ack"); this->handshake_buf_[0] = ota::OTA_RESPONSE_FEATURE_FLAGS; this->handshake_buf_[1] = (supports_compression ? SERVER_FEATURE_SUPPORTS_COMPRESSION : 0); +#ifdef USE_OTA_PARTITIONS + this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS; +#endif } else { this->handshake_buf_[0] = supports_compression ? ota::OTA_RESPONSE_SUPPORTS_COMPRESSION : ota::OTA_RESPONSE_HEADER_OK; @@ -347,10 +378,12 @@ void ESPHomeOTAComponent::handle_data_() { (static_cast(buf[2]) << 8) | buf[3]; ESP_LOGV(TAG, "Size is %u bytes", ota_size); +#ifndef USE_OTA_PARTITIONS if (ota_type != ota::OTA_TYPE_UPDATE_APP) { error_code = ota::OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE; goto error; // NOLINT(cppcoreguidelines-avoid-goto) } +#endif // Now that we've passed authentication and are actually // starting the update, set the warning status and notify @@ -362,8 +395,8 @@ void ESPHomeOTAComponent::handle_data_() { this->notify_state_(ota::OTA_STARTED, 0.0f, 0); #endif - // This will block for a few seconds as it locks flash - error_code = this->backend_->begin(ota_size); + // begin() may block for a few seconds while it locks flash. + error_code = this->backend_->begin(ota_size, ota_type); if (error_code != ota::OTA_RESPONSE_OK) goto error; // NOLINT(cppcoreguidelines-avoid-goto) update_started = true; @@ -465,6 +498,13 @@ void ESPHomeOTAComponent::handle_data_() { this->notify_state_(ota::OTA_COMPLETED, 100.0f, 0); #endif delay(100); // NOLINT +#ifdef USE_OTA_PARTITIONS + if (ota_type == ota::OTA_TYPE_UPDATE_PARTITION_TABLE) { + // Skip on_safe_shutdown: nvs_flash_deinit() has already invalidated open NVS handles, so + // preferences flush would emit ESP_ERR_NVS_INVALID_HANDLE for every entry. Reboot directly. + App.reboot(); + } +#endif App.safe_reboot(); error: diff --git a/esphome/components/esphome/ota/ota_esphome.h b/esphome/components/esphome/ota/ota_esphome.h index 5043bc33ef..0053ca6969 100644 --- a/esphome/components/esphome/ota/ota_esphome.h +++ b/esphome/components/esphome/ota/ota_esphome.h @@ -98,6 +98,10 @@ class ESPHomeOTAComponent final : public ota::OTAComponent { uint32_t client_connect_time_{0}; static constexpr size_t HANDSHAKE_BUF_SIZE = 5; +#ifdef USE_OTA_PARTITIONS + uint32_t running_app_offset_{0}; + size_t running_app_size_{0}; +#endif uint16_t port_; uint8_t handshake_buf_[HANDSHAKE_BUF_SIZE]; OTAState ota_state_{OTAState::IDLE}; diff --git a/esphome/components/ota/ota_backend.h b/esphome/components/ota/ota_backend.h index 7e7b0f6523..5888a8e12d 100644 --- a/esphome/components/ota/ota_backend.h +++ b/esphome/components/ota/ota_backend.h @@ -42,6 +42,8 @@ enum OTAResponseTypes { OTA_RESPONSE_ERROR_RP2040_NOT_ENOUGH_SPACE = 0x8C, OTA_RESPONSE_ERROR_SIGNATURE_INVALID = 0x8D, OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE = 0x8E, + OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY = 0x8F, + OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE = 0x90, OTA_RESPONSE_ERROR_UNKNOWN = 0xFF, }; @@ -55,6 +57,7 @@ enum OTAState { enum OTAType : uint8_t { OTA_TYPE_UPDATE_APP = 0x00, + OTA_TYPE_UPDATE_PARTITION_TABLE = 0x01, }; /** Listener interface for OTA state changes. diff --git a/esphome/components/ota/ota_backend_arduino_libretiny.cpp b/esphome/components/ota/ota_backend_arduino_libretiny.cpp index dcd71e92dd..4cc99202a7 100644 --- a/esphome/components/ota/ota_backend_arduino_libretiny.cpp +++ b/esphome/components/ota/ota_backend_arduino_libretiny.cpp @@ -13,7 +13,10 @@ static const char *const TAG = "ota.arduino_libretiny"; std::unique_ptr make_ota_backend() { return make_unique(); } -OTAResponseTypes ArduinoLibreTinyOTABackend::begin(size_t image_size) { +OTAResponseTypes ArduinoLibreTinyOTABackend::begin(size_t image_size, OTAType ota_type) { + if (ota_type != OTA_TYPE_UPDATE_APP) { + return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE; + } // Handle UPDATE_SIZE_UNKNOWN (0) which is used by web server OTA // where the exact firmware size is unknown due to multipart encoding if (image_size == 0) { diff --git a/esphome/components/ota/ota_backend_arduino_libretiny.h b/esphome/components/ota/ota_backend_arduino_libretiny.h index 3d426e6759..c2716a44d1 100644 --- a/esphome/components/ota/ota_backend_arduino_libretiny.h +++ b/esphome/components/ota/ota_backend_arduino_libretiny.h @@ -8,7 +8,7 @@ namespace esphome::ota { class ArduinoLibreTinyOTABackend final { public: - OTAResponseTypes begin(size_t image_size); + OTAResponseTypes begin(size_t image_size, OTAType ota_type = OTA_TYPE_UPDATE_APP); void set_update_md5(const char *md5); OTAResponseTypes write(uint8_t *data, size_t len); OTAResponseTypes end(); diff --git a/esphome/components/ota/ota_backend_arduino_rp2040.cpp b/esphome/components/ota/ota_backend_arduino_rp2040.cpp index bc8ef812e6..0ca0602519 100644 --- a/esphome/components/ota/ota_backend_arduino_rp2040.cpp +++ b/esphome/components/ota/ota_backend_arduino_rp2040.cpp @@ -15,7 +15,10 @@ static const char *const TAG = "ota.arduino_rp2040"; std::unique_ptr make_ota_backend() { return make_unique(); } -OTAResponseTypes ArduinoRP2040OTABackend::begin(size_t image_size) { +OTAResponseTypes ArduinoRP2040OTABackend::begin(size_t image_size, OTAType ota_type) { + if (ota_type != OTA_TYPE_UPDATE_APP) { + return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE; + } // OTA size of 0 is not currently handled, but // web_server is not supported for RP2040, so this is not an issue. bool ret = Update.begin(image_size, U_FLASH); diff --git a/esphome/components/ota/ota_backend_arduino_rp2040.h b/esphome/components/ota/ota_backend_arduino_rp2040.h index 05bd2f5cc4..d04d5c1a84 100644 --- a/esphome/components/ota/ota_backend_arduino_rp2040.h +++ b/esphome/components/ota/ota_backend_arduino_rp2040.h @@ -10,7 +10,7 @@ namespace esphome::ota { class ArduinoRP2040OTABackend final { public: - OTAResponseTypes begin(size_t image_size); + OTAResponseTypes begin(size_t image_size, OTAType ota_type = OTA_TYPE_UPDATE_APP); void set_update_md5(const char *md5); OTAResponseTypes write(uint8_t *data, size_t len); OTAResponseTypes end(); diff --git a/esphome/components/ota/ota_backend_esp8266.cpp b/esphome/components/ota/ota_backend_esp8266.cpp index 7c9d392532..6a678fb419 100644 --- a/esphome/components/ota/ota_backend_esp8266.cpp +++ b/esphome/components/ota/ota_backend_esp8266.cpp @@ -50,7 +50,10 @@ static const char *const TAG = "ota.esp8266"; std::unique_ptr make_ota_backend() { return make_unique(); } -OTAResponseTypes ESP8266OTABackend::begin(size_t image_size) { +OTAResponseTypes ESP8266OTABackend::begin(size_t image_size, OTAType ota_type) { + if (ota_type != OTA_TYPE_UPDATE_APP) { + return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE; + } // Handle UPDATE_SIZE_UNKNOWN (0) by calculating available space if (image_size == 0) { // Round down to sector boundary: subtract one sector, then mask to sector alignment diff --git a/esphome/components/ota/ota_backend_esp8266.h b/esphome/components/ota/ota_backend_esp8266.h index b364e216a3..21b5c12c2d 100644 --- a/esphome/components/ota/ota_backend_esp8266.h +++ b/esphome/components/ota/ota_backend_esp8266.h @@ -14,7 +14,7 @@ namespace esphome::ota { /// by not having a global Update object in .bss. class ESP8266OTABackend final { public: - OTAResponseTypes begin(size_t image_size); + OTAResponseTypes begin(size_t image_size, OTAType ota_type = OTA_TYPE_UPDATE_APP); void set_update_md5(const char *md5); OTAResponseTypes write(uint8_t *data, size_t len); OTAResponseTypes end(); diff --git a/esphome/components/ota/ota_backend_esp_idf.cpp b/esphome/components/ota/ota_backend_esp_idf.cpp index b4b38a192f..42d106bf1f 100644 --- a/esphome/components/ota/ota_backend_esp_idf.cpp +++ b/esphome/components/ota/ota_backend_esp_idf.cpp @@ -16,7 +16,30 @@ static const char *const TAG = "ota.idf"; std::unique_ptr make_ota_backend() { return make_unique(); } -OTAResponseTypes IDFOTABackend::begin(size_t image_size) { +OTAResponseTypes IDFOTABackend::begin(size_t image_size, ota::OTAType ota_type) { +#ifdef USE_OTA_PARTITIONS + this->ota_type_ = ota_type; + if (this->ota_type_ == ota::OTA_TYPE_UPDATE_PARTITION_TABLE) { + // Reject any size other than ESP_PARTITION_TABLE_MAX_LEN: under- leaves stale bytes from the + // previous table; over- can't fit the reserved region. + if (image_size != ESP_PARTITION_TABLE_MAX_LEN) { + ESP_LOGE(TAG, "Wrong partition table size: expected %u bytes, got %zu", ESP_PARTITION_TABLE_MAX_LEN, image_size); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + memset(this->buf_, 0xFF, sizeof this->buf_); + this->buf_written_ = 0; + this->image_size_ = image_size; + this->md5_.init(); + return OTA_RESPONSE_OK; + } + if (this->ota_type_ != ota::OTA_TYPE_UPDATE_APP) { + return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE; + } +#else + if (ota_type != ota::OTA_TYPE_UPDATE_APP) { + return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE; + } +#endif #ifdef USE_OTA_ROLLBACK // If we're starting an OTA, the current boot is good enough - mark it valid // to prevent rollback and allow the OTA to proceed even if the safe mode @@ -52,6 +75,21 @@ void IDFOTABackend::set_update_md5(const char *expected_md5) { } OTAResponseTypes IDFOTABackend::write(uint8_t *data, size_t len) { +#ifdef USE_OTA_PARTITIONS + if (this->ota_type_ == ota::OTA_TYPE_UPDATE_PARTITION_TABLE) { + if (len > PARTITION_TABLE_BUFFER_SIZE - this->buf_written_) { + ESP_LOGE(TAG, "Wrong partition table size"); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + memcpy(this->buf_ + this->buf_written_, data, len); + this->buf_written_ += len; + this->md5_.add(data, len); + return OTA_RESPONSE_OK; + } + if (this->ota_type_ != ota::OTA_TYPE_UPDATE_APP) { + return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE; + } +#endif esp_err_t err = esp_ota_write(this->update_handle_, data, len); this->md5_.add(data, len); if (err != ESP_OK) { @@ -73,6 +111,14 @@ OTAResponseTypes IDFOTABackend::end() { return OTA_RESPONSE_ERROR_MD5_MISMATCH; } } +#ifdef USE_OTA_PARTITIONS + if (this->ota_type_ == ota::OTA_TYPE_UPDATE_PARTITION_TABLE) { + return this->update_partition_table(); + } + if (this->ota_type_ != ota::OTA_TYPE_UPDATE_APP) { + return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE; + } +#endif esp_err_t err = esp_ota_end(this->update_handle_); this->update_handle_ = 0; if (err == ESP_OK) { @@ -96,6 +142,14 @@ OTAResponseTypes IDFOTABackend::end() { } void IDFOTABackend::abort() { +#ifdef USE_OTA_PARTITIONS + if (this->partition_table_part_ != nullptr) { + esp_partition_deregister_external(this->partition_table_part_); + this->partition_table_part_ = nullptr; + } +#endif + // esp_ota_abort with handle 0 returns ESP_ERR_INVALID_ARG harmlessly, so this is safe whether + // or not an update is in flight. esp_ota_abort(this->update_handle_); this->update_handle_ = 0; } diff --git a/esphome/components/ota/ota_backend_esp_idf.h b/esphome/components/ota/ota_backend_esp_idf.h index d007bcd128..54fdd24f93 100644 --- a/esphome/components/ota/ota_backend_esp_idf.h +++ b/esphome/components/ota/ota_backend_esp_idf.h @@ -9,21 +9,54 @@ namespace esphome::ota { +#ifdef USE_OTA_PARTITIONS +// Staging buffer holds the entire partition table for verification before any flash op. +static constexpr size_t PARTITION_TABLE_BUFFER_SIZE = ESP_PARTITION_TABLE_MAX_LEN; // 0xC00 + +void get_running_app_position(uint32_t &offset, size_t &size); +#endif + class IDFOTABackend final { public: - OTAResponseTypes begin(size_t image_size); + OTAResponseTypes begin(size_t image_size, ota::OTAType ota_type = ota::OTA_TYPE_UPDATE_APP); void set_update_md5(const char *md5); OTAResponseTypes write(uint8_t *data, size_t len); OTAResponseTypes end(); void abort(); bool supports_compression() { return false; } + protected: +#ifdef USE_OTA_PARTITIONS + // copy_dest_part non-null means the running app must be copied INTO this slot of the current + // table before the new partition table is committed. The destination is in the current table + // because that's where esp_partition_copy can write; once the new table replaces it, the same + // flash region becomes target_app_index in the new table. + struct PartitionTablePlan { + int target_app_index{-1}; + const esp_partition_t *copy_dest_part{nullptr}; + }; + + OTAResponseTypes validate_new_partition_table_(uint32_t running_app_offset, size_t running_app_size, + PartitionTablePlan &plan); + OTAResponseTypes update_partition_table(); +#endif + private: esp_ota_handle_t update_handle_{0}; const esp_partition_t *partition_; md5::MD5Digest md5_{}; char expected_bin_md5_[32]; bool md5_set_{false}; +#ifdef USE_OTA_PARTITIONS + // Buffer first so it packs tightly after the preceding `bool md5_set_` with no alignment + // padding. Only resident during an active OTA: the backend is constructed per connection and + // destroyed on cleanup_connection_(). + uint8_t buf_[PARTITION_TABLE_BUFFER_SIZE]; + size_t buf_written_{0}; + size_t image_size_{0}; + const esp_partition_t *partition_table_part_{nullptr}; + ota::OTAType ota_type_{ota::OTA_TYPE_UPDATE_APP}; +#endif }; std::unique_ptr make_ota_backend(); diff --git a/esphome/components/ota/ota_backend_host.cpp b/esphome/components/ota/ota_backend_host.cpp index 2e2132418d..a2c9f2cc33 100644 --- a/esphome/components/ota/ota_backend_host.cpp +++ b/esphome/components/ota/ota_backend_host.cpp @@ -10,7 +10,9 @@ namespace esphome::ota { std::unique_ptr make_ota_backend() { return make_unique(); } -OTAResponseTypes HostOTABackend::begin(size_t image_size) { return OTA_RESPONSE_ERROR_UPDATE_PREPARE; } +OTAResponseTypes HostOTABackend::begin(size_t image_size, OTAType ota_type) { + return OTA_RESPONSE_ERROR_UPDATE_PREPARE; +} void HostOTABackend::set_update_md5(const char *expected_md5) {} diff --git a/esphome/components/ota/ota_backend_host.h b/esphome/components/ota/ota_backend_host.h index 300facf72f..4451fdfe18 100644 --- a/esphome/components/ota/ota_backend_host.h +++ b/esphome/components/ota/ota_backend_host.h @@ -9,7 +9,7 @@ namespace esphome::ota { /// OTA triggers to compile for host platform during development. class HostOTABackend final { public: - OTAResponseTypes begin(size_t image_size); + OTAResponseTypes begin(size_t image_size, OTAType ota_type = OTA_TYPE_UPDATE_APP); void set_update_md5(const char *md5); OTAResponseTypes write(uint8_t *data, size_t len); OTAResponseTypes end(); diff --git a/esphome/components/ota/ota_partitions_esp_idf.cpp b/esphome/components/ota/ota_partitions_esp_idf.cpp new file mode 100644 index 0000000000..2a2ed577f1 --- /dev/null +++ b/esphome/components/ota/ota_partitions_esp_idf.cpp @@ -0,0 +1,327 @@ +#ifdef USE_ESP32 +#include "ota_backend_esp_idf.h" + +#include "esphome/core/defines.h" + +#ifdef USE_OTA_PARTITIONS +#include "esphome/components/watchdog/watchdog.h" +#include "esphome/core/log.h" + +#include +#include +#include + +#include + +namespace esphome::ota { + +static const char *const TAG = "ota.idf"; + +static inline bool check_overlap(uint32_t a_offset, size_t a_size, uint32_t b_offset, size_t b_size) { + return (a_offset + a_size > b_offset && b_offset + b_size > a_offset); +} + +// Wraps esp_partition_find/_get/_next/_release. Returns nullptr if no APP partition at `address` +// is at least `min_size` bytes. +static const esp_partition_t *find_app_partition_at(uint32_t address, size_t min_size) { + const esp_partition_t *found = nullptr; + esp_partition_iterator_t it = esp_partition_find(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_ANY, nullptr); + while (it != nullptr) { + const esp_partition_t *p = esp_partition_get(it); + if (p->address == address && p->size >= min_size) { + found = p; + break; + } + it = esp_partition_next(it); + } + esp_partition_iterator_release(it); + return found; +} + +// Validates the staged partition table and picks the post-update boot slot. All non-destructive +// checks live here; the destructive write is in update_partition_table(). +// Side effect: registers the live partition-table region as partition_table_part_ so the caller +// can write to it; abort() releases it on error. +OTAResponseTypes IDFOTABackend::validate_new_partition_table_(uint32_t running_app_offset, size_t running_app_size, + PartitionTablePlan &plan) { + esp_err_t err = esp_partition_register_external( + nullptr, ESP_PRIMARY_PARTITION_TABLE_OFFSET, ESP_PARTITION_TABLE_SIZE, "PrimaryPrtTable", + ESP_PARTITION_TYPE_PARTITION_TABLE, ESP_PARTITION_SUBTYPE_PARTITION_TABLE_PRIMARY, &this->partition_table_part_); + if (err != ESP_OK) { + ESP_LOGE(TAG, "esp_partition_register_external failed (err=0x%X)", err); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + + int num_partitions = 0; + const esp_partition_info_t *existing_partition_table = nullptr; + esp_partition_mmap_handle_t partition_table_map; + err = esp_partition_mmap(this->partition_table_part_, 0, ESP_PARTITION_TABLE_MAX_LEN, ESP_PARTITION_MMAP_DATA, + reinterpret_cast(&existing_partition_table), &partition_table_map); + if (err != ESP_OK) { + ESP_LOGE(TAG, "esp_partition_mmap failed (err=0x%X)", err); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + err = esp_partition_table_verify(existing_partition_table, true, &num_partitions); + esp_partition_munmap(partition_table_map); + if (err != ESP_OK) { + ESP_LOGE(TAG, "esp_partition_table_verify failed (existing partition table) (err=0x%X)", err); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + + const esp_partition_info_t *new_partition_table = reinterpret_cast(this->buf_); + err = esp_partition_table_verify(new_partition_table, true, &num_partitions); + if (err != ESP_OK) { + ESP_LOGE(TAG, "esp_partition_table_verify failed (new partition table) (err=0x%X)", err); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + + // esp_partition_table_verify does not catch a missing MD5 entry, but the bootloader refuses + // to boot from a table without one. + bool checksum_found = false; + for (size_t i = 0; i < ESP_PARTITION_TABLE_MAX_ENTRIES; i++) { + if (new_partition_table[i].magic == ESP_PARTITION_MAGIC_MD5) { + checksum_found = true; + break; + } + } + if (!checksum_found) { + ESP_LOGE(TAG, "New partition table has no checksum"); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + + // Slot-selection policy when multiple slots can host the running app: pick the FIRST eligible + // slot in table order, preferring the no-copy path (matching offset) over the copy path. + // Deterministic and table-ordering-stable. + int app_partitions_found = 0; + int new_app_part_index = -1; + int new_app_part_index_with_copy = -1; + const esp_partition_t *app_copy_dest_part = nullptr; + bool otadata_partition_found = false; + bool otadata_overlap = false; + bool nvs_partition_found = false; + for (int i = 0; i < num_partitions; i++) { + const esp_partition_info_t *new_part = &new_partition_table[i]; + if (new_part->type == ESP_PARTITION_TYPE_APP) { + app_partitions_found++; + if (new_part->pos.size >= running_app_size) { + if (new_part->pos.offset == running_app_offset) { + if (new_app_part_index == -1) { + new_app_part_index = i; + } + } else if (new_app_part_index_with_copy == -1 && + !check_overlap(running_app_offset, running_app_size, new_part->pos.offset, running_app_size)) { + // esp_partition_copy writes into a registered partition; need one at this offset in the + // current table. + const esp_partition_t *p = find_app_partition_at(new_part->pos.offset, running_app_size); + if (p != nullptr) { + new_app_part_index_with_copy = i; + app_copy_dest_part = p; + } + } + } + } else if (new_part->type == ESP_PARTITION_TYPE_DATA) { + if (new_part->subtype == ESP_PARTITION_SUBTYPE_DATA_OTA) { + otadata_partition_found = true; + otadata_overlap = check_overlap(running_app_offset, running_app_size, new_part->pos.offset, new_part->pos.size); + } else if (new_part->subtype == ESP_PARTITION_SUBTYPE_DATA_NVS && + strncmp(reinterpret_cast(new_part->label), "nvs", sizeof(new_part->label)) == 0) { + nvs_partition_found = true; + } + } + } + + if (new_app_part_index == -1 && new_app_part_index_with_copy == -1) { + // Most likely cause: the user picked the wrong migration .bin for their running app's size. + // Rejecting here is non-destructive (no flash op has run yet); the user can safely retry with + // a different .bin. Log enough info that they can pick the right method without guessing. + ESP_LOGE(TAG, + "Running app at 0x%X (%u bytes used) does not fit any compatible slot in the new " + "partition table. Pick a migration method whose size limit is at least %u bytes and " + "retry; no flash content was modified.", + running_app_offset, running_app_size, running_app_size); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + if (app_partitions_found < 2) { + ESP_LOGE(TAG, "New partition table needs at least 2 app partitions, found %d", app_partitions_found); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + if (!otadata_partition_found) { + ESP_LOGE(TAG, "New partition table is missing the required otadata partition"); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + if (!nvs_partition_found) { + ESP_LOGE(TAG, "New partition table is missing the required nvs partition"); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + if (otadata_overlap) { + ESP_LOGE(TAG, + "New otadata partition overlaps with the running app at 0x%X (size %u). The chosen " + "partition table is not compatible with this device's current flash layout; pick a " + "different migration method.", + running_app_offset, running_app_size); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + + if (new_app_part_index != -1) { + plan.target_app_index = new_app_part_index; + plan.copy_dest_part = nullptr; + } else { + plan.target_app_index = new_app_part_index_with_copy; + plan.copy_dest_part = app_copy_dest_part; + } + return OTA_RESPONSE_OK; +} + +OTAResponseTypes IDFOTABackend::update_partition_table() { + if (this->buf_written_ == 0 || this->image_size_ != this->buf_written_) { + ESP_LOGE(TAG, "Not enough data received"); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + + // Without a valid running-app size we cannot compute overlap or copy bounds. zero indicates + // esp_ota_get_running_partition() failed (e.g. cache unloaded by a previous aborted OTA). + uint32_t running_app_offset; + size_t running_app_size; + get_running_app_position(running_app_offset, running_app_size); + if (running_app_size == 0) { + ESP_LOGE(TAG, "Failed to determine running app position"); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_VERIFY; + } + + PartitionTablePlan plan; + OTAResponseTypes validate_result = this->validate_new_partition_table_(running_app_offset, running_app_size, plan); + if (validate_result != OTA_RESPONSE_OK) { + return validate_result; + } + + // ERROR severity so the warning shows up in default log filters; any failure past this point + // can leave the device unbootable until it is recovered with a serial flash. + ESP_LOGE(TAG, "Starting partition table update.\n" + " DO NOT REMOVE POWER until the device reboots successfully.\n" + " Loss of power during this operation may render the device unable to boot until\n" + " it is recovered via a serial flash."); + + // One guard over the whole critical section in case an IDF call takes longer than expected on + // some chip variant. + watchdog::WatchdogManager watchdog(15000); + + esp_err_t err; + const esp_partition_info_t *new_partition_table = reinterpret_cast(this->buf_); + + if (plan.copy_dest_part != nullptr) { + // Resolve the source via running_app_offset rather than esp_ota_get_running_partition() in + // case a prior aborted partition-table OTA called esp_partition_unload_all() in this boot, + // which leaves esp_ota_get_running_partition() returning nullptr. + const esp_partition_t *running_app_part = find_app_partition_at(running_app_offset, running_app_size); + if (running_app_part == nullptr) { + ESP_LOGE(TAG, "Cannot resolve running app partition at offset 0x%X", running_app_offset); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE; + } + ESP_LOGD(TAG, "Copying running app from 0x%X to 0x%X (size: 0x%X)", running_app_part->address, + plan.copy_dest_part->address, running_app_size); + err = esp_partition_copy(plan.copy_dest_part, 0, running_app_part, 0, running_app_size); + if (err != ESP_OK) { + ESP_LOGE(TAG, "esp_partition_copy failed (err=0x%X)", err); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE; + } + } + + // Deinit NVS only just before the first destructive write so verify/copy failure paths return + // with NVS still functional. From this point on, components that hold open NVS handles + // (e.g. preferences) will fail with ESP_ERR_NVS_INVALID_HANDLE on success or failure; + // nvs_flash_init() can re-init the subsystem but cannot revive existing handles. On the + // success path the device reboots immediately afterwards so this doesn't matter; on the + // failure path the user must reboot the device before retrying. + nvs_flash_deinit(); + + // Update the partition table + err = esp_ota_begin(this->partition_table_part_, ESP_PARTITION_TABLE_MAX_LEN, &this->update_handle_); + if (err != ESP_OK) { + esp_ota_abort(this->update_handle_); + this->update_handle_ = 0; + ESP_LOGE(TAG, "esp_ota_begin failed (err=0x%X)", err); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE; + } + err = esp_ota_write(this->update_handle_, this->buf_, ESP_PARTITION_TABLE_MAX_LEN); + if (err != ESP_OK) { + esp_ota_abort(this->update_handle_); + this->update_handle_ = 0; + ESP_LOGE(TAG, "esp_ota_write failed (err=0x%X)", err); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE; + } + err = esp_ota_end(this->update_handle_); + this->update_handle_ = 0; // esp_ota_end releases the handle internally + if (err != ESP_OK) { + ESP_LOGE(TAG, "esp_ota_end failed (err=0x%X)", err); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE; + } + // unload first, then null the member pointer; if abort() ran between the two steps it would + // see a freed pointer. esp_partition_unload_all() invalidates partition_table_part_ too, so + // an explicit deregister would be redundant. + esp_partition_unload_all(); + this->partition_table_part_ = nullptr; + + // Write otadata to set the new boot partition + const esp_partition_info_t *new_part = &new_partition_table[plan.target_app_index]; + const esp_partition_t *new_boot_partition = find_app_partition_at(new_part->pos.offset, 0); + if (new_boot_partition == nullptr) { + ESP_LOGE(TAG, "Selected app partition not found after partition table update"); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE; + } + ESP_LOGD(TAG, "Setting next boot partition to 0x%X", new_boot_partition->address); + err = esp_ota_set_boot_partition(new_boot_partition); + if (err != ESP_OK) { + ESP_LOGE(TAG, "esp_ota_set_boot_partition failed (err=0x%X)", err); + return OTA_RESPONSE_ERROR_PARTITION_TABLE_UPDATE; + } + return OTA_RESPONSE_OK; +} + +// Process-scoped cache. Cannot be a backend member: backends are per-connection but the cache +// must outlive a connection that called esp_partition_unload_all(), after which +// esp_ota_get_running_partition() no longer returns valid data. +static bool s_running_app_initialized = false; +static uint32_t s_running_app_cached_offset = 0; +static size_t s_running_app_cached_size = 0; + +// Flag-gated rather than size==0 so a failed first call doesn't poison the cache. +void get_running_app_position(uint32_t &offset, size_t &size) { + if (!s_running_app_initialized) { + const esp_partition_t *running_app_part = esp_ota_get_running_partition(); + if (running_app_part == nullptr || running_app_part->erase_size == 0) { + // Surface zeros without committing to the cache so a later call has a chance to succeed. + offset = 0; + size = 0; + return; + } + + uint32_t pending_offset = running_app_part->address; + size_t pending_size = running_app_part->size; + + const esp_partition_pos_t running_app_pos = { + .offset = running_app_part->address, + .size = running_app_part->size, + }; + esp_image_metadata_t image_metadata = {}; + image_metadata.start_addr = running_app_part->address; + if (esp_image_verify(ESP_IMAGE_VERIFY_SILENT, &running_app_pos, &image_metadata) == ESP_OK && + image_metadata.image_len < running_app_part->size) { + pending_size = image_metadata.image_len; + } + // Round up to a full flash sector so the copy spans complete erase blocks. + pending_size = ((pending_size + running_app_part->erase_size - 1) / running_app_part->erase_size) * + running_app_part->erase_size; + + s_running_app_cached_offset = pending_offset; + s_running_app_cached_size = pending_size; + s_running_app_initialized = true; + } + + offset = s_running_app_cached_offset; + size = s_running_app_cached_size; +} + +} // namespace esphome::ota + +#endif // USE_OTA_PARTITIONS +#endif // USE_ESP32 diff --git a/esphome/core/__init__.py b/esphome/core/__init__.py index 956d10ad81..05f1038350 100644 --- a/esphome/core/__init__.py +++ b/esphome/core/__init__.py @@ -815,6 +815,17 @@ class EsphomeCore: return self.relative_pioenvs_path(self.name, "firmware.uf2") return self.relative_pioenvs_path(self.name, "firmware.bin") + @property + def partition_table_bin(self) -> Path: + # Native ESP-IDF (--native-idf): the partition table image is emitted under + # build/partition_table/partition-table.bin alongside firmware.bin. PlatformIO writes the + # equivalent file as partitions.bin in the env-specific .pioenvs directory. + if self.data.get(KEY_NATIVE_IDF): + return self.relative_build_path( + "build", "partition_table", "partition-table.bin" + ) + return self.relative_pioenvs_path(self.name, "partitions.bin") + @property def target_platform(self): return self.data[KEY_CORE][KEY_TARGET_PLATFORM] diff --git a/esphome/espota2.py b/esphome/espota2.py index f4c0c73589..a45a6ef234 100644 --- a/esphome/espota2.py +++ b/esphome/espota2.py @@ -16,6 +16,7 @@ from esphome.core import EsphomeError from esphome.helpers import ProgressBar, resolve_ip_address OTA_TYPE_UPDATE_APP = 0x00 +OTA_TYPE_UPDATE_PARTITION_TABLE = 0x01 RESPONSE_OK = 0x00 RESPONSE_REQUEST_AUTH = 0x01 @@ -46,6 +47,8 @@ RESPONSE_ERROR_MD5_MISMATCH = 0x8B RESPONSE_ERROR_RP2040_NOT_ENOUGH_SPACE = 0x8C RESPONSE_ERROR_SIGNATURE_INVALID = 0x8D RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE = 0x8E +RESPONSE_ERROR_PARTITION_TABLE_VERIFY = 0x8F +RESPONSE_ERROR_PARTITION_TABLE_UPDATE = 0x90 RESPONSE_ERROR_UNKNOWN = 0xFF OTA_VERSION_1_0 = 1 @@ -62,7 +65,9 @@ SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS = 0x02 # OTA types this client knows how to send. Future PRs that add bootloader/partition # updates extend this set. Anything outside the set is rejected up front so callers # of perform_ota/run_ota get a clear error instead of a post-auth 0x8E from the device. -_SUPPORTED_OTA_TYPES: frozenset[int] = frozenset({OTA_TYPE_UPDATE_APP}) +_SUPPORTED_OTA_TYPES: frozenset[int] = frozenset( + {OTA_TYPE_UPDATE_APP, OTA_TYPE_UPDATE_PARTITION_TABLE} +) UPLOAD_BLOCK_SIZE = 8192 UPLOAD_BUFFER_SIZE = UPLOAD_BLOCK_SIZE * 8 @@ -128,6 +133,16 @@ _ERROR_MESSAGES: dict[int, str] = { RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE: ( "The requested OTA type is not supported by the device." ), + RESPONSE_ERROR_PARTITION_TABLE_VERIFY: ( + "The partition table update could not be verified. No changes were " + "made to the flash content. Check the logs for more information and retry." + ), + RESPONSE_ERROR_PARTITION_TABLE_UPDATE: ( + "An error occurred while updating the partition table. The device is now " + "in a degraded state (NVS handles are invalid; many components will fail) " + "and may not be able to boot. Check the logs, reboot the device, and " + "retry the update. If the device fails to boot, recover it via a serial flash." + ), RESPONSE_ERROR_UNKNOWN: "Unknown error from ESP", } diff --git a/tests/components/ota/test-partition_access.esp32-idf.yaml b/tests/components/ota/test-partition_access.esp32-idf.yaml new file mode 100644 index 0000000000..0cbf854952 --- /dev/null +++ b/tests/components/ota/test-partition_access.esp32-idf.yaml @@ -0,0 +1,5 @@ +ota: + - platform: esphome + allow_partition_access: true + +<<: !include common.yaml diff --git a/tests/unit_tests/fixtures/partition_tables/esp_idf_hello_world.bin b/tests/unit_tests/fixtures/partition_tables/esp_idf_hello_world.bin new file mode 100644 index 0000000000..b8fa03b4b3 Binary files /dev/null and b/tests/unit_tests/fixtures/partition_tables/esp_idf_hello_world.bin differ diff --git a/tests/unit_tests/fixtures/partition_tables/esphome_dashboard_firmware.bin b/tests/unit_tests/fixtures/partition_tables/esphome_dashboard_firmware.bin new file mode 100644 index 0000000000..e648fa3270 Binary files /dev/null and b/tests/unit_tests/fixtures/partition_tables/esphome_dashboard_firmware.bin differ diff --git a/tests/unit_tests/fixtures/partition_tables/esphome_default.bin b/tests/unit_tests/fixtures/partition_tables/esphome_default.bin new file mode 100644 index 0000000000..d39bf337c9 Binary files /dev/null and b/tests/unit_tests/fixtures/partition_tables/esphome_default.bin differ diff --git a/tests/unit_tests/test_espota2.py b/tests/unit_tests/test_espota2.py index b114f17e6c..2cad1d2ec8 100644 --- a/tests/unit_tests/test_espota2.py +++ b/tests/unit_tests/test_espota2.py @@ -193,6 +193,14 @@ def test_receive_exactly_socket_error(mock_socket: Mock) -> None: espota2.RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE, "Error: The requested OTA type is not supported by the device", ), + ( + espota2.RESPONSE_ERROR_PARTITION_TABLE_VERIFY, + "Error: The partition table update could not be verified", + ), + ( + espota2.RESPONSE_ERROR_PARTITION_TABLE_UPDATE, + "Error: An error occurred while updating the partition table", + ), (espota2.RESPONSE_ERROR_UNKNOWN, "Unknown error from ESP"), ], ) @@ -831,6 +839,66 @@ def test_perform_ota_extended_protocol_app( ) +@pytest.mark.usefixtures("mock_time") +def test_perform_ota_successful_partition_table( + mock_socket: Mock, mock_file: io.BytesIO +) -> None: + """Test OTA partition table update. + + The mocked server advertises both COMPRESSION and PARTITION_ACCESS to exercise + the full extended-protocol negotiation path. Real IDFOTABackend devices return + ``supports_compression() == false`` and never set the COMPRESSION flag for a + partition-table OTA; the flag here is intentional protocol-coverage, not a + description of on-device behaviour. + """ + recv_responses = [ + bytes([espota2.RESPONSE_OK]), # First byte of version response + bytes([espota2.OTA_VERSION_2_0]), # Version number + bytes([espota2.RESPONSE_FEATURE_FLAGS]), # Device supports extended protocol + bytes( + [ + espota2.SERVER_FEATURE_SUPPORTS_COMPRESSION + | espota2.SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS + ] + ), # Device feature flags (compression flag is unrealistic; see docstring) + bytes([espota2.RESPONSE_AUTH_OK]), # No auth required + bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]), # Binary size OK + bytes([espota2.RESPONSE_BIN_MD5_OK]), # MD5 checksum OK + bytes([espota2.RESPONSE_CHUNK_OK]), # Chunk OK + bytes([espota2.RESPONSE_RECEIVE_OK]), # Receive OK + bytes([espota2.RESPONSE_UPDATE_END_OK]), # Update end OK + ] + + mock_socket.recv.side_effect = recv_responses + + espota2.perform_ota( + mock_socket, + "testpass", + mock_file, + "partitions.bin", + espota2.OTA_TYPE_UPDATE_PARTITION_TABLE, + ) + + # Verify magic bytes were sent + assert mock_socket.sendall.call_args_list[0] == call(bytes(espota2.MAGIC_BYTES)) + + # Verify features were sent (compression + SHA256 support + extended protocol) + assert mock_socket.sendall.call_args_list[1] == call( + bytes( + [ + espota2.CLIENT_FEATURE_SUPPORTS_COMPRESSION + | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH + | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL + ] + ) + ) + + # Verify ota type was sent + assert mock_socket.sendall.call_args_list[2] == call( + bytes([espota2.OTA_TYPE_UPDATE_PARTITION_TABLE]) + ) + + @pytest.mark.usefixtures("mock_time") def test_perform_ota_device_rejects_with_unsupported_ota_type( mock_socket: Mock, mock_file: io.BytesIO diff --git a/tests/unit_tests/test_main.py b/tests/unit_tests/test_main.py index 186d8a9573..798a43a4ce 100644 --- a/tests/unit_tests/test_main.py +++ b/tests/unit_tests/test_main.py @@ -24,6 +24,7 @@ from esphome.__main__ import ( _get_configured_xtal_freq, _make_crystal_freq_callback, _resolve_network_devices, + _validate_partition_table_binary, choose_upload_log_host, command_analyze_memory, command_bundle, @@ -83,7 +84,7 @@ from esphome.const import ( PLATFORM_RP2040, ) from esphome.core import CORE, EsphomeError -from esphome.espota2 import OTA_TYPE_UPDATE_APP +from esphome.espota2 import OTA_TYPE_UPDATE_APP, OTA_TYPE_UPDATE_PARTITION_TABLE from esphome.util import BootselResult from esphome.zeroconf import _await_discovery, discover_mdns_devices @@ -1112,6 +1113,7 @@ class MockArgs: reset: bool = False list_only: bool = False output: str | None = None + partition_table: bool = False def test_upload_program_serial_esp32( @@ -1629,6 +1631,237 @@ def test_upload_program_ota_with_file_arg( ) +_PARTITION_TABLE_LEN = 0xC00 + + +def _make_partition_table_bytes() -> bytes: + """Build a minimal partition table image accepted by _validate_partition_table_binary.""" + table = bytearray(b"\xff" * _PARTITION_TABLE_LEN) + # First entry: ESP_PARTITION_MAGIC (0x50AA) little-endian -> bytes 0xAA, 0x50. + table[0] = 0xAA + table[1] = 0x50 + # MD5 checksum entry at offset 32: ESP_PARTITION_MAGIC_MD5 (0xEBEB) little-endian. + table[32] = 0xEB + table[33] = 0xEB + return bytes(table) + + +def test_upload_program_ota_partition_table_with_file_arg( + mock_run_ota: Mock, + mock_get_port_type: Mock, + tmp_path: Path, +) -> None: + """Test upload_program with OTA and partition table.""" + setup_core(platform=PLATFORM_ESP32, tmp_path=tmp_path) + + mock_get_port_type.return_value = "NETWORK" + mock_run_ota.return_value = (0, "192.168.1.100") + + partition_file = tmp_path / "partitions.bin" + partition_file.write_bytes(_make_partition_table_bytes()) + + config = { + CONF_OTA: [ + { + CONF_PLATFORM: CONF_ESPHOME, + CONF_PORT: 3232, + "allow_partition_access": True, + } + ] + } + args = MockArgs(file=str(partition_file), partition_table=True) + devices = ["192.168.1.100"] + + exit_code, host = upload_program(config, args, devices) + + assert exit_code == 0 + assert host == "192.168.1.100" + mock_run_ota.assert_called_once_with( + ["192.168.1.100"], + 3232, + None, + partition_file, + OTA_TYPE_UPDATE_PARTITION_TABLE, + ) + + +def test_upload_program_serial_partition_table( + mock_upload_using_esptool: Mock, + mock_get_port_type: Mock, +) -> None: + """Test serial upload with partition table option (unsupported).""" + setup_core(platform=PLATFORM_ESP32) + mock_get_port_type.return_value = "SERIAL" + mock_upload_using_esptool.return_value = 0 + + config = {} + args = MockArgs(partition_table=True) + devices = ["/dev/ttyUSB0"] + + with pytest.raises( + EsphomeError, + match="The option --partition-table can only be used for Over The Air updates", + ): + upload_program(config, args, devices) + + +def test_upload_program_ota_partition_table_mqttip( + mock_run_ota: Mock, + mock_get_port_type: Mock, + tmp_path: Path, +) -> None: + """--partition-table is allowed for MQTTIP devices; they resolve to a real IP at OTA time.""" + setup_core(platform=PLATFORM_ESP32, tmp_path=tmp_path) + + mock_get_port_type.return_value = "MQTTIP" + mock_run_ota.return_value = (0, "192.168.1.100") + + partition_file = tmp_path / "partitions.bin" + partition_file.write_bytes(_make_partition_table_bytes()) + + config = { + CONF_OTA: [ + { + CONF_PLATFORM: CONF_ESPHOME, + CONF_PORT: 3232, + "allow_partition_access": True, + } + ] + } + args = MockArgs(file=str(partition_file), partition_table=True) + + with patch( + "esphome.__main__._resolve_network_devices", return_value=["192.168.1.100"] + ): + exit_code, host = upload_program(config, args, ["MQTTIP"]) + + assert exit_code == 0 + assert host == "192.168.1.100" + mock_run_ota.assert_called_once_with( + ["192.168.1.100"], + 3232, + None, + partition_file, + OTA_TYPE_UPDATE_PARTITION_TABLE, + ) + + +def test_validate_partition_table_binary_accepts_valid(tmp_path: Path) -> None: + f = tmp_path / "partitions.bin" + f.write_bytes(_make_partition_table_bytes()) + _validate_partition_table_binary(f) + + +_PARTITION_FIXTURE_DIR = Path(__file__).parent / "fixtures" / "partition_tables" + + +@pytest.mark.parametrize( + "fixture", + [ + # Stock ESP-IDF gen_esp32part.py output for an ESPHome build. + "esphome_default.bin", + # ESP-IDF Hello-world example partition table (vendored from espressif/esp-serial-flasher). + "esp_idf_hello_world.bin", + # Partition table shipped with esphome_dashboard's prebuilt firmware. + "esphome_dashboard_firmware.bin", + ], +) +def test_validate_partition_table_binary_accepts_real_binaries(fixture: str) -> None: + """Real-world partition-table binaries from ESP-IDF / ESPHome tooling pass validation.""" + _validate_partition_table_binary(_PARTITION_FIXTURE_DIR / fixture) + + +def test_validate_partition_table_binary_rejects_wrong_size(tmp_path: Path) -> None: + f = tmp_path / "partitions.bin" + f.write_bytes(b"\xaa\x50" + b"\xff" * 100) + with pytest.raises(EsphomeError, match="wrong size"): + _validate_partition_table_binary(f) + + +def test_validate_partition_table_binary_rejects_wrong_magic(tmp_path: Path) -> None: + data = bytearray(_make_partition_table_bytes()) + data[0] = 0x00 + data[1] = 0x00 + f = tmp_path / "partitions.bin" + f.write_bytes(bytes(data)) + with pytest.raises(EsphomeError, match="partition magic"): + _validate_partition_table_binary(f) + + +def test_validate_partition_table_binary_rejects_missing_md5(tmp_path: Path) -> None: + data = bytearray(_make_partition_table_bytes()) + data[32] = 0xFF + data[33] = 0xFF + f = tmp_path / "partitions.bin" + f.write_bytes(bytes(data)) + with pytest.raises(EsphomeError, match="missing the MD5 checksum entry"): + _validate_partition_table_binary(f) + + +def test_validate_partition_table_binary_missing_file(tmp_path: Path) -> None: + with pytest.raises(EsphomeError, match="Cannot read partition table file"): + _validate_partition_table_binary(tmp_path / "does-not-exist.bin") + + +def test_upload_program_ota_partition_table_invalid_file( + mock_run_ota: Mock, + mock_get_port_type: Mock, + tmp_path: Path, +) -> None: + """--partition-table must fail before calling run_ota when the file is not a partition table.""" + setup_core(platform=PLATFORM_ESP32, tmp_path=tmp_path) + + mock_get_port_type.return_value = "NETWORK" + + bad_file = tmp_path / "firmware.bin" + bad_file.write_bytes(b"\x00" * 4096) + + config = { + CONF_OTA: [ + { + CONF_PLATFORM: CONF_ESPHOME, + CONF_PORT: 3232, + "allow_partition_access": True, + } + ] + } + args = MockArgs(file=str(bad_file), partition_table=True) + devices = ["192.168.1.100"] + + with pytest.raises(EsphomeError, match="wrong size"): + upload_program(config, args, devices) + mock_run_ota.assert_not_called() + + +def test_upload_program_ota_partition_table_without_allow_flag( + mock_run_ota: Mock, + mock_get_port_type: Mock, + tmp_path: Path, +) -> None: + """--partition-table must fail fast when allow_partition_access is not enabled in YAML.""" + setup_core(platform=PLATFORM_ESP32, tmp_path=tmp_path) + + mock_get_port_type.return_value = "NETWORK" + + config = { + CONF_OTA: [ + { + CONF_PLATFORM: CONF_ESPHOME, + CONF_PORT: 3232, + } + ] + } + args = MockArgs(file="partitions.bin", partition_table=True) + devices = ["192.168.1.100"] + + with pytest.raises( + EsphomeError, + match="requires 'allow_partition_access: true'", + ): + upload_program(config, args, devices) + mock_run_ota.assert_not_called() + + def test_upload_program_ota_no_config( mock_get_port_type: Mock, ) -> None: