From c9ea30e6c09229acde6d47fbace6b08ff827e170 Mon Sep 17 00:00:00 2001
From: Konstantin Kondrashov <konstantin@espressif.com>
Date: Fri, 20 Dec 2024 14:14:44 +0200
Subject: [PATCH] fix(app_update): Invalidate ota data slot of last boot app in
 esp_ota_begin

Closes https://github.com/espressif/esp-idf/issues/14808
---
 components/app_update/esp_ota_ops.c         | 34 +++++++++++++++++----
 components/app_update/include/esp_ota_ops.h | 15 +++++++++
 2 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/components/app_update/esp_ota_ops.c b/components/app_update/esp_ota_ops.c
index e4f7055345..63e33b9c09 100644
--- a/components/app_update/esp_ota_ops.c
+++ b/components/app_update/esp_ota_ops.c
@@ -182,8 +182,18 @@ esp_err_t esp_ota_begin(const esp_partition_t *partition, size_t image_size, esp
         } else {
             erase_size = ALIGN_UP(image_size, partition->erase_size);
         }
-        return esp_partition_erase_range(partition, 0, erase_size);
+        esp_err_t err = esp_partition_erase_range(partition, 0, erase_size);
+        if (err != ESP_OK) {
+            return err;
+        }
     }
+
+#ifdef CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE
+    if (is_ota_partition(partition)) {
+        esp_ota_invalidate_inactive_ota_data_slot();
+    }
+#endif
+
     return ESP_OK;
 }
 
@@ -955,7 +965,7 @@ esp_err_t esp_ota_get_state_partition(const esp_partition_t *partition, esp_ota_
     return ESP_OK;
 }
 
-esp_err_t esp_ota_erase_last_boot_app_partition(void)
+static esp_err_t erase_last_boot_app_partition(bool skip_app_part_erase)
 {
     esp_ota_select_entry_t otadata[2];
     const esp_partition_t* ota_data_partition = read_otadata(otadata);
@@ -987,13 +997,15 @@ esp_err_t esp_ota_erase_last_boot_app_partition(void)
         return ESP_FAIL;
     }
 
-    esp_err_t err = esp_partition_erase_range(last_boot_app_partition_from_otadata, 0, last_boot_app_partition_from_otadata->size);
-    if (err != ESP_OK) {
-        return err;
+    if (!skip_app_part_erase) {
+        esp_err_t err = esp_partition_erase_range(last_boot_app_partition_from_otadata, 0, last_boot_app_partition_from_otadata->size);
+        if (err != ESP_OK) {
+            return err;
+        }
     }
 
     int sec_id = inactive_otadata;
-    err = esp_partition_erase_range(ota_data_partition, sec_id * ota_data_partition->erase_size, ota_data_partition->erase_size);
+    esp_err_t err = esp_partition_erase_range(ota_data_partition, sec_id * ota_data_partition->erase_size, ota_data_partition->erase_size);
     if (err != ESP_OK) {
         return err;
     }
@@ -1001,6 +1013,16 @@ esp_err_t esp_ota_erase_last_boot_app_partition(void)
     return ESP_OK;
 }
 
+esp_err_t esp_ota_erase_last_boot_app_partition(void)
+{
+    return erase_last_boot_app_partition(false);
+}
+
+esp_err_t esp_ota_invalidate_inactive_ota_data_slot(void)
+{
+    return erase_last_boot_app_partition(true);
+}
+
 #if SOC_SUPPORT_SECURE_BOOT_REVOKE_KEY && CONFIG_SECURE_BOOT_V2_ENABLED
 
 // Validates the image at "app_pos" with the secure boot digests other than "revoked_key_index"
diff --git a/components/app_update/include/esp_ota_ops.h b/components/app_update/include/esp_ota_ops.h
index 0916f1f896..5f4a8a05cf 100644
--- a/components/app_update/include/esp_ota_ops.h
+++ b/components/app_update/include/esp_ota_ops.h
@@ -83,6 +83,7 @@ int esp_ota_get_app_elf_sha256(char* dst, size_t size) __attribute__((deprecated
  * use esp_ota_mark_app_valid_cancel_rollback() function for it (this should be done as early as possible when you first download a new application).
  *
  * Note: Rollback is applicable only for app type partitions.
+ * Note: For Rollback - The OTA data slot corresponding to the last boot application partition will be invalidated.
  *
  * @param partition  Pointer to info for partition which will receive the OTA update. Required.
  *                   This is considered as the staging partition (where OTA is downloaded), be default this also considered as the final partition which supposed to be updated.
@@ -302,6 +303,20 @@ esp_err_t esp_ota_get_partition_description(const esp_partition_t *partition, es
  */
 esp_err_t esp_ota_get_bootloader_description(const esp_partition_t *bootloader_partition, esp_bootloader_desc_t *desc);
 
+/**
+ * @brief Invalidate the OTA data slot associated with the last boot application partition.
+ *
+ * This function erases the OTA data slot corresponding to the last boot application partition,
+ * making the partition invalid for booting in future. The application partition itself
+ * is not erased, preserving its contents.
+ *
+ * @return
+ *      - ESP_OK: Successfully invalidated the OTA data slot.
+ *      - ESP_FAIL: Failed to invalidate the OTA data slot (e.g., invalid parameters, no OTA data partition, or other errors).
+ *      - Other error codes from `esp_partition_erase_range`.
+ */
+esp_err_t esp_ota_invalidate_inactive_ota_data_slot(void);
+
 /**
  * @brief Returns number of ota partitions provided in partition table.
  *