From a76882aa8ca3bf97a9da10d60a4f09abfa67322b Mon Sep 17 00:00:00 2001
From: John Bland <john.bland@wolfssl.com>
Date: Fri, 6 Dec 2024 20:15:34 -0500
Subject: [PATCH] properly unlock flash on re-entry of
 wolfBoot_swap_and_final_erase

and move it out of the lock logic of update and delta update
---
 src/update_flash.c | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/update_flash.c b/src/update_flash.c
index 20cc67cc0..b0c9368b2 100644
--- a/src/update_flash.c
+++ b/src/update_flash.c
@@ -245,6 +245,11 @@ static int wolfBoot_swap_and_final_erase(int resume)
     if ((resume == 1) && (swapDone == 0) && (st != IMG_STATE_FINAL_FLAGS)) {
         return -1;
     }
+    hal_flash_unlock();
+#ifdef EXT_FLASH
+    ext_flash_unlock();
+#endif
+
     if (swapDone == 0) {
         /* IMG_STATE_FINAL_FLAGS allows re-entry without blowing away swap */
         if (st != IMG_STATE_FINAL_FLAGS) {
@@ -282,6 +287,12 @@ static int wolfBoot_swap_and_final_erase(int resume)
     wolfBoot_set_partition_state(PART_BOOT, IMG_STATE_TESTING);
     /* erase the last sector(s) of update */
     wb_flash_erase(update, WOLFBOOT_PARTITION_SIZE - eraseLen, eraseLen);
+
+#ifdef EXT_FLASH
+    ext_flash_lock();
+#endif
+    hal_flash_lock();
+
     return 0;
 }
 #endif
@@ -480,15 +491,17 @@ static int wolfBoot_delta_update(struct wolfBoot_image *boot,
         sector++;
     }
 #ifndef DISABLE_BACKUP
-    /* start re-entrant final erase, return code is only for resumption in
-     * wolfBoot_start */
-    wolfBoot_swap_and_final_erase(0);
 #endif
 out:
 #ifdef EXT_FLASH
     ext_flash_lock();
 #endif
     hal_flash_lock();
+    /* start re-entrant final erase, return code is only for resumption in
+     * wolfBoot_start */
+    if (ret == 0) {
+        wolfBoot_swap_and_final_erase(0);
+    }
     /* encryption key was not erased, will be erased by success */
     return ret;
 }
@@ -741,14 +754,14 @@ static int RAMFUNCTION wolfBoot_update(int fallback_allowed)
     }
 #endif /* WOLFBOOT_FLASH_MULTI_SECTOR_ERASE */
 
-    /* start re-entrant final erase, return code is only for resumption in
-     * wolfBoot_start*/
-    wolfBoot_swap_and_final_erase(0);
     /* encryption key was not erased, will be erased by success */
     #ifdef EXT_FLASH
     ext_flash_lock();
     #endif
     hal_flash_lock();
+    /* start re-entrant final erase, return code is only for resumption in
+     * wolfBoot_start*/
+    wolfBoot_swap_and_final_erase(0);
 
 #else /* DISABLE_BACKUP */
     /* Direct Swap without power fail safety */