bootutil: Store image encrypted in scratch area

boot/boot_serial/src/boot_serial_encryption.c

@@ -175,7 +175,7 @@ decrypt_region_inplace(struct boot_loader_state *state,
                     blk_sz = tlv_off - (off + bytes_copied);
                 }
             }
-            boot_encrypt(BOOT_CURR_ENC(state), image_index, fap,
+            boot_encrypt(BOOT_CURR_ENC(state), image_index, flash_area_get_id(fap),
                     (off + bytes_copied + idx) - hdr->ih_hdr_size, blk_sz,
                     blk_off, &buf[idx]);
         }

boot/bootutil/include/bootutil/enc_key.h

@@ -59,7 +59,7 @@ int boot_enc_decrypt(const uint8_t *buf, uint8_t *enckey);
 bool boot_enc_valid(struct enc_key_data *enc_state, int image_index,
         const struct flash_area *fap);
 void boot_encrypt(struct enc_key_data *enc_state, int image_index,
-        const struct flash_area *fap, uint32_t off, uint32_t sz,
+        int fa_id, uint32_t off, uint32_t sz,
         uint32_t blk_off, uint8_t *buf);
 void boot_enc_zeroize(struct enc_key_data *enc_state);
 

boot/bootutil/src/encrypted.c

@@ -681,7 +681,7 @@ boot_enc_valid(struct enc_key_data *enc_state, int image_index,
 
 void
 boot_encrypt(struct enc_key_data *enc_state, int image_index,
-        const struct flash_area *fap, uint32_t off, uint32_t sz,
+        int fa_id, uint32_t off, uint32_t sz,
         uint32_t blk_off, uint8_t *buf)
 {
     struct enc_key_data *enc;
@@ -701,7 +701,7 @@ boot_encrypt(struct enc_key_data *enc_state, int image_index,
     nonce[14] = (uint8_t)(off >> 8);
     nonce[15] = (uint8_t)off;
 
-    rc = flash_area_id_to_multi_image_slot(image_index, flash_area_get_id(fap));
+    rc = flash_area_id_to_multi_image_slot(image_index, fa_id);
     if (rc < 0) {
         assert(0);
         return;

boot/bootutil/src/image_validate.c

@@ -150,7 +150,7 @@ bootutil_img_hash(struct enc_key_data *enc_state, int image_index,
             /* Only payload is encrypted (area between header and TLVs) */
             if (off >= hdr_size && off < tlv_off) {
                 blk_off = (off - hdr_size) & 0xf;
-                boot_encrypt(enc_state, image_index, fap, off - hdr_size,
+                boot_encrypt(enc_state, image_index, flash_area_get_id(fap), off - hdr_size,
                         blk_sz, blk_off, tmp_buf);
             }
         }

boot/bootutil/src/loader.c

@@ -1218,10 +1218,13 @@ boot_copy_region(struct boot_loader_state *state,
     uint32_t off;
     uint32_t tlv_off;
     size_t blk_off;
+    int enc_area_id;
     struct image_header *hdr;
     uint16_t idx;
     uint32_t blk_sz;
     uint8_t image_index;
+    bool encrypted_src;
+    bool encrypted_dst;
 #endif
 
     TARGET_STATIC uint8_t buf[BUF_SZ] __attribute__((aligned(4)));
@@ -1245,25 +1248,22 @@ boot_copy_region(struct boot_loader_state *state,
 
 #ifdef MCUBOOT_ENC_IMAGES
         image_index = BOOT_CURR_IMG(state);
-        if ((flash_area_get_id(fap_src) == FLASH_AREA_IMAGE_SECONDARY(image_index) ||
-             flash_area_get_id(fap_dst) == FLASH_AREA_IMAGE_SECONDARY(image_index)) &&
-            !(flash_area_get_id(fap_src) == FLASH_AREA_IMAGE_SECONDARY(image_index) &&
-              flash_area_get_id(fap_dst) == FLASH_AREA_IMAGE_SECONDARY(image_index))) {
-            /* assume the secondary slot as src, needs decryption */
-            hdr = boot_img_hdr(state, BOOT_SECONDARY_SLOT);
-#if !defined(MCUBOOT_SWAP_USING_MOVE)
-            off = off_src;
-            if (flash_area_get_id(fap_dst) == FLASH_AREA_IMAGE_SECONDARY(image_index)) {
-                /* might need encryption (metadata from the primary slot) */
-                hdr = boot_img_hdr(state, BOOT_PRIMARY_SLOT);
-                off = off_dst;
-            }
-#else
+        encrypted_src = (flash_area_get_id(fap_src) != FLASH_AREA_IMAGE_PRIMARY(image_index));
+        encrypted_dst = (flash_area_get_id(fap_dst) != FLASH_AREA_IMAGE_PRIMARY(image_index));
+
+        if (encrypted_src != encrypted_dst) {
             off = off_dst;
-            if (flash_area_get_id(fap_dst) == FLASH_AREA_IMAGE_SECONDARY(image_index)) {
+
+            if (encrypted_dst) {
+                /* Need encryption, metadata from the primary slot */
                 hdr = boot_img_hdr(state, BOOT_PRIMARY_SLOT);
+                enc_area_id = FLASH_AREA_IMAGE_PRIMARY(image_index);
+            } else {
+                /* Need decryption, metadata from the secondary slot */
+                hdr = boot_img_hdr(state, BOOT_SECONDARY_SLOT);
+                enc_area_id = FLASH_AREA_IMAGE_SECONDARY(image_index);
             }
-#endif
+
             if (IS_ENCRYPTED(hdr)) {
                 uint32_t abs_off = off + bytes_copied;
                 if (abs_off < hdr->ih_hdr_size) {
@@ -1294,7 +1294,7 @@ boot_copy_region(struct boot_loader_state *state,
                             blk_sz = tlv_off - abs_off;
                         }
                     }
-                    boot_encrypt(BOOT_CURR_ENC(state), image_index, fap_src,
+                    boot_encrypt(BOOT_CURR_ENC(state), image_index, enc_area_id,
                             (abs_off + idx) - hdr->ih_hdr_size, blk_sz,
                             blk_off, &buf[idx]);
                 }
@@ -1872,7 +1872,7 @@ boot_prepare_image_for_update(struct boot_loader_state *state,
         }
 #endif
 
-#ifdef MCUBOOT_SWAP_USING_MOVE
+#if defined(MCUBOOT_SWAP_USING_SCRATCH) || defined(MCUBOOT_SWAP_USING_MOVE)
         /*
          * Must re-read image headers because the boot status might
          * have been updated in the previous function call.
@@ -2774,13 +2774,13 @@ boot_decrypt_and_copy_image_to_sram(struct boot_loader_state *state,
              * Part of the chunk is encrypted payload */
             blk_off = ((bytes_copied) - hdr->ih_hdr_size) & 0xf;
             blk_sz = tlv_off - (bytes_copied);
-            boot_encrypt(BOOT_CURR_ENC(state), image_index, fap_src,
+            boot_encrypt(BOOT_CURR_ENC(state), image_index, area_id,
                 (bytes_copied + idx) - hdr->ih_hdr_size, blk_sz,
                 blk_off, cur_dst);
         } else {
             /* Image encrypted payload section */
             blk_off = ((bytes_copied) - hdr->ih_hdr_size) & 0xf;
-            boot_encrypt(BOOT_CURR_ENC(state), image_index, fap_src,
+            boot_encrypt(BOOT_CURR_ENC(state), image_index, area_id,
                     (bytes_copied + idx) - hdr->ih_hdr_size, blk_sz,
                     blk_off, cur_dst);
         }

boot/bootutil/src/swap_scratch.c

@@ -47,35 +47,6 @@ int boot_status_fails = 0;
 #define BOOT_STATUS_ASSERT(x) ASSERT(x)
 #endif
 
-int
-boot_read_image_header(struct boot_loader_state *state, int slot,
-                       struct image_header *out_hdr, struct boot_status *bs)
-{
-    const struct flash_area *fap;
-    int area_id;
-    int rc = 0;
-
-    (void)bs;
-
-#if (BOOT_IMAGE_NUMBER == 1)
-    (void)state;
-#endif
-
-    area_id = flash_area_id_from_multi_image_slot(BOOT_CURR_IMG(state), slot);
-
-    rc = flash_area_open(area_id, &fap);
-    if (rc == 0) {
-        rc = flash_area_read(fap, 0, out_hdr, sizeof *out_hdr);
-        flash_area_close(fap);
-    }
-
-    if (rc != 0) {
-        rc = BOOT_EFLASH;
-    }
-
-    return rc;
-}
-
 #if !defined(MCUBOOT_DIRECT_XIP) && !defined(MCUBOOT_RAM_LOAD)
 /**
  * Reads the status of a partially-completed swap, if any.  This is necessary
@@ -470,6 +441,84 @@ boot_copy_sz(const struct boot_loader_state *state, int last_sector_idx,
     return sz;
 }
 
+/**
+ * Finds the index of the last sector in the primary slot that needs swapping.
+ *
+ * @param state     Current bootloader's state.
+ * @param copy_size Total number of bytes to swap.
+ *
+ * @return          Index of the last sector in the primary slot that needs swapping.
+ */
+static int
+find_last_sector_idx(const struct boot_loader_state *state, uint32_t copy_size)
+{
+    int last_sector_idx;
+    uint32_t primary_slot_size;
+    uint32_t secondary_slot_size;
+
+    primary_slot_size = 0;
+    secondary_slot_size = 0;
+    last_sector_idx = 0;
+
+    /*
+     * Knowing the size of the largest image between both slots, here we
+     * find what is the last sector in the primary slot that needs swapping.
+     * Since we already know that both slots are compatible, the secondary
+     * slot's last sector is not really required after this check is finished.
+     */
+    while (1) {
+        if ((primary_slot_size < copy_size) ||
+            (primary_slot_size < secondary_slot_size)) {
+           primary_slot_size += boot_img_sector_size(state,
+                                                     BOOT_PRIMARY_SLOT,
+                                                     last_sector_idx);
+        }
+        if ((secondary_slot_size < copy_size) ||
+            (secondary_slot_size < primary_slot_size)) {
+           secondary_slot_size += boot_img_sector_size(state,
+                                                       BOOT_SECONDARY_SLOT,
+                                                       last_sector_idx);
+        }
+        if (primary_slot_size >= copy_size &&
+                secondary_slot_size >= copy_size &&
+                primary_slot_size == secondary_slot_size) {
+            break;
+        }
+        last_sector_idx++;
+    }
+
+    return last_sector_idx;
+}
+
+/**
+ * Finds the number of swap operations that have to be performed to swap the two images.
+ *
+ * @param state     Current bootloader's state.
+ * @param copy_size Total number of bytes to swap.
+ *
+ * @return          The number of swap operations that have to be performed.
+*/
+static uint32_t
+find_swap_count(const struct boot_loader_state *state, uint32_t copy_size)
+{
+    int first_sector_idx;
+    int last_sector_idx;
+    uint32_t swap_count;
+
+    last_sector_idx = find_last_sector_idx(state, copy_size);
+
+    swap_count = 0;
+
+    while (last_sector_idx >= 0) {
+        boot_copy_sz(state, last_sector_idx, &first_sector_idx);
+
+        last_sector_idx = first_sector_idx - 1;
+        swap_count++;
+    }
+
+    return swap_count;
+}
+
 /**
  * Swaps the contents of two flash regions within the two image slots.
  *
@@ -691,43 +740,10 @@ swap_run(struct boot_loader_state *state, struct boot_status *bs,
     int first_sector_idx;
     int last_sector_idx;
     uint32_t swap_idx;
-    int last_idx_secondary_slot;
-    uint32_t primary_slot_size;
-    uint32_t secondary_slot_size;
-    primary_slot_size = 0;
-    secondary_slot_size = 0;
-    last_sector_idx = 0;
-    last_idx_secondary_slot = 0;
 
     BOOT_LOG_INF("Starting swap using scratch algorithm.");
 
-    /*
-     * Knowing the size of the largest image between both slots, here we
-     * find what is the last sector in the primary slot that needs swapping.
-     * Since we already know that both slots are compatible, the secondary
-     * slot's last sector is not really required after this check is finished.
-     */
-    while (1) {
-        if ((primary_slot_size < copy_size) ||
-            (primary_slot_size < secondary_slot_size)) {
-           primary_slot_size += boot_img_sector_size(state,
-                                                     BOOT_PRIMARY_SLOT,
-                                                     last_sector_idx);
-        }
-        if ((secondary_slot_size < copy_size) ||
-            (secondary_slot_size < primary_slot_size)) {
-           secondary_slot_size += boot_img_sector_size(state,
-                                                       BOOT_SECONDARY_SLOT,
-                                                       last_idx_secondary_slot);
-        }
-        if (primary_slot_size >= copy_size &&
-                secondary_slot_size >= copy_size &&
-                primary_slot_size == secondary_slot_size) {
-            break;
-        }
-        last_sector_idx++;
-        last_idx_secondary_slot++;
-    }
+    last_sector_idx = find_last_sector_idx(state, copy_size);
 
     swap_idx = 0;
     while (last_sector_idx >= 0) {
@@ -857,4 +873,95 @@ int app_max_size(struct boot_loader_state *state)
 
 #endif /* !MCUBOOT_DIRECT_XIP && !MCUBOOT_RAM_LOAD */
 
+int
+boot_read_image_header(struct boot_loader_state *state, int slot,
+                       struct image_header *out_hdr, struct boot_status *bs)
+{
+    const struct flash_area *fap;
+#ifdef MCUBOOT_SWAP_USING_SCRATCH
+    uint32_t swap_count;
+    uint32_t swap_size;
+#endif
+    int area_id;
+    int hdr_slot;
+    int rc = 0;
+
+#ifndef MCUBOOT_SWAP_USING_SCRATCH
+    (void)bs;
+#endif
+
+#if (BOOT_IMAGE_NUMBER == 1)
+    (void)state;
+#endif
+
+    hdr_slot = slot;
+
+#ifdef MCUBOOT_SWAP_USING_SCRATCH
+    /* If the slots are being swapped, the headers might have been moved to scratch area or to the
+     * other slot depending on the progress of the swap process.
+     */
+    if (bs && !boot_status_is_reset(bs)) {
+        rc = boot_find_status(BOOT_CURR_IMG(state), &fap);
+
+        if (rc != 0) {
+            rc = BOOT_EFLASH;
+            goto done;
+        }
+
+        rc = boot_read_swap_size(fap, &swap_size);
+        flash_area_close(fap);
+
+        if (rc != 0) {
+            rc = BOOT_EFLASH;
+            goto done;
+        }
+
+        swap_count = find_swap_count(state, swap_size);
+
+        if (bs->idx - BOOT_STATUS_IDX_0 >= swap_count) {
+            /* If all segments have been swapped, the header is located in the other slot */
+            hdr_slot = (slot == BOOT_PRIMARY_SLOT) ? BOOT_SECONDARY_SLOT : BOOT_PRIMARY_SLOT;
+        } else if (bs->idx - BOOT_STATUS_IDX_0 == swap_count - 1) {
+            /* If the last swap operation is in progress, the headers are currently being swapped
+             * since the first segment of each slot is the last to be processed.
+             */
+
+            if (slot == BOOT_SECONDARY_SLOT && bs->state >= BOOT_STATUS_STATE_1) {
+                /* After BOOT_STATUS_STATE_1, the secondary image's header has been moved to the
+                 * scratch area.
+                 */
+                hdr_slot = BOOT_NUM_SLOTS;
+            } else if (slot == BOOT_PRIMARY_SLOT && bs->state >= BOOT_STATUS_STATE_2) {
+                /* After BOOT_STATUS_STATE_2, the primary image's header has been moved to the
+                 * secondary slot.
+                 */
+                hdr_slot = BOOT_SECONDARY_SLOT;
+            }
+        }
+    }
+
+    if (hdr_slot == BOOT_NUM_SLOTS) {
+        area_id = FLASH_AREA_IMAGE_SCRATCH;
+    } else {
+        area_id = flash_area_id_from_multi_image_slot(BOOT_CURR_IMG(state), hdr_slot);
+    }
+#else
+    area_id = flash_area_id_from_multi_image_slot(BOOT_CURR_IMG(state), hdr_slot);
+#endif
+
+    rc = flash_area_open(area_id, &fap);
+    if (rc == 0) {
+        rc = flash_area_read(fap, 0, out_hdr, sizeof *out_hdr);
+        flash_area_close(fap);
+    }
+
+    if (rc != 0) {
+        rc = BOOT_EFLASH;
+        goto done;
+    }
+
+done:
+    return rc;
+}
+
 #endif /* !MCUBOOT_SWAP_USING_MOVE */