1 files changed, 8 insertions, 4 deletions

diff --git a/bootloader.c b/bootloader.c
index de441e1df..bc79bee76 100644
--- a/bootloader.c
+++ b/bootloader.c
@@ -198,7 +198,7 @@ int write_update_for_bootloader(
     header.version = UPDATE_VERSION;
     header.size = header_size;
 
-    header.image_offset = mtd_erase_blocks(write, 0);
+    off_t image_start_pos = mtd_erase_blocks(write, 0);
     header.image_length = update_length;
     if ((int) header.image_offset == -1 ||
         mtd_write_data(write, update, update_length) != update_length) {
@@ -206,6 +206,8 @@ int write_update_for_bootloader(
         mtd_write_close(write);
         return -1;
     }
+    off_t busy_start_pos = mtd_erase_blocks(write, 0);
+    header.image_offset = mtd_find_write_start(write, image_start_pos);
 
     header.bitmap_width = bitmap_width;
     header.bitmap_height = bitmap_height;
@@ -213,7 +215,6 @@ int write_update_for_bootloader(
 
     int bitmap_length = (bitmap_bpp + 7) / 8 * bitmap_width * bitmap_height;
 
-    header.busy_bitmap_offset = mtd_erase_blocks(write, 0);
     header.busy_bitmap_length = busy_bitmap != NULL ? bitmap_length : 0;
     if ((int) header.busy_bitmap_offset == -1 ||
         mtd_write_data(write, busy_bitmap, bitmap_length) != bitmap_length) {
@@ -221,8 +222,9 @@ int write_update_for_bootloader(
         mtd_write_close(write);
         return -1;
     }
+    off_t fail_start_pos = mtd_erase_blocks(write, 0);
+    header.busy_bitmap_offset = mtd_find_write_start(write, busy_start_pos);
 
-    header.fail_bitmap_offset = mtd_erase_blocks(write, 0);
     header.fail_bitmap_length = fail_bitmap != NULL ? bitmap_length : 0;
     if ((int) header.fail_bitmap_offset == -1 ||
         mtd_write_data(write, fail_bitmap, bitmap_length) != bitmap_length) {
@@ -230,6 +232,8 @@ int write_update_for_bootloader(
         mtd_write_close(write);
         return -1;
     }
+    mtd_erase_blocks(write, 0);
+    header.fail_bitmap_offset = mtd_find_write_start(write, fail_start_pos);
 
     /* Write the header last, after all the blocks it refers to, so that
      * when the magic number is installed everything is valid.
@@ -252,7 +256,7 @@ int write_update_for_bootloader(
         return -1;
     }
 
-    if (mtd_erase_blocks(write, 0) != (off_t) header.image_offset) {
+    if (mtd_erase_blocks(write, 0) != image_start_pos) {
         LOGE("Misalignment rewriting %s\n(%s)\n", CACHE_NAME, strerror(errno));
         mtd_write_close(write);
         return -1;