Add partition table update functionality to ota component

esphome/__main__.py

@@ -1009,15 +1009,19 @@ def upload_program(
 
     remote_port = int(ota_conf[CONF_PORT])
     password = ota_conf.get(CONF_PASSWORD)
-    if getattr(args, "file", None) is not None:
-        binary = Path(args.file)
-    else:
-        binary = CORE.firmware_bin
 
     # Resolve MQTT magic strings to actual IP addresses
     network_devices = _resolve_network_devices(devices, config, args)
 
-    return espota2.run_ota(network_devices, remote_port, password, binary)
+    binary = CORE.firmware_bin
+    ota_type = espota2.OTA_TYPE_UPDATE_APP
+    if getattr(args, "partition_table", False):
+        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)
+
+    return espota2.run_ota(network_devices, remote_port, password, binary, ota_type)
 
 
 def show_logs(config: ConfigType, args: ArgsProtocol, devices: list[str]) -> int | None:
@@ -1646,6 +1650,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",
+        action="store_true",
+    )
 
     parser_logs = subparsers.add_parser(
         "logs",

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 config[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)
 
@@ -117,6 +123,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"
@@ -147,6 +154,8 @@ async def to_code(config: ConfigType) -> None:
         cg.add(var.set_auth_password(config[CONF_PASSWORD]))
         cg.add_define("USE_OTA_PASSWORD")
     cg.add_define("USE_OTA_VERSION", config[CONF_VERSION])
+    if config[CONF_ALLOW_PARTITION_ACCESS]:
+        cg.add_define("USE_OTA_PARTITIONS")
 
     await cg.register_component(var, config)
     await ota_to_code(var, config)

esphome/components/esphome/ota/ota_esphome.cpp

@@ -90,8 +90,11 @@ void ESPHomeOTAComponent::loop() {
   }
 }
 
-static const uint8_t FEATURE_SUPPORTS_COMPRESSION = 0x01;
-static const uint8_t FEATURE_SUPPORTS_SHA256_AUTH = 0x02;
+static const uint8_t CLIENT_FEATURE_SUPPORTS_COMPRESSION = 0x01;
+static const uint8_t CLIENT_FEATURE_SUPPORTS_SHA256_AUTH = 0x02;
+static const uint8_t CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL = 0x04;
+static const uint8_t SERVER_FEATURE_SUPPORTS_COMPRESSION = 0x01;
+static const uint8_t SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS = 0x02;
 
 void ESPHomeOTAComponent::handle_handshake_() {
   /// Handle the OTA handshake and authentication.
@@ -177,18 +180,40 @@ void ESPHomeOTAComponent::handle_handshake_() {
       this->ota_features_ = this->handshake_buf_[0];
       ESP_LOGV(TAG, "Features: 0x%02X", this->ota_features_);
       this->transition_ota_state_(OTAState::FEATURE_ACK);
-      this->handshake_buf_[0] =
-          ((this->ota_features_ & FEATURE_SUPPORTS_COMPRESSION) != 0 && this->backend_->supports_compression())
-              ? ota::OTA_RESPONSE_SUPPORTS_COMPRESSION
-              : ota::OTA_RESPONSE_HEADER_OK;
+#ifdef USE_OTA_PARTITIONS
+      this->extended_proto_ = this->ota_features_ & CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL;
+      if (this->extended_proto_) {
+        this->handshake_buf_[0] = ota::OTA_RESPONSE_FEATURE_FLAGS;
+        this->handshake_buf_[1] = 0;
+        if ((this->ota_features_ & CLIENT_FEATURE_SUPPORTS_COMPRESSION) != 0 && this->backend_->supports_compression()) {
+          this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_COMPRESSION;
+        }
+        this->handshake_buf_[1] |= SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS;
+
+      } else {
+#endif
+        this->handshake_buf_[0] =
+            ((this->ota_features_ & CLIENT_FEATURE_SUPPORTS_COMPRESSION) != 0 && this->backend_->supports_compression())
+                ? ota::OTA_RESPONSE_SUPPORTS_COMPRESSION
+                : ota::OTA_RESPONSE_HEADER_OK;
+#ifdef USE_OTA_PARTITIONS
+      }
+#endif
       [[fallthrough]];
     }
 
     case OTAState::FEATURE_ACK: {
       // Acknowledge header - 1 byte
+#ifdef USE_OTA_PARTITIONS
+      if (!this->try_write_(this->extended_proto_ ? 2 : 1, LOG_STR("ack feature"))) {
+        return;
+      }
+#else
       if (!this->try_write_(1, LOG_STR("ack feature"))) {
         return;
       }
+#endif
+
 #ifdef USE_OTA_PASSWORD
       // If password is set, move to auth phase
       if (!this->password_.empty()) {
@@ -271,6 +296,9 @@ void ESPHomeOTAComponent::handle_data_() {
   uint8_t buf[OTA_BUFFER_SIZE];
   char *sbuf = reinterpret_cast<char *>(buf);
   size_t ota_size;
+#ifdef USE_OTA_PARTITIONS
+  ota::OTAType ota_type = ota::OTA_TYPE_UPDATE_APP;
+#endif
 #if USE_OTA_VERSION == 2
   size_t size_acknowledged = 0;
 #endif
@@ -286,6 +314,18 @@ void ESPHomeOTAComponent::handle_data_() {
   // Acknowledge auth OK - 1 byte
   this->write_byte_(ota::OTA_RESPONSE_AUTH_OK);
 
+#ifdef USE_OTA_PARTITIONS
+  if (this->extended_proto_) {
+    // Read ota type, 1 byte
+    if (!this->readall_(buf, 1)) {
+      this->log_read_error_(LOG_STR("OTA type"));
+      goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
+    }
+    ota_type = static_cast<ota::OTAType>(buf[0]);
+  }
+  ESP_LOGV(TAG, "OTA type is 0x%02x", ota_type);
+#endif
+
   // Read size, 4 bytes MSB first
   if (!this->readall_(buf, 4)) {
     this->log_read_error_(LOG_STR("size"));
@@ -306,7 +346,11 @@ void ESPHomeOTAComponent::handle_data_() {
 #endif
 
   // This will block for a few seconds as it locks flash
+#ifdef USE_OTA_PARTITIONS
+  error_code = this->backend_->begin(ota_size, ota_type);
+#else
   error_code = this->backend_->begin(ota_size);
+#endif
   if (error_code != ota::OTA_RESPONSE_OK)
     goto error;  // NOLINT(cppcoreguidelines-avoid-goto)
   update_started = true;
@@ -577,7 +621,7 @@ void ESPHomeOTAComponent::yield_and_feed_watchdog_() {
 void ESPHomeOTAComponent::log_auth_warning_(const LogString *msg) { ESP_LOGW(TAG, "Auth: %s", LOG_STR_ARG(msg)); }
 
 bool ESPHomeOTAComponent::select_auth_type_() {
-  bool client_supports_sha256 = (this->ota_features_ & FEATURE_SUPPORTS_SHA256_AUTH) != 0;
+  bool client_supports_sha256 = (this->ota_features_ & CLIENT_FEATURE_SUPPORTS_SHA256_AUTH) != 0;
 
   // Require SHA256
   if (!client_supports_sha256) {

esphome/components/esphome/ota/ota_esphome.h

@@ -83,6 +83,9 @@ class ESPHomeOTAComponent final : public ota::OTAComponent {
   std::string password_;
   std::unique_ptr<uint8_t[]> auth_buf_;
 #endif  // USE_OTA_PASSWORD
+#ifdef USE_OTA_PARTITIONS
+  bool extended_proto_{false};
+#endif
 
   socket::ListenSocket *server_{nullptr};
   std::unique_ptr<socket::Socket> client_;

esphome/components/ota/ota_backend.h

@@ -23,6 +23,7 @@ enum OTAResponseTypes {
   OTA_RESPONSE_UPDATE_END_OK = 0x45,
   OTA_RESPONSE_SUPPORTS_COMPRESSION = 0x46,
   OTA_RESPONSE_CHUNK_OK = 0x47,
+  OTA_RESPONSE_FEATURE_FLAGS = 0x48,
 
   OTA_RESPONSE_ERROR_MAGIC = 0x80,
   OTA_RESPONSE_ERROR_UPDATE_PREPARE = 0x81,
@@ -38,6 +39,7 @@ enum OTAResponseTypes {
   OTA_RESPONSE_ERROR_MD5_MISMATCH = 0x8B,
   OTA_RESPONSE_ERROR_RP2040_NOT_ENOUGH_SPACE = 0x8C,
   OTA_RESPONSE_ERROR_SIGNATURE_INVALID = 0x8D,
+  OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE = 0x8E,
   OTA_RESPONSE_ERROR_UNKNOWN = 0xFF,
 };
 
@@ -49,6 +51,11 @@ enum OTAState {
   OTA_ERROR,
 };
 
+enum OTAType {
+  OTA_TYPE_UPDATE_APP = 0x00,
+  OTA_TYPE_UPDATE_PARTITION_TABLE = 0x01,
+};
+
 /** Listener interface for OTA state changes.
  *
  * Components can implement this interface to receive OTA state updates

esphome/components/ota/ota_backend_esp_idf.cpp

@@ -9,60 +9,80 @@
 #include <esp_task_wdt.h>
 #include <spi_flash_mmap.h>
 
+#ifdef USE_OTA_PARTITIONS
+#include <esp_image_format.h>
+#endif
+
 namespace esphome::ota {
 
 static const char *const TAG = "ota.idf";
 
 std::unique_ptr<IDFOTABackend> make_ota_backend() { return make_unique<IDFOTABackend>(); }
 
-OTAResponseTypes IDFOTABackend::begin(size_t image_size) {
+OTAResponseTypes IDFOTABackend::begin(size_t image_size, ota::OTAType ota_type) {
+  this->ota_type_ = ota_type;
+  if (this->ota_type_ == ota::OTA_TYPE_UPDATE_APP) {
 #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
-  // timer hasn't expired yet.
-  esp_ota_mark_app_valid_cancel_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
+    // timer hasn't expired yet.
+    esp_ota_mark_app_valid_cancel_rollback();
 #endif
 
-  this->partition_ = esp_ota_get_next_update_partition(nullptr);
-  if (this->partition_ == nullptr) {
-    return OTA_RESPONSE_ERROR_NO_UPDATE_PARTITION;
-  }
+    this->partition_ = esp_ota_get_next_update_partition(nullptr);
+    if (this->partition_ == nullptr) {
+      return OTA_RESPONSE_ERROR_NO_UPDATE_PARTITION;
+    }
 
 #if CONFIG_ESP_TASK_WDT_TIMEOUT_S < 15
-  // The following function takes longer than the 5 seconds timeout of WDT
-  esp_task_wdt_config_t wdtc;
-  wdtc.idle_core_mask = 0;
+    // The following function takes longer than the 5 seconds timeout of WDT
+    esp_task_wdt_config_t wdtc;
+    wdtc.idle_core_mask = 0;
 #if CONFIG_ESP_TASK_WDT_CHECK_IDLE_TASK_CPU0
-  wdtc.idle_core_mask |= (1 << 0);
+    wdtc.idle_core_mask |= (1 << 0);
 #endif
 #if CONFIG_ESP_TASK_WDT_CHECK_IDLE_TASK_CPU1
-  wdtc.idle_core_mask |= (1 << 1);
+    wdtc.idle_core_mask |= (1 << 1);
 #endif
-  wdtc.timeout_ms = 15000;
-  wdtc.trigger_panic = false;
-  esp_task_wdt_reconfigure(&wdtc);
+    wdtc.timeout_ms = 15000;
+    wdtc.trigger_panic = false;
+    esp_task_wdt_reconfigure(&wdtc);
 #endif
 
-  esp_err_t err = esp_ota_begin(this->partition_, image_size, &this->update_handle_);
+    esp_err_t err = esp_ota_begin(this->partition_, image_size, &this->update_handle_);
 
 #if CONFIG_ESP_TASK_WDT_TIMEOUT_S < 15
-  // Set the WDT back to the configured timeout
-  wdtc.timeout_ms = CONFIG_ESP_TASK_WDT_TIMEOUT_S * 1000;
-  esp_task_wdt_reconfigure(&wdtc);
+    // Set the WDT back to the configured timeout
+    wdtc.timeout_ms = CONFIG_ESP_TASK_WDT_TIMEOUT_S * 1000;
+    esp_task_wdt_reconfigure(&wdtc);
 #endif
 
-  if (err != ESP_OK) {
-    esp_ota_abort(this->update_handle_);
-    this->update_handle_ = 0;
-    if (err == ESP_ERR_INVALID_SIZE) {
-      return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
-    } else if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
-      return OTA_RESPONSE_ERROR_WRITING_FLASH;
+    if (err != ESP_OK) {
+      esp_ota_abort(this->update_handle_);
+      this->update_handle_ = 0;
+      if (err == ESP_ERR_INVALID_SIZE) {
+        return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
+      } else if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
+        return OTA_RESPONSE_ERROR_WRITING_FLASH;
+      }
+      return OTA_RESPONSE_ERROR_UNKNOWN;
     }
-    return OTA_RESPONSE_ERROR_UNKNOWN;
+    this->md5_.init();
+    return OTA_RESPONSE_OK;
   }
-  this->md5_.init();
-  return OTA_RESPONSE_OK;
+#ifdef USE_OTA_PARTITIONS
+  if (this->ota_type_ == ota::OTA_TYPE_UPDATE_PARTITION_TABLE) {
+    if (image_size > ESP_PARTITION_TABLE_SIZE || image_size > ESP_PARTITION_TABLE_MAX_LEN || image_size > OTA_BUFFER_SIZE) {
+      return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
+    }
+    memset(this->buf_, 0xFF, sizeof this->buf_);
+    this->buf_written_ = 0;
+    this->image_size_ = image_size;
+    this->md5_.init();
+    return OTA_RESPONSE_OK;
+  }
+#endif
+  return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
 }
 
 void IDFOTABackend::set_update_md5(const char *expected_md5) {
@@ -71,17 +91,31 @@ void IDFOTABackend::set_update_md5(const char *expected_md5) {
 }
 
 OTAResponseTypes IDFOTABackend::write(uint8_t *data, size_t len) {
-  esp_err_t err = esp_ota_write(this->update_handle_, data, len);
-  this->md5_.add(data, len);
-  if (err != ESP_OK) {
-    if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
-      return OTA_RESPONSE_ERROR_MAGIC;
-    } else if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
-      return OTA_RESPONSE_ERROR_WRITING_FLASH;
+  if (this->ota_type_ == ota::OTA_TYPE_UPDATE_APP) {
+    esp_err_t err = esp_ota_write(this->update_handle_, data, len);
+    this->md5_.add(data, len);
+    if (err != ESP_OK) {
+      if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
+        return OTA_RESPONSE_ERROR_MAGIC;
+      } else if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
+        return OTA_RESPONSE_ERROR_WRITING_FLASH;
+      }
+      return OTA_RESPONSE_ERROR_UNKNOWN;
     }
-    return OTA_RESPONSE_ERROR_UNKNOWN;
+    return OTA_RESPONSE_OK;
   }
-  return OTA_RESPONSE_OK;
+#ifdef USE_OTA_PARTITIONS
+  if (this->ota_type_ == ota::OTA_TYPE_UPDATE_PARTITION_TABLE) {
+    if (len > OTA_BUFFER_SIZE - this->buf_written_) {
+      return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
+    }
+    memcpy(this->buf_ + this->buf_written_, data, len);
+    this->buf_written_ += len;
+    this->md5_.add(data, len);
+    return OTA_RESPONSE_OK;
+  }
+#endif
+  return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
 }
 
 OTAResponseTypes IDFOTABackend::end() {
@@ -92,32 +126,176 @@ OTAResponseTypes IDFOTABackend::end() {
       return OTA_RESPONSE_ERROR_MD5_MISMATCH;
     }
   }
-  esp_err_t err = esp_ota_end(this->update_handle_);
-  this->update_handle_ = 0;
-  if (err == ESP_OK) {
-    err = esp_ota_set_boot_partition(this->partition_);
+  if (this->ota_type_ == ota::OTA_TYPE_UPDATE_APP) {
+    esp_err_t err = esp_ota_end(this->update_handle_);
+    this->update_handle_ = 0;
     if (err == ESP_OK) {
-      return OTA_RESPONSE_OK;
+      err = esp_ota_set_boot_partition(this->partition_);
+      if (err == ESP_OK) {
+        return OTA_RESPONSE_OK;
+      }
     }
+    if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
+  #ifdef USE_OTA_SIGNED_VERIFICATION
+      ESP_LOGE(TAG, "OTA validation failed (err=0x%X) - possible signature verification failure", err);
+      return OTA_RESPONSE_ERROR_SIGNATURE_INVALID;
+  #else
+      return OTA_RESPONSE_ERROR_UPDATE_END;
+  #endif
+    }
+    if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
+      return OTA_RESPONSE_ERROR_WRITING_FLASH;
+    }
+    return OTA_RESPONSE_ERROR_UNKNOWN;
+  }
+#ifdef USE_OTA_PARTITIONS
+  if (this->ota_type_ == ota::OTA_TYPE_UPDATE_PARTITION_TABLE) {
+    return this->update_partition_table();
   }
-  if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
-#ifdef USE_OTA_SIGNED_VERIFICATION
-    ESP_LOGE(TAG, "OTA validation failed (err=0x%X) - possible signature verification failure", err);
-    return OTA_RESPONSE_ERROR_SIGNATURE_INVALID;
-#else
-    return OTA_RESPONSE_ERROR_UPDATE_END;
 #endif
-  }
-  if (err == ESP_ERR_FLASH_OP_TIMEOUT || err == ESP_ERR_FLASH_OP_FAIL) {
-    return OTA_RESPONSE_ERROR_WRITING_FLASH;
-  }
-  return OTA_RESPONSE_ERROR_UNKNOWN;
+  return OTA_RESPONSE_ERROR_UNSUPPORTED_OTA_TYPE;
 }
 
 void IDFOTABackend::abort() {
-  esp_ota_abort(this->update_handle_);
-  this->update_handle_ = 0;
+  if (this->ota_type_ == ota::OTA_TYPE_UPDATE_APP) {
+    esp_ota_abort(this->update_handle_);
+    this->update_handle_ = 0;
+  }
+#ifdef USE_OTA_PARTITIONS
+  if (this->partition_table_part_ != nullptr) {
+    esp_partition_deregister_external(this->partition_table_part_);
+    this->partition_table_part_ = nullptr;
+  }
+#endif
 }
 
+#ifdef USE_OTA_PARTITIONS
+OTAResponseTypes IDFOTABackend::update_partition_table() {
+  esp_err_t err;
+  int num_partitions;
+  if (this->buf_written_ == 0 || this->image_size_ != this->buf_written_) {
+    ESP_LOGE(TAG, "not enough data received (%d/%d bytes)", this->buf_written_, this->image_size_);
+    return OTA_RESPONSE_ERROR_UNKNOWN;
+  }
+  ESP_LOGD(TAG, "partition table size %d", this->image_size_);
+
+  // Get running app partition and used size
+  const esp_partition_t *running_app_part = esp_ota_get_running_partition();
+  size_t running_app_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;
+  err = esp_image_verify(ESP_IMAGE_VERIFY_SILENT, &running_app_pos, &image_metadata);
+  if (err == ESP_OK && image_metadata.image_len < running_app_part->size) {
+    running_app_size = image_metadata.image_len;
+  }
+
+  // Get partition table partition
+  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_UNKNOWN;
+  }
+  // Verify existing partition table
+  const esp_partition_info_t *existing_partition_table = NULL;
+  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, (const void**)&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_UNKNOWN;
+  }
+  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_UNKNOWN;
+  }
+
+  // Verify new partition table
+  const esp_partition_info_t *new_partition_table = (const esp_partition_info_t *)this->buf_;
+  // esp_partition_table_verify expects ESP_PARTITION_TABLE_MAX_LEN bytes of data
+  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_UNKNOWN;
+  }
+
+  // Check if the required app and otadata partitions exist in the new partition table
+  // Check which app slot to boot from in the new partition table
+  int app_partitions_found = 0;
+  int app_index = -1;
+  int app_index_with_copy = -1;
+  int otadata_index = -1;
+  bool otadata_no_overlap = false;
+  for (int i = 0; i < num_partitions; i++) {
+    const esp_partition_info_t *part = &new_partition_table[i];
+    if (part->type == ESP_PARTITION_TYPE_APP) {
+      app_partitions_found++;
+      if (part->pos.size >= running_app_size) {
+        if (part->pos.offset == running_app_part->address) {
+          app_index = i;
+        } else if (part->pos.offset >= running_app_part->address + running_app_size || running_app_part->address >= part->pos.offset + part->pos.size) {
+          // No overlap with running app
+          app_index_with_copy = i;
+        }
+      }
+    } else if (part->type == ESP_PARTITION_TYPE_DATA && part->subtype == ESP_PARTITION_SUBTYPE_DATA_OTA) {
+      otadata_index = i;
+      otadata_no_overlap = part->pos.offset >= running_app_part->address + running_app_size || running_app_part->address >= part->pos.offset + part->pos.size;
+    }
+  }
+  if (app_index == -1 && app_index_with_copy == -1) {
+    // Can't move running app to new partition layout
+    ESP_LOGE(TAG, "No compatible app partition found in the new partition table");
+    return OTA_RESPONSE_ERROR_UNKNOWN;
+  }
+  if (app_partitions_found < 2 || otadata_index == -1) {
+    // OTA would be impossible with new partition table
+    ESP_LOGE(TAG, "New partition table is missing the required partitions for OTA");
+    return OTA_RESPONSE_ERROR_UNKNOWN;
+  }
+  if (!otadata_no_overlap) {
+    // Can't write to new otadata partition because it overlaps with the running app
+    ESP_LOGE(TAG, "New otadata partition overlaps with running app");
+    return OTA_RESPONSE_ERROR_UNKNOWN;
+  }
+
+  ESP_LOGD(TAG, "Checks passed, starting partition table update", err);
+
+  // TODO: Copy the running app partition to new position if needed
+  if (app_index == -1) {
+    ESP_LOGE(TAG, "Moving the app partition is required but not implemented");
+    return OTA_RESPONSE_ERROR_UNKNOWN;
+  }
+
+  // Update the partition table
+  err = esp_ota_begin(this->partition_table_part_, this->image_size_, &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_UNKNOWN;
+  }
+  err = esp_ota_write(this->update_handle_, this->buf_, this->image_size_);
+  if (err != ESP_OK) {
+    ESP_LOGE(TAG, "esp_ota_write failed (err=0x%X) ", err);
+    return OTA_RESPONSE_ERROR_UNKNOWN;
+  }
+  err = esp_ota_end(this->update_handle_);
+  this->update_handle_ = 0;
+  if (err != ESP_OK) {
+    ESP_LOGE(TAG, "esp_ota_end failed (err=0x%X) ", err);
+    return OTA_RESPONSE_ERROR_UNKNOWN;
+  }
+
+  // TODO: Reload partition table and rewrite otadata
+
+  return OTA_RESPONSE_OK;
+}
+#endif
+
 }  // namespace esphome::ota
 #endif  // USE_ESP32

esphome/components/ota/ota_backend_esp_idf.h

@@ -9,21 +9,37 @@
 
 namespace esphome::ota {
 
+#ifdef USE_OTA_PARTITIONS
+static constexpr size_t OTA_BUFFER_SIZE = ESP_PARTITION_TABLE_MAX_LEN; // 0xC00
+#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
+  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};
+  ota::OTAType ota_type_{ota::OTA_TYPE_UPDATE_APP};
+#ifdef USE_OTA_PARTITIONS
+  uint8_t buf_[OTA_BUFFER_SIZE];
+  size_t buf_written_{0};
+  size_t image_size_{0};
+  const esp_partition_t *partition_table_part_{nullptr};
+#endif
 };
 
 std::unique_ptr<IDFOTABackend> make_ota_backend();

esphome/core/__init__.py

@@ -775,6 +775,10 @@ 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):
+        return self.relative_pioenvs_path(self.name, "partitions.bin")
+
     @property
     def target_platform(self):
         return self.data[KEY_CORE][KEY_TARGET_PLATFORM]

esphome/espota2.py

@@ -15,6 +15,9 @@ from typing import Any
 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
 RESPONSE_REQUEST_SHA256_AUTH = 0x02
@@ -27,6 +30,7 @@ RESPONSE_RECEIVE_OK = 0x44
 RESPONSE_UPDATE_END_OK = 0x45
 RESPONSE_SUPPORTS_COMPRESSION = 0x46
 RESPONSE_CHUNK_OK = 0x47
+RESPONSE_FEATURE_FLAGS = 0x48
 
 RESPONSE_ERROR_MAGIC = 0x80
 RESPONSE_ERROR_UPDATE_PREPARE = 0x81
@@ -49,9 +53,11 @@ OTA_VERSION_2_0 = 2
 
 MAGIC_BYTES = [0x6C, 0x26, 0xF7, 0x5C, 0x45]
 
-FEATURE_SUPPORTS_COMPRESSION = 0x01
-FEATURE_SUPPORTS_SHA256_AUTH = 0x02
-
+CLIENT_FEATURE_SUPPORTS_COMPRESSION = 0x01
+CLIENT_FEATURE_SUPPORTS_SHA256_AUTH = 0x02
+CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL = 0x04
+SERVER_FEATURE_SUPPORTS_COMPRESSION = 0x01
+SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS = 0x02
 
 UPLOAD_BLOCK_SIZE = 8192
 UPLOAD_BUFFER_SIZE = UPLOAD_BLOCK_SIZE * 8
@@ -232,7 +238,11 @@ def send_check(
 
 
 def perform_ota(
-    sock: socket.socket, password: str | None, file_handle: io.IOBase, filename: Path
+    sock: socket.socket,
+    password: str | None,
+    file_handle: io.IOBase,
+    filename: Path,
+    ota_type: int,
 ) -> None:
     file_contents = file_handle.read()
     file_size = len(file_contents)
@@ -251,7 +261,11 @@ def perform_ota(
         )
 
     # Features - send both compression and SHA256 auth support
-    features_to_send = FEATURE_SUPPORTS_COMPRESSION | FEATURE_SUPPORTS_SHA256_AUTH
+    features_to_send = (
+        CLIENT_FEATURE_SUPPORTS_COMPRESSION
+        | CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
+        | CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
+    )
     send_check(sock, features_to_send, "features")
     features = receive_exactly(
         sock,
@@ -260,7 +274,26 @@ def perform_ota(
         None,  # Accept any response
     )[0]
 
-    if features == RESPONSE_SUPPORTS_COMPRESSION:
+    extended_proto = False
+    if features == RESPONSE_FEATURE_FLAGS:
+        extended_proto = True
+        features = receive_exactly(
+            sock,
+            1,
+            "feature flags",
+            None,  # Accept any response
+        )[0]
+    elif features == RESPONSE_SUPPORTS_COMPRESSION:
+        features = SERVER_FEATURE_SUPPORTS_COMPRESSION
+    else:
+        features = 0
+
+    if ota_type != 0 and not features & SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS:
+        raise OTAError(
+            f"Device only supports app updates"
+        )
+
+    if features & SERVER_FEATURE_SUPPORTS_COMPRESSION:
         upload_contents = gzip.compress(file_contents, compresslevel=9)
         _LOGGER.info("Compressed to %s bytes", len(upload_contents))
     else:
@@ -315,6 +348,9 @@ def perform_ota(
     # Timeout must match device-side OTA_SOCKET_TIMEOUT_DATA to prevent premature failures
     sock.settimeout(90.0)
 
+    if extended_proto:
+        send_check(sock, ota_type, "ota type")
+
     upload_size = len(upload_contents)
     upload_size_encoded = [
         (upload_size >> 24) & 0xFF,
@@ -375,7 +411,11 @@ def perform_ota(
 
 
 def run_ota_impl_(
-    remote_host: str | list[str], remote_port: int, password: str | None, filename: Path
+    remote_host: str | list[str],
+    remote_port: int,
+    password: str | None,
+    filename: Path,
+    ota_type: int,
 ) -> tuple[int, str | None]:
     from esphome.core import CORE
 
@@ -413,7 +453,7 @@ def run_ota_impl_(
         _LOGGER.info("Connected to %s", sa[0])
         with open(filename, "rb") as file_handle:
             try:
-                perform_ota(sock, password, file_handle, filename)
+                perform_ota(sock, password, file_handle, filename, ota_type)
             except OTAError as err:
                 _LOGGER.error(str(err))
                 return 1, None
@@ -428,10 +468,14 @@ def run_ota_impl_(
 
 
 def run_ota(
-    remote_host: str | list[str], remote_port: int, password: str | None, filename: Path
+    remote_host: str | list[str],
+    remote_port: int,
+    password: str | None,
+    filename: Path,
+    ota_type: int,
 ) -> tuple[int, str | None]:
     try:
-        return run_ota_impl_(remote_host, remote_port, password, filename)
+        return run_ota_impl_(remote_host, remote_port, password, filename, ota_type)
     except OTAError as err:
         _LOGGER.error(err)
         return 1, None