3 files changed, 51 insertions, 8 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;
diff --git a/mtdutils/mtdutils.c b/mtdutils/mtdutils.c
index 8d3252059..18e6a5d63 100644
--- a/mtdutils/mtdutils.c
+++ b/mtdutils/mtdutils.c
@@ -47,6 +47,10 @@ struct MtdWriteContext {
     char *buffer;
     size_t stored;
     int fd;
+
+    off_t* bad_block_offsets;
+    int bad_block_alloc;
+    int bad_block_count;
 };
 
 typedef struct {
@@ -366,6 +370,10 @@ MtdWriteContext *mtd_write_partition(const MtdPartition *partition)
     MtdWriteContext *ctx = (MtdWriteContext*) malloc(sizeof(MtdWriteContext));
     if (ctx == NULL) return NULL;
 
+    ctx->bad_block_offsets = NULL;
+    ctx->bad_block_alloc = 0;
+    ctx->bad_block_count = 0;
+
     ctx->buffer = malloc(partition->erase_size);
     if (ctx->buffer == NULL) {
         free(ctx);
@@ -386,8 +394,20 @@ MtdWriteContext *mtd_write_partition(const MtdPartition *partition)
     return ctx;
 }
 
-static int write_block(const MtdPartition *partition, int fd, const char *data)
+static void add_bad_block_offset(MtdWriteContext *ctx, off_t pos) {
+    if (ctx->bad_block_count + 1 > ctx->bad_block_alloc) {
+        ctx->bad_block_alloc = (ctx->bad_block_alloc*2) + 1;
+        ctx->bad_block_offsets = realloc(ctx->bad_block_offsets,
+                                         ctx->bad_block_alloc * sizeof(off_t));
+    }
+    ctx->bad_block_offsets[ctx->bad_block_count++] = pos;
+}
+
+static int write_block(MtdWriteContext *ctx, const char *data)
 {
+    const MtdPartition *partition = ctx->partition;
+    int fd = ctx->fd;
+
     off_t pos = lseek(fd, 0, SEEK_CUR);
     if (pos == (off_t) -1) return 1;
 
@@ -395,6 +415,7 @@ static int write_block(const MtdPartition *partition, int fd, const char *data)
     while (pos + size <= (int) partition->size) {
         loff_t bpos = pos;
         if (ioctl(fd, MEMGETBADBLOCK, &bpos) > 0) {
+            add_bad_block_offset(ctx, pos);
             fprintf(stderr, "mtd: not writing bad block at 0x%08lx\n", pos);
             pos += partition->erase_size;
             continue;  // Don't try to erase known factory-bad blocks.
@@ -436,6 +457,7 @@ static int write_block(const MtdPartition *partition, int fd, const char *data)
         }
 
         // Try to erase it once more as we give up on this block
+        add_bad_block_offset(ctx, pos);
         fprintf(stderr, "mtd: skipping write block at 0x%08lx\n", pos);
         ioctl(fd, MEMERASE, &erase_info);
         pos += partition->erase_size;
@@ -461,13 +483,13 @@ ssize_t mtd_write_data(MtdWriteContext *ctx, const char *data, size_t len)
 
         // If a complete block was accumulated, write it
         if (ctx->stored == ctx->partition->erase_size) {
-            if (write_block(ctx->partition, ctx->fd, ctx->buffer)) return -1;
+            if (write_block(ctx, ctx->buffer)) return -1;
             ctx->stored = 0;
         }
 
         // Write complete blocks directly from the user's buffer
         while (ctx->stored == 0 && len - wrote >= ctx->partition->erase_size) {
-            if (write_block(ctx->partition, ctx->fd, data + wrote)) return -1;
+            if (write_block(ctx, data + wrote)) return -1;
             wrote += ctx->partition->erase_size;
         }
     }
@@ -481,7 +503,7 @@ off_t mtd_erase_blocks(MtdWriteContext *ctx, int blocks)
     if (ctx->stored > 0) {
         size_t zero = ctx->partition->erase_size - ctx->stored;
         memset(ctx->buffer + ctx->stored, 0, zero);
-        if (write_block(ctx->partition, ctx->fd, ctx->buffer)) return -1;
+        if (write_block(ctx, ctx->buffer)) return -1;
         ctx->stored = 0;
     }
 
@@ -522,7 +544,23 @@ int mtd_write_close(MtdWriteContext *ctx)
     // Make sure any pending data gets written
     if (mtd_erase_blocks(ctx, 0) == (off_t) -1) r = -1;
     if (close(ctx->fd)) r = -1;
+    free(ctx->bad_block_offsets);
     free(ctx->buffer);
     free(ctx);
     return r;
 }
+
+/* Return the offset of the first good block at or after pos (which
+ * might be pos itself).
+ */
+off_t mtd_find_write_start(MtdWriteContext *ctx, off_t pos) {
+    int i;
+    for (i = 0; i < ctx->bad_block_count; ++i) {
+        if (ctx->bad_block_offsets[i] == pos) {
+            pos += ctx->partition->erase_size;
+        } else if (ctx->bad_block_offsets[i] > pos) {
+            return pos;
+        }
+    }
+    return pos;
+}
diff --git a/mtdutils/mtdutils.h b/mtdutils/mtdutils.h
index 8d2cb56b1..528a5bbe8 100644
--- a/mtdutils/mtdutils.h
+++ b/mtdutils/mtdutils.h
@@ -49,6 +49,7 @@ void mtd_read_close(MtdReadContext *);
 MtdWriteContext *mtd_write_partition(const MtdPartition *);
 ssize_t mtd_write_data(MtdWriteContext *, const char *data, size_t data_len);
 off_t mtd_erase_blocks(MtdWriteContext *, int blocks);  /* 0 ok, -1 for all */
+off_t mtd_find_write_start(MtdWriteContext *ctx, off_t pos);
 int mtd_write_close(MtdWriteContext *);
 
 #endif  // MTDUTILS_H_