Fix power-fail edge case in flash updater

.gitignore

@@ -180,6 +180,7 @@ tools/unit-tests/unit-update-flash
 tools/unit-tests/unit-update-flash-enc
 tools/unit-tests/unit-update-flash-db
 tools/unit-tests/unit-update-flash-enc-full
+tools/unit-tests/unit-update-flash-enc-intswap
 tools/unit-tests/unit-update-ram
 tools/unit-tests/unit-boot-x86-fsp
 tools/unit-tests/unit-image-rsa

src/update_flash.c

@@ -950,6 +950,58 @@ static uint32_t wolfBoot_get_total_size(struct wolfBoot_image* boot,
     return total_size;
 }
 
+#if !defined(DISABLE_BACKUP) && !defined(DELTA_UPDATES)
+/* Read the UPDATE image size from the copy of UPDATE sector 0 in SWAP. Used to
+ * resume a swap after a power cut during the BOOT->UPDATE copy of sector 0,
+ * which erases the UPDATE header. */
+static uint32_t wolfBoot_swap_update_size(void)
+{
+    struct wolfBoot_image swap;
+    uint32_t              hdr[2] = {0, 0}; /* magic and image size */
+    uint32_t              size;
+    uint8_t               flag = SECT_FLAG_NEW;
+#if defined(EXT_ENCRYPTED) && defined(EXT_FLASH) && PARTN_IS_EXT(PART_SWAP)
+    uint8_t key[ENCRYPT_KEY_SIZE];
+    uint8_t nonce[ENCRYPT_NONCE_SIZE];
+#endif
+
+    if ((wolfBoot_get_update_sector_flag(0, &flag) != 0) ||
+        (flag != SECT_FLAG_SWAPPING)) {
+        return 0;
+    }
+    wolfBoot_open_image(&swap, PART_SWAP);
+#if defined(EXT_FLASH) && PARTN_IS_EXT(PART_SWAP)
+#ifdef EXT_ENCRYPTED
+    /* External SWAP holds the sector still encrypted as at UPDATE offset 0,
+     * with the normal IV mapping */
+    if (wolfBoot_initialize_encryption() < 0) {
+        return 0;
+    }
+    wolfBoot_get_encrypt_key(key, nonce);
+    wolfBoot_enable_fallback_iv(0);
+    wolfBoot_crypto_set_iv(nonce, 0);
+    ext_flash_check_read((uintptr_t)swap.hdr, (uint8_t*)hdr, sizeof(hdr));
+    wolfBoot_zeroize(key, sizeof(key));
+    wolfBoot_zeroize(nonce, sizeof(nonce));
+#else
+    ext_flash_read((uintptr_t)swap.hdr, (uint8_t*)hdr, sizeof(hdr));
+#endif
+#else
+    /* Internal SWAP holds the sector as plaintext, also in encrypted builds */
+    memcpy(hdr, swap.hdr, sizeof(hdr));
+#endif
+    size = wolfBoot_image_size((uint8_t*)hdr);
+    if ((WOLFBOOT_HDR_GET_U32(hdr) != WOLFBOOT_MAGIC) ||
+        (size > MAX_UPDATE_SIZE)) {
+        return 0;
+    }
+    return size;
+}
+#else
+/* SWAP doesn't hold a copy of the UPDATE header for these cases */
+#define wolfBoot_swap_update_size() 0
+#endif
+
 #ifdef WOLFBOOT_PERSIST_FAILURE_STATUS
 /* Record a failure according to image verification flags */
 static void RAMFUNCTION wolfBoot_record_verify_failure(uint8_t phase,
@@ -1033,16 +1085,24 @@ static int RAMFUNCTION wolfBoot_update(int fallback_allowed)
             int prev = wolfBoot_enable_fallback_iv(1);
             (void)prev;
             update_open = wolfBoot_open_image(&update, PART_UPDATE);
-            if (update_open < 0) {
+            if (update_open >= 0) {
+                fallback_image = 1;
+            }
+        }
+#endif
+        if (update_open < 0) {
+            /* A power cut mid-swap may have erased the UPDATE header, try to
+             * resume using the copy in SWAP */
+            update.fw_size = wolfBoot_swap_update_size();
+            if (update.fw_size == 0) {
+#ifdef EXT_ENCRYPTED
                 wolfBoot_enable_fallback_iv(0);
+#endif
                 return -1;
             }
-            fallback_image = 1;
         }
+#ifdef EXT_ENCRYPTED
         wolfBoot_enable_fallback_iv(fallback_image);
-#else
-        if (update_open < 0)
-            return -1;
 #endif
         wolfBoot_open_image(&boot, PART_BOOT);
 #ifndef DISABLE_BACKUP

tools/test.mk

@@ -1243,52 +1243,54 @@ test-size-all:
 	# Re-measured 2026-09-29: NONE +24B (5204), ECC256/ECC384 +12B each, from
 	# the uart printf %u/is_signed change compiled into every build; RSA,
 	# RSAPSS, LMS and ML_DSA measured 8B smaller than their limits.
-	make test-size SIGN=NONE LIMIT=5204 NO_ARM_ASM=1
+	# Re-measured 2026-10-08: every configuration grew 96-108B from the
+	# power-fail swap resume (wolfBoot_swap_update_size) in update_flash.c.
+	make test-size SIGN=NONE LIMIT=5284 NO_ARM_ASM=1
 	make keysclean
-	make test-size SIGN=ED25519 LIMIT=12356 NO_ARM_ASM=1
+	make test-size SIGN=ED25519 LIMIT=12436 NO_ARM_ASM=1
 	make keysclean
-	make test-size SIGN=ECC256  LIMIT=19076 NO_ARM_ASM=1
+	make test-size SIGN=ECC256  LIMIT=19156 NO_ARM_ASM=1
 	make clean
-	make test-size SIGN=ECC256 NO_ASM=1 LIMIT=14120 NO_ARM_ASM=1
+	make test-size SIGN=ECC256 NO_ASM=1 LIMIT=14200 NO_ARM_ASM=1
 	make keysclean
-	make test-size SIGN=RSA2048 LIMIT=11984 NO_ARM_ASM=1
+	make test-size SIGN=RSA2048 LIMIT=11992 NO_ARM_ASM=1
 	make clean
-	make test-size SIGN=RSA2048 NO_ASM=1 LIMIT=12532 NO_ARM_ASM=1
+	make test-size SIGN=RSA2048 NO_ASM=1 LIMIT=12548 NO_ARM_ASM=1
 	make keysclean
-	make test-size SIGN=RSA4096 LIMIT=12284 NO_ARM_ASM=1
+	make test-size SIGN=RSA4096 LIMIT=12292 NO_ARM_ASM=1
 	make clean
-	make test-size SIGN=RSA4096 NO_ASM=1 LIMIT=12828 NO_ARM_ASM=1
+	make test-size SIGN=RSA4096 NO_ASM=1 LIMIT=12840 NO_ARM_ASM=1
 	make keysclean
-	make test-size SIGN=ECC384 LIMIT=19760 NO_ARM_ASM=1
+	make test-size SIGN=ECC384 LIMIT=19840 NO_ARM_ASM=1
 	make clean
-	make test-size SIGN=ECC384 NO_ASM=1 LIMIT=15476 NO_ARM_ASM=1
+	make test-size SIGN=ECC384 NO_ASM=1 LIMIT=15556 NO_ARM_ASM=1
 	make keysclean
-	make test-size SIGN=ED448 LIMIT=14424 NO_ARM_ASM=1
+	make test-size SIGN=ED448 LIMIT=14456 NO_ARM_ASM=1
 	make keysclean
-	make test-size SIGN=RSA3072 LIMIT=12124 NO_ARM_ASM=1
+	make test-size SIGN=RSA3072 LIMIT=12132 NO_ARM_ASM=1
 	make clean
-	make test-size SIGN=RSA3072 NO_ASM=1 LIMIT=12648 NO_ARM_ASM=1
+	make test-size SIGN=RSA3072 NO_ASM=1 LIMIT=12660 NO_ARM_ASM=1
 	make keysclean
-	make test-size SIGN=RSAPSS2048 LIMIT=13912 NO_ARM_ASM=1
+	make test-size SIGN=RSAPSS2048 LIMIT=13984 NO_ARM_ASM=1
 	make clean
-	make test-size SIGN=RSAPSS2048 NO_ASM=1 LIMIT=14460 NO_ARM_ASM=1
+	make test-size SIGN=RSAPSS2048 NO_ASM=1 LIMIT=14540 NO_ARM_ASM=1
 	make keysclean
-	make test-size SIGN=RSAPSS3072 LIMIT=14076 NO_ARM_ASM=1
+	make test-size SIGN=RSAPSS3072 LIMIT=14148 NO_ARM_ASM=1
 	make clean
-	make test-size SIGN=RSAPSS3072 NO_ASM=1 LIMIT=14600 NO_ARM_ASM=1
+	make test-size SIGN=RSAPSS3072 NO_ASM=1 LIMIT=14676 NO_ARM_ASM=1
 	make keysclean
-	make test-size SIGN=RSAPSS4096 LIMIT=14248 NO_ARM_ASM=1
+	make test-size SIGN=RSAPSS4096 LIMIT=14320 NO_ARM_ASM=1
 	make clean
-	make test-size SIGN=RSAPSS4096 NO_ASM=1 LIMIT=14792 NO_ARM_ASM=1
+	make test-size SIGN=RSAPSS4096 NO_ASM=1 LIMIT=14868 NO_ARM_ASM=1
 	make keysclean
 	make test-size SIGN=LMS LMS_LEVELS=2 LMS_HEIGHT=5 LMS_WINTERNITZ=8 \
 		WOLFBOOT_SMALL_STACK=0 IMAGE_SIGNATURE_SIZE=2644 \
-		IMAGE_HEADER_SIZE?=5288 LIMIT=8220 NO_ARM_ASM=1
+		IMAGE_HEADER_SIZE?=5288 LIMIT=8296 NO_ARM_ASM=1
 	make keysclean
 	make test-size SIGN=XMSS XMSS_PARAMS='XMSS-SHA2_10_256' \
 		IMAGE_SIGNATURE_SIZE=2500 IMAGE_HEADER_SIZE?=4096 \
-		LIMIT=8852 NO_ARM_ASM=1
+		LIMIT=8932 NO_ARM_ASM=1
 	make keysclean
 	make clean
-	make test-size SIGN=ML_DSA ML_DSA_LEVEL=2 LIMIT=19694 \
+	make test-size SIGN=ML_DSA ML_DSA_LEVEL=2 LIMIT=19798 \
 		IMAGE_SIGNATURE_SIZE=2420 IMAGE_HEADER_SIZE?=8192

tools/unit-tests/Makefile

@@ -67,7 +67,7 @@ TESTS:=unit-parser unit-parser-large-header unit-fdt unit-extflash unit-string \
        unit-update-flash-hook \
        unit-update-flash-self-update \
        unit-nsc-update \
-       unit-update-flash-enc unit-update-flash-enc-full unit-update-ram unit-update-ram-uboot unit-update-ram-enc unit-update-ram-enc-nopart unit-update-ram-nofixed unit-update-ram-nofixed-noramboot unit-update-ram-noramboot unit-update-ram-custom-trailer unit-custom-trailer-nopart unit-update-flash-hwswap unit-pkcs11_store unit-psa_store unit-wolfhsm_flash_hal unit-disk \
+       unit-update-flash-enc unit-update-flash-enc-full unit-update-flash-enc-intswap unit-update-ram unit-update-ram-uboot unit-update-ram-enc unit-update-ram-enc-nopart unit-update-ram-nofixed unit-update-ram-nofixed-noramboot unit-update-ram-noramboot unit-update-ram-custom-trailer unit-custom-trailer-nopart unit-update-flash-hwswap unit-pkcs11_store unit-psa_store unit-wolfhsm_flash_hal unit-disk \
        unit-update-disk unit-update-disk-confirm unit-update-disk-golden unit-update-disk-fsp unit-update-disk-oob unit-update-disk-fit unit-multiboot unit-boot-x86-fsp unit-loader-tpm-init unit-qspi-flash unit-fwtpm-stub unit-tpm-rsa-exp \
        unit-image-nopart unit-image-sha384 unit-image-sha3-384 unit-image-dts \
        unit-image-dts-sha384 unit-image-dts-sha3-384 unit-store-sbrk \
@@ -1064,6 +1064,19 @@ unit-update-flash-enc-full: ../../include/target.h unit-update-flash.c
 	$(WOLFBOOT_LIB_WOLFSSL)/wolfcrypt/src/chacha.c \
 	$(CFLAGS) $(LDFLAGS)
 
+# Same as unit-update-flash-enc-full but with SWAP in internal flash
+# (NO_SWAP_EXT=1), where the swap keeps the UPDATE sector as plaintext.
+unit-update-flash-enc-intswap:CFLAGS+=-DMOCK_PARTITIONS -DWOLFBOOT_NO_SIGN -DUNIT_TEST_AUTH \
+	-DWOLFBOOT_HASH_SHA256 -DPRINTF_ENABLED -DEXT_FLASH -DPART_UPDATE_EXT \
+	-DEXT_ENCRYPTED -DENCRYPT_WITH_CHACHA -DHAVE_CHACHA \
+	-DCUSTOM_ENCRYPT_KEY \
+	-DWOLFBOOT_ORIGIN=MOCK_ADDRESS_BOOT -DBOOTLOADER_PARTITION_SIZE=WOLFBOOT_PARTITION_SIZE
+unit-update-flash-enc-intswap: ../../include/target.h unit-update-flash.c
+	gcc -o $@ unit-update-flash.c ../../src/image.c \
+	$(WOLFBOOT_LIB_WOLFSSL)/wolfcrypt/src/sha256.c \
+	$(WOLFBOOT_LIB_WOLFSSL)/wolfcrypt/src/chacha.c \
+	$(CFLAGS) $(LDFLAGS)
+
 unit-update-ram: ../../include/target.h unit-update-ram.c
 	gcc -o $@ unit-update-ram.c ../../src/image.c $(WOLFBOOT_LIB_WOLFSSL)/wolfcrypt/src/sha256.c $(CFLAGS) $(LDFLAGS)
 

tools/unit-tests/unit-mock-flash.c

@@ -158,7 +158,7 @@ int hal_flash_write(haladdr_t address, const uint8_t *data, int len)
         return -1;
     }
     if ((address >= WOLFBOOT_PARTITION_SWAP_ADDRESS) &&
-            (address < WOLFBOOT_PARTITION_UPDATE_ADDRESS + WOLFBOOT_SECTOR_SIZE)) {
+            (address < WOLFBOOT_PARTITION_SWAP_ADDRESS + WOLFBOOT_SECTOR_SIZE)) {
         for (i = 0; i < len; i++) {
             a[i] = data[i];
         }