Fix corruption of last boot sector with encryption

include/wolfboot/wolfboot.h

@@ -407,6 +407,7 @@ int wolfBoot_ram_decrypt(uint8_t *src, uint8_t *dst);
 int wolfBoot_get_diffbase_hdr(uint8_t part, uint8_t **ptr);
 #endif
 
+int wolfBoot_initialize_encryption(void);
 int wolfBoot_set_encrypt_key(const uint8_t *key, const uint8_t *nonce);
 int wolfBoot_get_encrypt_key(uint8_t *key, uint8_t *nonce);
 int wolfBoot_erase_encrypt_key(void);

src/libwolfboot.c

@@ -64,6 +64,7 @@
 
 #if defined(EXT_ENCRYPTED)
 static int encrypt_initialized = 0;
+
 static uint8_t encrypt_iv_nonce[ENCRYPT_NONCE_SIZE] XALIGNED(4);
     #if defined(__WOLFBOOT)
         #include "encrypt.h"
@@ -73,8 +74,26 @@ static uint8_t encrypt_iv_nonce[ENCRYPT_NONCE_SIZE] XALIGNED(4);
         #define XMEMCPY memcpy
         #define XMEMCMP memcmp
     #endif
+
+#ifdef __WOLFBOOT
+int wolfBoot_initialize_encryption(void)
+{
+    if (!encrypt_initialized) {
+        if (crypto_init() != 0) {
+            return -1;
+        }
+        encrypt_initialized = 1;
+    }
+    return 0;
+}
 #endif
 
+#else
+    #define wolfBoot_initialize_encryption() (0)
+#endif /* EXT_ENCRYPTED */
+
+
+
 #if defined(EXT_FLASH) && defined(EXT_ENCRYPTED)
     #define ENCRYPT_TMP_SECRET_OFFSET (WOLFBOOT_PARTITION_SIZE - \
                          (TRAILER_SKIP + ENCRYPT_KEY_SIZE + ENCRYPT_NONCE_SIZE))
@@ -1039,9 +1058,8 @@ uint32_t wolfBoot_get_blob_version(uint8_t *blob)
     if (blob == NULL)
         return 0;
 #if defined(EXT_ENCRYPTED) && defined(MMU)
-    if (!encrypt_initialized)
-        if (crypto_init() < 0)
-            return 0;
+    if (wolfBoot_initialize_encryption() < 0)
+        return 0;
     decrypt_header(blob);
     img_bin = dec_hdr;
 #endif
@@ -1073,9 +1091,8 @@ uint16_t wolfBoot_get_blob_type(uint8_t *blob)
     uint32_t *magic = NULL;
     uint8_t *img_bin = blob;
 #if defined(EXT_ENCRYPTED) && defined(MMU)
-    if (!encrypt_initialized)
-        if (crypto_init() < 0)
-            return 0;
+    if (wolfBoot_initialize_encryption() < 0)
+        return 0;
     decrypt_header(blob);
     img_bin = dec_hdr;
 #endif
@@ -1111,9 +1128,8 @@ uint32_t wolfBoot_get_blob_diffbase_version(uint8_t *blob)
     uint32_t *magic = NULL;
     uint8_t *img_bin = blob;
 #if defined(EXT_ENCRYPTED) && defined(MMU)
-    if (!encrypt_initialized)
-        if (crypto_init() < 0)
-            return 0;
+    if (wolfBoot_initialize_encryption() < 0)
+        return 0;
     decrypt_header(blob);
     img_bin = dec_hdr;
 #endif
@@ -1790,10 +1806,9 @@ int RAMFUNCTION ext_flash_encrypt_write(uintptr_t address, const uint8_t *data,
                     ENCRYPT_BLOCK_SIZE) {
                 return ext_flash_write(address, data, len);
             }
-            if (!encrypt_initialized) {
-                if (crypto_init() < 0)
-                    return -1;
-            }
+            if (wolfBoot_initialize_encryption() < 0)
+                return -1;
+
             crypto_set_iv(encrypt_iv_nonce, iv_counter);
             break;
         case PART_SWAP:

src/update_flash.c

@@ -28,6 +28,7 @@
 #include "hal.h"
 #include "spi_flash.h"
 #include "wolfboot/wolfboot.h"
+#include "target.h"
 
 #include "delta.h"
 #include "printf.h"
@@ -152,7 +153,11 @@ static int RAMFUNCTION wolfBoot_copy_sector(struct wolfBoot_image *src,
         dst_sector_offset = 0;
 
 #ifdef EXT_ENCRYPTED
+    if (wolfBoot_initialize_encryption() < 0) {
+        return -1;
+    }
     wolfBoot_get_encrypt_key(key, nonce);
+
     if (src->part == PART_SWAP)
         iv_counter = dst_sector_offset;
     else
@@ -202,6 +207,56 @@ static int RAMFUNCTION wolfBoot_copy_sector(struct wolfBoot_image *src,
     return pos;
 }
 
+#ifdef EXT_ENCRYPTED
+static int RAMFUNCTION wolfBoot_backup_last_boot_sector(uint32_t sector)
+{
+    uint32_t pos = 0;
+    uint32_t src_sector_offset = (sector * WOLFBOOT_SECTOR_SIZE);
+    uint32_t dst_sector_offset = 0;
+    uint8_t key[ENCRYPT_KEY_SIZE];
+    uint8_t nonce[ENCRYPT_NONCE_SIZE];
+    uint32_t iv_counter;
+    uint8_t block[ENCRYPT_BLOCK_SIZE], encrypted_block[ENCRYPT_BLOCK_SIZE];
+    struct wolfBoot_image src[1], dst[1];
+
+    wolfBoot_open_image(src, PART_BOOT);
+    wolfBoot_open_image(dst, PART_SWAP);
+
+
+    wolfBoot_printf("Copy sector %d (part %d->%d)\n",
+        sector, src->part, dst->part);
+
+    wolfBoot_get_encrypt_key(key, nonce);
+    wolfBoot_printf("In function wolfBoot_backup_last_boot_sector (sector # %lu)\n", sector);
+
+    iv_counter = src_sector_offset;
+    iv_counter /= ENCRYPT_BLOCK_SIZE;
+    if (wolfBoot_initialize_encryption() < 0)
+        return -1;
+    crypto_set_iv(nonce, iv_counter);
+
+    /* Erase swap space */
+    wb_flash_erase(dst, dst_sector_offset, WOLFBOOT_SECTOR_SIZE);
+    if (PART_IS_EXT(dst)) {
+        uint32_t sz = 0;
+        uint32_t step = 0;
+        uint8_t *orig = (uint8_t *)(WOLFBOOT_PARTITION_BOOT_ADDRESS) + src_sector_offset; 
+        while (pos < WOLFBOOT_SECTOR_SIZE) {
+            uint32_t len = ENCRYPT_BLOCK_SIZE;
+            XMEMCPY(block, orig + pos, ENCRYPT_BLOCK_SIZE);
+            crypto_encrypt(encrypted_block, block, ENCRYPT_BLOCK_SIZE);
+            wb_flash_write(dst, dst_sector_offset + pos, encrypted_block, ENCRYPT_BLOCK_SIZE);
+            step++;
+            pos += ENCRYPT_BLOCK_SIZE;
+        }
+        return 0;
+    } else
+        return wolfBoot_copy_sector(src, dst, sector);
+}
+#else
+#define wolfBoot_backup_last_boot_sector(sec) wolfBoot_copy_sector(boot, swap, sec)
+#endif
+
 #if !defined(DISABLE_BACKUP) && !defined(CUSTOM_PARTITION_TRAILER)
 
 #ifdef EXT_ENCRYPTED
@@ -293,11 +348,13 @@ static int RAMFUNCTION wolfBoot_swap_and_final_erase(int resume)
     if (updateState != IMG_STATE_FINAL_FLAGS) {
         /* First, backup the staging sector (sector at tmpBootPos) into swap partition */
         /* This sector will be modified with the magic trailer, so we need to preserve it */
-        wolfBoot_copy_sector(boot, swap, tmpBootPos / WOLFBOOT_SECTOR_SIZE);
+        wolfBoot_backup_last_boot_sector(tmpBootPos / WOLFBOOT_SECTOR_SIZE);
+        wolfBoot_printf("Copied boot sector to swap\n");
         /* Mark update as being in final swap phase to allow resumption if power fails */
         wolfBoot_set_partition_state(PART_UPDATE, IMG_STATE_FINAL_FLAGS);
     }
 #ifdef EXT_ENCRYPTED
+    wolfBoot_printf("In function wolfBoot_final_swap: swapDone = %d\n", swapDone);
     if (swapDone == 0) {
         /* For encrypted images: Get the encryption key and IV */
         wolfBoot_get_encrypt_key((uint8_t*)tmpBuffer,
@@ -321,6 +378,7 @@ static int RAMFUNCTION wolfBoot_swap_and_final_erase(int resume)
 #endif
     /* Restore the original contents of the staging sector (with the magic trailer if encrypted) */
     if (tmpBootPos < boot->fw_size + IMAGE_HEADER_SIZE) {
+        wolfBoot_printf("Restoring last boot sector from swap\n");
         wolfBoot_copy_sector(swap, boot, tmpBootPos / WOLFBOOT_SECTOR_SIZE);
     }
     else {
@@ -461,6 +519,10 @@ static int wolfBoot_delta_update(struct wolfBoot_image *boot,
 #ifdef EXT_ENCRYPTED
                     uint32_t iv_counter = sector * WOLFBOOT_SECTOR_SIZE + len;
                     int wr_ret;
+                    if (wolfBoot_initialize_encryption() < 0) {
+                        ret = -1;
+                        goto out;
+                    }
                     iv_counter /= ENCRYPT_BLOCK_SIZE;
                     /* Encrypt + send */
                     crypto_set_iv(nonce, iv_counter);