Adds total size of a swap to the trailer

When starting a swap upgrade, the total size of data to be swapped is
calculated only at the beginning and saved to the trailer. This avoids
having to use complicated heuristics to find the total swap size, which
might depend on data that was already moved. When resuming a swap, the
size is found in the trailer and used.

Also includes some small comment fixes and refactors.

Signed-off-by: Fabio Utzig <utzig@apache.org>
Signed-off-by: David Brown <david.brown@linaro.org>
diff --git a/boot/bootutil/src/bootutil_misc.c b/boot/bootutil/src/bootutil_misc.c
index a71df72..bf4e9b8 100644
--- a/boot/bootutil/src/bootutil_misc.c
+++ b/boot/bootutil/src/bootutil_misc.c
@@ -120,7 +120,7 @@
 {
     return /* state for all sectors */
            BOOT_STATUS_MAX_ENTRIES * BOOT_STATUS_STATE_COUNT * min_write_sz +
-           BOOT_MAX_ALIGN * 2 /* copy_done + image_ok */                    +
+           BOOT_MAX_ALIGN * 3 /* copy_done + image_ok + swap_size */        +
            BOOT_MAGIC_SZ;
 }
 
@@ -128,7 +128,7 @@
 boot_scratch_trailer_sz(uint8_t min_write_sz)
 {
     return BOOT_STATUS_STATE_COUNT * min_write_sz +  /* state for one sector */
-           BOOT_MAX_ALIGN                         +  /* image_ok */
+           BOOT_MAX_ALIGN * 2                     +  /* image_ok + swap_size */
            BOOT_MAGIC_SZ;
 }
 
@@ -186,6 +186,20 @@
     return fap->fa_size - BOOT_MAGIC_SZ - BOOT_MAX_ALIGN;
 }
 
+static uint32_t
+boot_swap_size_off(const struct flash_area *fap)
+{
+    /*
+     * The "swap_size" field if located just before the trailer.
+     * The scratch slot doesn't store "copy_done"...
+     */
+    if (fap->fa_id == FLASH_AREA_IMAGE_SCRATCH) {
+        return fap->fa_size - BOOT_MAGIC_SZ - BOOT_MAX_ALIGN * 2;
+    }
+
+    return fap->fa_size - BOOT_MAGIC_SZ - BOOT_MAX_ALIGN * 3;
+}
+
 int
 boot_read_swap_state(const struct flash_area *fap,
                      struct boot_swap_state *state)
@@ -203,14 +217,14 @@
 
     if (fap->fa_id != FLASH_AREA_IMAGE_SCRATCH) {
         off = boot_copy_done_off(fap);
-        rc = flash_area_read(fap, off, &state->copy_done, 1);
+        rc = flash_area_read(fap, off, &state->copy_done, sizeof state->copy_done);
         if (rc != 0) {
             return BOOT_EFLASH;
         }
     }
 
     off = boot_image_ok_off(fap);
-    rc = flash_area_read(fap, off, &state->image_ok, 1);
+    rc = flash_area_read(fap, off, &state->image_ok, sizeof state->image_ok);
     if (rc != 0) {
         return BOOT_EFLASH;
     }
@@ -246,6 +260,68 @@
 }
 
 int
+boot_read_swap_size(uint32_t *swap_size)
+{
+    uint32_t magic[BOOT_MAGIC_SZ];
+    uint32_t off;
+    const struct flash_area *fap;
+    int rc;
+
+    /*
+     * In the middle a swap, tries to locate the saved swap size. Looks
+     * for a valid magic, first on Slot 0, then on scratch. Both "slots"
+     * can end up being temporary storage for a swap and it is assumed
+     * that if magic is valid then swap size is too, because magic is
+     * always written in the last step.
+     */
+
+    rc = flash_area_open(FLASH_AREA_IMAGE_0, &fap);
+    if (rc != 0) {
+        return BOOT_EFLASH;
+    }
+
+    off = boot_magic_off(fap);
+    rc = flash_area_read(fap, off, magic, BOOT_MAGIC_SZ);
+    if (rc != 0) {
+        rc = BOOT_EFLASH;
+        goto out;
+    }
+
+    if (memcmp(magic, boot_img_magic, BOOT_MAGIC_SZ) != 0) {
+        /*
+         * If Slot 0 's magic is not valid, try scratch...
+         */
+
+        flash_area_close(fap);
+
+        rc = flash_area_open(FLASH_AREA_IMAGE_SCRATCH, &fap);
+        if (rc != 0) {
+            return BOOT_EFLASH;
+        }
+
+        off = boot_magic_off(fap);
+        rc = flash_area_read(fap, off, magic, BOOT_MAGIC_SZ);
+        if (rc != 0) {
+            rc = BOOT_EFLASH;
+            goto out;
+        }
+
+        assert(memcmp(magic, boot_img_magic, BOOT_MAGIC_SZ) == 0);
+    }
+
+    off = boot_swap_size_off(fap);
+    rc = flash_area_read(fap, off, swap_size, sizeof *swap_size);
+    if (rc != 0) {
+        rc = BOOT_EFLASH;
+    }
+
+out:
+    flash_area_close(fap);
+    return rc;
+}
+
+
+int
 boot_write_magic(const struct flash_area *fap)
 {
     uint32_t off;
@@ -306,6 +382,31 @@
 }
 
 int
+boot_write_swap_size(const struct flash_area *fap, uint32_t swap_size)
+{
+    uint32_t off;
+    int rc;
+    uint8_t buf[BOOT_MAX_ALIGN];
+    uint8_t align;
+
+    off = boot_swap_size_off(fap);
+    align = hal_flash_align(fap->fa_device_id);
+    assert(align <= BOOT_MAX_ALIGN);
+    if (align < sizeof swap_size) {
+        align = sizeof swap_size;
+    }
+    memset(buf, 0xFF, BOOT_MAX_ALIGN);
+    memcpy(buf, (uint8_t *)&swap_size, sizeof swap_size);
+
+    rc = flash_area_write(fap, off, buf, align);
+    if (rc != 0) {
+        return BOOT_EFLASH;
+    }
+
+    return 0;
+}
+
+int
 boot_swap_type(void)
 {
     const struct boot_swap_table *table;
diff --git a/boot/bootutil/src/bootutil_priv.h b/boot/bootutil/src/bootutil_priv.h
index 3530338..c1cf779 100644
--- a/boot/bootutil/src/bootutil_priv.h
+++ b/boot/bootutil/src/bootutil_priv.h
@@ -47,6 +47,7 @@
     uint32_t idx;         /* Which area we're operating on */
     uint8_t state;        /* Which part of the swapping process are we at */
     uint8_t use_scratch;  /* Are status bytes ever written to scratch? */
+    uint32_t swap_size;   /* Total size of swapped image */
 };
 
 #define BOOT_MAGIC_GOOD  1
@@ -59,15 +60,19 @@
  *  0                   1                   2                   3
  *  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
  * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
- * ~                        MAGIC (16 octets)                      ~
- * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
  * ~                                                               ~
  * ~                Swap status (variable, aligned)                ~
  * ~                                                               ~
  * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
- * |   Copy done   |     0xff padding (up to min-write-sz - 1)     ~
+ * |                          Swap size                            |
  * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
- * |   Image OK    |     0xff padding (up to min-write-sz - 1)     ~
+ * ~                  0xff padding (MAX ALIGN - 4)                 ~
+ * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+ * |   Copy done   |          0xff padding (MAX ALIGN - 1)         ~
+ * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+ * |   Image OK    |          0xff padding (MAX ALIGN - 1)         ~
+ * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
+ * ~                        MAGIC (16 octets)                      ~
  * +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
  */
 
@@ -140,6 +145,8 @@
 int boot_schedule_test_swap(void);
 int boot_write_copy_done(const struct flash_area *fap);
 int boot_write_image_ok(const struct flash_area *fap);
+int boot_write_swap_size(const struct flash_area *fap, uint32_t swap_size);
+int boot_read_swap_size(uint32_t *swap_size);
 
 /*
  * Accessors for the contents of struct boot_loader_state.
diff --git a/boot/bootutil/src/loader.c b/boot/bootutil/src/loader.c
index 3a30bc8..b03b745 100644
--- a/boot/bootutil/src/loader.c
+++ b/boot/bootutil/src/loader.c
@@ -206,28 +206,6 @@
     return BOOT_SWAP_TYPE_FAIL;
 }
 
-#ifndef MCUBOOT_OVERWRITE_ONLY
-static int
-boot_read_header_from_scratch(struct image_header *out_hdr)
-{
-    const struct flash_area *fap;
-    int rc;
-
-    rc = flash_area_open(FLASH_AREA_IMAGE_SCRATCH, &fap);
-    if (rc != 0) {
-        return BOOT_EFLASH;
-    }
-
-    rc = flash_area_read(fap, 0, out_hdr, sizeof *out_hdr);
-    if (rc != 0) {
-        rc = BOOT_EFLASH;
-    }
-
-    flash_area_close(fap);
-    return rc;
-}
-#endif /* !MCUBOOT_OVERWRITE_ONLY */
-
 /*
  * Compute the total size of the given image.  Includes the size of
  * the TLVs.
@@ -477,7 +455,10 @@
         return BOOT_EFLASH;
     }
 
-    return boot_read_status_bytes(fap, bs);
+    rc = boot_read_status_bytes(fap, bs);
+
+    flash_area_close(fap);
+    return rc;
 }
 
 /**
@@ -797,7 +778,7 @@
 }
 
 static inline int
-boot_status_init_by_id(int flash_area_id)
+boot_status_init_by_id(int flash_area_id, const struct boot_status *bs)
 {
     const struct flash_area *fap;
     struct boot_swap_state swap_state;
@@ -814,6 +795,9 @@
         assert(rc == 0);
     }
 
+    rc = boot_write_swap_size(fap, bs->swap_size);
+    assert(rc == 0);
+
     rc = boot_write_magic(fap);
     assert(rc == 0);
 
@@ -907,7 +891,7 @@
 
         if (bs->idx == 0) {
             if (bs->use_scratch) {
-                boot_status_init_by_id(FLASH_AREA_IMAGE_SCRATCH);
+                boot_status_init_by_id(FLASH_AREA_IMAGE_SCRATCH, bs);
             } else {
                 /* Prepare the status area... here it is known that the
                  * last sector is not being used by the image data so it's
@@ -916,7 +900,7 @@
                 rc = boot_erase_last_sector_by_id(FLASH_AREA_IMAGE_0);
                 assert(rc == 0);
 
-                boot_status_init_by_id(FLASH_AREA_IMAGE_0);
+                boot_status_init_by_id(FLASH_AREA_IMAGE_0, bs);
             }
         }
 
@@ -982,6 +966,9 @@
                 assert(rc == 0);
             }
 
+            rc = boot_write_swap_size(fap, bs->swap_size);
+            assert(rc == 0);
+
             rc = boot_write_magic(fap);
             assert(rc == 0);
 
@@ -1057,42 +1044,43 @@
     struct image_header *hdr;
     uint32_t size;
     uint32_t copy_size;
-    struct image_header tmp_hdr;
     int rc;
 
     /* FIXME: just do this if asked by user? */
 
     size = copy_size = 0;
 
-    hdr = boot_img_hdr(&boot_data, 0);
-    if (hdr->ih_magic == IMAGE_MAGIC) {
-        rc = boot_read_image_size(0, hdr, &copy_size);
-        assert(rc == 0);
-    }
-
-    hdr = boot_img_hdr(&boot_data, 1);
-    if (hdr->ih_magic == IMAGE_MAGIC) {
-        rc = boot_read_image_size(1, hdr, &size);
-        assert(rc == 0);
-    }
-
-    if (!size || !copy_size || size == copy_size) {
-        rc = boot_read_header_from_scratch(&tmp_hdr);
-        assert(rc == 0);
-
-        hdr = &tmp_hdr;
+    if (bs->idx == 0 && bs->state == 0) {
+        /*
+         * No swap ever happened, so need to find the largest image which
+         * will be used to determine the amount of sectors to swap.
+         */
+        hdr = boot_img_hdr(&boot_data, 0);
         if (hdr->ih_magic == IMAGE_MAGIC) {
-            if (!size) {
-                rc = boot_read_image_size(2, hdr, &size);
-            } else {
-                rc = boot_read_image_size(2, hdr, &size);
-            }
+            rc = boot_read_image_size(0, hdr, &copy_size);
             assert(rc == 0);
         }
-    }
 
-    if (size > copy_size) {
-        copy_size = size;
+        hdr = boot_img_hdr(&boot_data, 1);
+        if (hdr->ih_magic == IMAGE_MAGIC) {
+            rc = boot_read_image_size(1, hdr, &size);
+            assert(rc == 0);
+        }
+
+        if (size > copy_size) {
+            copy_size = size;
+        }
+
+        bs->swap_size = copy_size;
+    } else {
+        /*
+         * If a swap was under way, the swap_size should already be present
+         * in the trailer...
+         */
+        rc = boot_read_swap_size(&bs->swap_size);
+        assert(rc == 0);
+
+        copy_size = bs->swap_size;
     }
 
     size = 0;