Prefix erase-ahead macro with USE_ and test the erase target arithmetic

esphome/components/ota/ota_backend.h

@@ -66,6 +66,19 @@ enum OTAResponseTypes {
  */
 bool version_is_older(const char *candidate, const char *reference);
 
+// 64 KiB flash block; the erase granularity the ESP-IDF backend erases ahead with.
+static constexpr size_t OTA_BLOCK_ERASE_SIZE = 64 * 1024;
+
+/** Target erased watermark for lazy block erase-ahead.
+ *
+ * Rounds the write end offset up to a block boundary, clamped to the partition
+ * size. Platform-independent so the arithmetic is host-testable.
+ */
+constexpr size_t next_erase_end(size_t write_end, size_t partition_size) {
+  const size_t rounded = (write_end + OTA_BLOCK_ERASE_SIZE - 1) & ~(OTA_BLOCK_ERASE_SIZE - 1);
+  return rounded < partition_size ? rounded : partition_size;
+}
+
 enum OTAState {
   OTA_COMPLETED = 0,
   OTA_STARTED,

esphome/components/ota/ota_backend_esp_idf.cpp

@@ -9,8 +9,6 @@
 #include <esp_ota_ops.h>
 #include <sdkconfig.h>
 #include <spi_flash_mmap.h>
-
-#include <algorithm>
 #ifdef USE_OTA_DOWNGRADE_PROTECTION
 #include <esp_app_desc.h>
 #endif
@@ -70,7 +68,7 @@ OTAResponseTypes IDFOTABackend::begin(size_t image_size, ota::OTAType ota_type)
   }
   this->written_ = 0;
   esp_err_t err;
-#ifdef OTA_BLOCK_ERASE_AHEAD
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
   this->erased_end_ = 0;
 #ifdef CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE
   // esp_ota_begin() refuses to start while the running app is unconfirmed;
@@ -146,7 +144,7 @@ OTAResponseTypes IDFOTABackend::write(uint8_t *data, size_t len) {
   if (this->written_ + len > this->partition_->size) {
     return OTA_RESPONSE_ERROR_ESP32_NOT_ENOUGH_SPACE;
   }
-#ifdef OTA_BLOCK_ERASE_AHEAD
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
   OTAResponseTypes erase_result = this->erase_ahead_(len);
   if (erase_result != OTA_RESPONSE_OK) {
     return erase_result;
@@ -170,17 +168,15 @@ OTAResponseTypes IDFOTABackend::write(uint8_t *data, size_t len) {
   return OTA_RESPONSE_OK;
 }
 
-#ifdef OTA_BLOCK_ERASE_AHEAD
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
 OTAResponseTypes IDFOTABackend::erase_ahead_(size_t len) {
-  static constexpr size_t BLOCK_ERASE_SIZE = 64 * 1024;
   const size_t end = this->written_ + len;
   if (this->erased_end_ >= end) {
     return OTA_RESPONSE_OK;
   }
   // Round up to a block boundary, clamped to the partition end; IDF splits the
   // range into 64 KiB block erases where aligned, sector erases elsewhere.
-  const size_t erase_to =
-      std::min<size_t>(this->partition_->size, (end + BLOCK_ERASE_SIZE - 1) & ~(BLOCK_ERASE_SIZE - 1));
+  const size_t erase_to = next_erase_end(end, this->partition_->size);
   // A block erase is one uninterruptible flash op (typically ~150 ms, seconds
   // on aged flash) and the transfer loop may not have fed the WDT for ~1s.
   watchdog::WatchdogManager watchdog(15000);
@@ -286,7 +282,7 @@ void IDFOTABackend::abort() {
   esp_ota_abort(this->update_handle_);
   this->update_handle_ = 0;
   this->written_ = 0;
-#ifdef OTA_BLOCK_ERASE_AHEAD
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
   this->erased_end_ = 0;
 #endif
 }

esphome/components/ota/ota_backend_esp_idf.h

@@ -14,7 +14,7 @@
 // used as fallback on older IDF).
 #if ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(5, 4, 2) || \
     (ESP_IDF_VERSION >= ESP_IDF_VERSION_VAL(5, 3, 3) && ESP_IDF_VERSION < ESP_IDF_VERSION_VAL(5, 4, 0))
-#define OTA_BLOCK_ERASE_AHEAD
+#define USE_OTA_BLOCK_ERASE_AHEAD
 #endif
 
 namespace esphome::ota {
@@ -64,7 +64,7 @@ class IDFOTABackend final {
 #endif
 
  private:
-#ifdef OTA_BLOCK_ERASE_AHEAD
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
   OTAResponseTypes erase_ahead_(size_t len);
 #endif
 #ifdef USE_OTA_SIGNED_VERIFICATION_MULTI_KEY
@@ -77,7 +77,7 @@ class IDFOTABackend final {
   esp_ota_handle_t update_handle_{0};
   const esp_partition_t *partition_{nullptr};
   size_t written_{0};  // Bytes handed to esp_ota_write()
-#ifdef OTA_BLOCK_ERASE_AHEAD
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
   size_t erased_end_{0};  // Erased up to this partition offset; must stay >= written_
 #endif
   char expected_bin_md5_[32];

esphome/components/ota/ota_bootloader_esp_idf.cpp

@@ -75,7 +75,7 @@ OTAResponseTypes IDFOTABackend::setup_bootloader_staging_() {
     ESP_LOGW(TAG, "esp_partition_erase_range failed (err=0x%X)", err);
     // No critical error, don't return
   }
-#ifdef OTA_BLOCK_ERASE_AHEAD
+#ifdef USE_OTA_BLOCK_ERASE_AHEAD
   if (err == ESP_OK) {
     // Skip re-erasing the pre-erased staging region in erase_ahead_()
     this->erased_end_ = this->bootloader_part_->size;

tests/components/ota/test_erase_ahead.cpp

@@ -0,0 +1,41 @@
+// Pins the lazy erase-ahead arithmetic used by the ESP-IDF OTA backend: the
+// erased watermark must always cover the write end, stay 64 KiB block-aligned
+// until the clamp, and never exceed the partition.
+
+#include <gtest/gtest.h>
+
+#include "esphome/components/ota/ota_backend.h"
+
+namespace esphome::ota::testing {
+
+static constexpr size_t BLOCK = 64 * 1024;
+static constexpr size_t PART = 1835008;  // 0x1C0000, a real app slot size
+
+TEST(NextEraseEnd, FirstWriteRoundsUpToOneBlock) { EXPECT_EQ(next_erase_end(1024, PART), BLOCK); }
+
+TEST(NextEraseEnd, ExactBlockBoundaryDoesNotOverErase) { EXPECT_EQ(next_erase_end(BLOCK, PART), BLOCK); }
+
+TEST(NextEraseEnd, StraddlingWriteCoversNextBlock) { EXPECT_EQ(next_erase_end(BLOCK + 1, PART), 2 * BLOCK); }
+
+TEST(NextEraseEnd, ClampsToPartitionEnd) {
+  // Partition sizes are sector multiples but not always block multiples
+  constexpr size_t part = 27 * BLOCK + 4096;
+  EXPECT_EQ(next_erase_end(27 * BLOCK + 1, part), part);
+  EXPECT_EQ(next_erase_end(part, part), part);
+}
+
+// Bootloader staging seeds erased_end_ mid-block (e.g. 0x8000); the target for
+// a write past that seed must still cover the write end.
+TEST(NextEraseEnd, MidBlockSeedStillCovered) { EXPECT_EQ(next_erase_end(0x8000 + 1024, PART), BLOCK); }
+
+TEST(NextEraseEnd, SweepAlwaysCoversWriteEndWithinPartition) {
+  for (size_t end = 1; end <= PART; end += 4093) {
+    const size_t erased = next_erase_end(end, PART);
+    ASSERT_GE(erased, end);
+    ASSERT_LE(erased, PART);
+    // Block-aligned unless clamped at the partition end
+    ASSERT_TRUE(erased == PART || erased % BLOCK == 0);
+  }
+}
+
+}  // namespace esphome::ota::testing