Fix corruption of last boot sector with encryption

As reported by @reza-hdd in #589:

> When Firmware is large enough to consume all the allocated boot partition sectors, in the the final swap and erase operations (_wolfBoot_swap_and_final_erase_) the last sector of FW in Boot partition, which is used as a temporary sector for a copy operation, gets corrupted after being copied in external Swap partition and copied back. It looks like the contents are copied to the external Swap partition without being encrypted, but get decrypted when it is read back.

This was due to the function wolfBoot_copy_sector() assuming that all
copies having an external SWAP as destination would not need any
encryption, as the "normal" case during update is to copy already
encrypted sectors from UPDATE->SWAP.

In the final state, after the update is successfully applied, the
mechanism saves a copy of the last sector from BOOT->SWAP. In this case,
if the application is big enough, an extra function is needed to
encrypt the content of this last sector before writing it to SWAP.

Restoring the backup is OK (normal case SWAP->BOOT, decryption was
already there as also noted by the reporter).

Author: danielinux

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);