2 files changed, 43 insertions, 4 deletions

diff --git a/mtdutils/mtdutils.c b/mtdutils/mtdutils.c
index 8d32520..18e6a5d 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 8d2cb56..528a5bb 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_