Merge branch 'for-next/misc'

11 files changed, 163 insertions, 104 deletions

diff --git a/common/block.c b/common/block.c
index d90c98948c..97cf5dc4de 100644
--- a/common/block.c
+++ b/common/block.c
@@ -38,6 +38,11 @@ struct chunk {
 
 #define BUFSIZE (PAGE_SIZE * 4)
 
+static int writebuffer_io_len(struct block_device *blk, struct chunk *chunk)
+{
+	return min(blk->rdbufsize, blk->num_blocks - chunk->block_start);
+}
+
 /*
  * Write all dirty chunks back to the device
  */
@@ -51,7 +56,9 @@ static int writebuffer_flush(struct block_device *blk)
 
 	list_for_each_entry(chunk, &blk->buffered_blocks, list) {
 		if (chunk->dirty) {
-			ret = blk->ops->write(blk, chunk->data, chunk->block_start, blk->rdbufsize);
+			ret = blk->ops->write(blk, chunk->data,
+					      chunk->block_start,
+					      writebuffer_io_len(blk, chunk));
 			if (ret < 0)
 				return ret;
 
@@ -76,7 +83,8 @@ static struct chunk *chunk_get_cached(struct block_device *blk, int block)
 	list_for_each_entry(chunk, &blk->buffered_blocks, list) {
 		if (block >= chunk->block_start &&
 				block < chunk->block_start + blk->rdbufsize) {
-			debug("%s: found %d in %d\n", __func__, block, chunk->num);
+			dev_dbg(blk->dev, "%s: found %d in %d\n", __func__,
+				block, chunk->num);
 			/*
 			 * move most recently used entry to the head of the list
 			 */
@@ -117,22 +125,20 @@ static struct chunk *get_chunk(struct block_device *blk)
 		/* use last entry which is the most unused */
 		chunk = list_last_entry(&blk->buffered_blocks, struct chunk, list);
 		if (chunk->dirty) {
-			size_t num_blocks = min(blk->rdbufsize,
-					blk->num_blocks - chunk->block_start);
-			ret = blk->ops->write(blk, chunk->data, chunk->block_start,
-					      num_blocks);
+			ret = blk->ops->write(blk, chunk->data,
+					      chunk->block_start,
+					      writebuffer_io_len(blk, chunk));
 			if (ret < 0)
 				return ERR_PTR(ret);
 
 			chunk->dirty = 0;
 		}
-
-		list_del(&chunk->list);
 	} else {
 		chunk = list_first_entry(&blk->idle_blocks, struct chunk, list);
-		list_del(&chunk->list);
 	}
 
+	list_del(&chunk->list);
+
 	return chunk;
 }
 
@@ -144,7 +150,6 @@ static struct chunk *get_chunk(struct block_device *blk)
 static int block_cache(struct block_device *blk, int block)
 {
 	struct chunk *chunk;
-	size_t num_blocks;
 	int ret;
 
 	chunk = get_chunk(blk);
@@ -153,12 +158,11 @@ static int block_cache(struct block_device *blk, int block)
 
 	chunk->block_start = block & ~blk->blkmask;
 
-	debug("%s: %d to %d\n", __func__, chunk->block_start,
-			chunk->num);
-
-	num_blocks = min(blk->rdbufsize, blk->num_blocks - chunk->block_start);
+	dev_dbg(blk->dev, "%s: %d to %d\n", __func__, chunk->block_start,
+		chunk->num);
 
-	ret = blk->ops->read(blk, chunk->data, chunk->block_start, num_blocks);
+	ret = blk->ops->read(blk, chunk->data, chunk->block_start,
+			     writebuffer_io_len(blk, chunk));
 	if (ret) {
 		list_add_tail(&chunk->list, &blk->idle_blocks);
 		return ret;
@@ -329,13 +333,6 @@ static ssize_t block_op_write(struct cdev *cdev, const void *buf, size_t count,
 }
 #endif
 
-static int block_op_close(struct cdev *cdev)
-{
-	struct block_device *blk = cdev->priv;
-
-	return writebuffer_flush(blk);
-}
-
 static int block_op_flush(struct cdev *cdev)
 {
 	struct block_device *blk = cdev->priv;
@@ -343,6 +340,8 @@ static int block_op_flush(struct cdev *cdev)
 	return writebuffer_flush(blk);
 }
 
+static int block_op_close(struct cdev *cdev) __alias(block_op_flush);
+
 static struct cdev_operations block_ops = {
 	.read	= block_op_read,
 #ifdef CONFIG_BLOCK_WRITE
@@ -368,8 +367,8 @@ int blockdevice_register(struct block_device *blk)
 	INIT_LIST_HEAD(&blk->idle_blocks);
 	blk->blkmask = blk->rdbufsize - 1;
 
-	debug("%s: rdbufsize: %d blockbits: %d blkmask: 0x%08x\n", __func__, blk->rdbufsize, blk->blockbits,
-			blk->blkmask);
+	dev_dbg(blk->dev, "rdbufsize: %d blockbits: %d blkmask: 0x%08x\n",
+		blk->rdbufsize, blk->blockbits, blk->blkmask);
 
 	for (i = 0; i < 32; i++) {
 		struct chunk *chunk = xzalloc(sizeof(*chunk));
diff --git a/common/boot.c b/common/boot.c
index 41bf1ce64b..974eaf5d02 100644
--- a/common/boot.c
+++ b/common/boot.c
@@ -96,8 +96,8 @@ static int bootscript_boot(struct bootentry *entry, int verbose, int dryrun)
 
 	ret = run_command(bs->scriptpath);
 	if (ret) {
-		printf("Running %s failed\n", bs->scriptpath);
-		goto out;
+		pr_err("Running script '%s' failed: %s\n", bs->scriptpath, strerror(-ret));
+		return ret;
 	}
 
 	bootm_data_init_defaults(&data);
@@ -107,11 +107,7 @@ static int bootscript_boot(struct bootentry *entry, int verbose, int dryrun)
 	if (dryrun)
 		data.dryrun = dryrun;
 
-	ret = bootm_boot(&data);
-	if (ret)
-		pr_err("Booting '%s' failed: %s\n", basename(bs->scriptpath), strerror(-ret));
-out:
-	return ret;
+	return bootm_boot(&data);
 }
 
 static unsigned int boot_watchdog_timeout;
@@ -135,7 +131,7 @@ int boot_entry(struct bootentry *be, int verbose, int dryrun)
 {
 	int ret;
 
-	printf("booting '%s'\n", be->title);
+	printf("Booting entry '%s'\n", be->title);
 
 	if (IS_ENABLED(CONFIG_WATCHDOG) && boot_watchdog_timeout) {
 		ret = watchdog_set_timeout(boot_watchdog_timeout);
@@ -144,9 +140,8 @@ int boot_entry(struct bootentry *be, int verbose, int dryrun)
 	}
 
 	ret = be->boot(be, verbose, dryrun);
-
 	if (ret)
-		printf("booting '%s' failed: %s\n", be->title, strerror(-ret));
+		pr_err("Booting entry '%s' failed\n", be->title);
 
 	return ret;
 }
@@ -154,11 +149,8 @@ int boot_entry(struct bootentry *be, int verbose, int dryrun)
 static void bootsource_action(struct menu *m, struct menu_entry *me)
 {
 	struct bootentry *be = container_of(me, struct bootentry, me);
-	int ret;
 
-	ret = boot_entry(be, 0, 0);
-	if (ret)
-		printf("Booting failed with: %s\n", strerror(-ret));
+	boot_entry(be, 0, 0);
 
 	printf("Press any key to continue\n");
 
diff --git a/common/image-fit.c b/common/image-fit.c
index 87a55b7e27..6ac4644686 100644
--- a/common/image-fit.c
+++ b/common/image-fit.c
@@ -368,7 +368,7 @@ static int fit_verify_signature(struct device_node *sig_node, const void *fit)
  out_sl:
 	string_list_free(&inc_nodes);
 	string_list_free(&exc_props);
- 
+
 	return ret;
 }
 
diff --git a/common/memory.c b/common/memory.c
index 00fa7c50ff..21b2b4f63a 100644
--- a/common/memory.c
+++ b/common/memory.c
@@ -224,39 +224,78 @@ int memory_bank_first_find_space(resource_size_t *retstart,
 
 #ifdef CONFIG_OFTREE
 
-static int of_memory_fixup(struct device_node *node, void *unused)
+static int of_memory_fixup(struct device_node *root, void *unused)
 {
 	struct memory_bank *bank;
 	int err;
-	int addr_cell_len, size_cell_len, len = 0;
-	struct device_node *memnode;
-	u8 tmp[16 * 16]; /* Up to 64-bit address + 64-bit size */
+	int addr_cell_len, size_cell_len;
+	struct device_node *memnode, *tmp, *np;
+	char *memnode_name;
 
-	memnode = of_create_node(node, "/memory");
-	if (!memnode)
-		return -ENOMEM;
+	/*
+	 * Since kernel 4.16 the memory node got a @<reg> suffix. To support
+	 * the old and the new style delete any found memory node and add it
+	 * again to be sure that the memory node exists only once. It shouldn't
+	 * bother older kernels if the memory node has this suffix so adding it
+	 * following the new style.
+	 */
 
-	err = of_property_write_string(memnode, "device_type", "memory");
-	if (err)
-		return err;
+	for_each_child_of_node_safe(root, tmp, np) {
+		const char *device_type;
 
-	addr_cell_len = of_n_addr_cells(memnode);
-	size_cell_len = of_n_size_cells(memnode);
+		err = of_property_read_string(np, "device_type", &device_type);
+		if (err || of_node_cmp("memory", device_type))
+			continue;
+
+		/* delete every found memory node */
+		of_delete_node(np);
+	}
+
+	addr_cell_len = of_n_addr_cells(root);
+	size_cell_len = of_n_size_cells(root);
 
 	for_each_memory_bank(bank) {
-		of_write_number(tmp + len, bank->start, addr_cell_len);
+		u8 tmp[16]; /* Up to 64-bit address + 64-bit size */
+		int len = 0;
+
+		/* Create a /memory node for each bank */
+		memnode_name = basprintf("/memory@%lx", bank->start);
+		if (!memnode_name) {
+			err = -ENOMEM;
+			goto err_out;
+		}
+
+		memnode = of_create_node(root, memnode_name);
+		if (!memnode) {
+			err = -ENOMEM;
+			goto err_free;
+		}
+
+		err = of_property_write_string(memnode, "device_type",
+					       "memory");
+		if (err)
+			goto err_free;
+
+		of_write_number(tmp, bank->start, addr_cell_len);
 		len += addr_cell_len * 4;
 		of_write_number(tmp + len, bank->size, size_cell_len);
 		len += size_cell_len * 4;
-	}
 
-	err = of_set_property(memnode, "reg", tmp, len, 1);
-	if (err) {
-		pr_err("could not set reg %s.\n", strerror(-err));
-		return err;
+		err = of_set_property(memnode, "reg", tmp, len, 1);
+		if (err) {
+			pr_err("could not set reg %s.\n", strerror(-err));
+			goto err_free;
+		}
+
+		free(memnode_name);
 	}
 
 	return 0;
+
+err_free:
+	free(memnode_name);
+err_out:
+	return err;
 }
 
 static int of_register_memory_fixup(void)
diff --git a/common/state/backend_bucket_circular.c b/common/state/backend_bucket_circular.c
index 4b71d8751d..da7c8421ae 100644
--- a/common/state/backend_bucket_circular.c
+++ b/common/state/backend_bucket_circular.c
@@ -23,6 +23,10 @@
 #include <mtd/mtd-peb.h>
 #include <string.h>
 
+#ifndef __BAREBOX__
+#include <sys/param.h>
+#endif
+
 #include "state.h"
 
 /*
@@ -161,17 +165,17 @@ static int state_mtd_peb_read(struct state_backend_storage_bucket_circular *circ
 	ret = lseek(circ->fd, offset, SEEK_SET);
 	if (ret < 0) {
 		dev_err(circ->dev, "Failed to set circular read position to %lld, %d\n",
-			offset, ret);
+			(long long) offset, ret);
 		return ret;
 	}
 
-	dev_dbg(circ->dev, "Read state from %ld length %zd\n", offset,
+	dev_dbg(circ->dev, "Read state from %lld length %d\n", (long long) offset,
 		len);
 
 
 	ret = read_full(circ->fd, buf, len);
 	if (ret < 0) {
-		dev_err(circ->dev, "Failed to read circular storage len %zd, %d\n",
+		dev_err(circ->dev, "Failed to read circular storage len %d, %d\n",
 			len, ret);
 		free(buf);
 	}
@@ -189,15 +193,15 @@ static int state_mtd_peb_write(struct state_backend_storage_bucket_circular *cir
 
 	ret = lseek(circ->fd, offset, SEEK_SET);
 	if (ret < 0) {
-		dev_err(circ->dev, "Failed to set position for circular write %ld, %d\n",
-			offset, ret);
+		dev_err(circ->dev, "Failed to set position for circular write %lld, %d\n",
+			(long long) offset, ret);
 		return ret;
 	}
 
 	ret = write_full(circ->fd, buf, len);
 	if (ret < 0) {
-		dev_err(circ->dev, "Failed to write circular to %ld length %zd, %d\n",
-			offset, len, ret);
+		dev_err(circ->dev, "Failed to write circular to %lld length %d, %d\n",
+			(long long) offset, len, ret);
 		return ret;
 	}
 
@@ -207,8 +211,8 @@ static int state_mtd_peb_write(struct state_backend_storage_bucket_circular *cir
 	 */
 	flush(circ->fd);
 
-	dev_dbg(circ->dev, "Written state to offset %ld length %zd data length %zd\n",
-		offset, len, len);
+	dev_dbg(circ->dev, "Written state to offset %lld length %d data length %d\n",
+		(long long) offset, len, len);
 
 	return 0;
 }
@@ -265,8 +269,8 @@ static int state_backend_bucket_circular_read(struct state_backend_storage_bucke
 	if (!buf)
 		return -ENOMEM;
 
-	dev_dbg(circ->dev, "Read state from PEB %u global offset %ld length %zd\n",
-		circ->eraseblock, offset, read_len);
+	dev_dbg(circ->dev, "Read state from PEB %u global offset %lld length %zd\n",
+		circ->eraseblock, (long long) offset, read_len);
 
 	ret = state_mtd_peb_read(circ, buf, offset, read_len);
 	if (ret < 0 && ret != -EUCLEAN) {
@@ -298,7 +302,7 @@ static int state_backend_bucket_circular_write(struct state_backend_storage_buck
 	void *write_buf;
 
 	if (written_length > circ->max_size) {
-		dev_err(circ->dev, "Error, state data too big to be written, to write: %zd, writesize: %zd, length: %zd, available: %zd\n",
+		dev_err(circ->dev, "Error, state data too big to be written, to write: %u, writesize: %zd, length: %zd, available: %zd\n",
 			written_length, circ->writesize, len, circ->max_size);
 		return -E2BIG;
 	}
@@ -345,13 +349,13 @@ static int state_backend_bucket_circular_write(struct state_backend_storage_buck
 
 	ret = state_mtd_peb_write(circ, write_buf, offset, written_length);
 	if (ret < 0 && ret != -EUCLEAN) {
-		dev_err(circ->dev, "Failed to write circular to %ld length %zd, %d\n",
-			offset, written_length, ret);
+		dev_err(circ->dev, "Failed to write circular to %lld length %u, %d\n",
+			(long long) offset, written_length, ret);
 		goto out_free;
 	}
 
-	dev_dbg(circ->dev, "Written state to PEB %u offset %ld length %zd data length %zd\n",
-		circ->eraseblock, offset, written_length, len);
+	dev_dbg(circ->dev, "Written state to PEB %u offset %lld length %u data length %zd\n",
+		circ->eraseblock, (long long) offset, written_length, len);
 
 out_free:
 	free(write_buf);
@@ -445,8 +449,8 @@ static int bucket_circular_is_block_bad(struct state_backend_storage_bucket_circ
 
 	ret = ioctl(circ->fd, MEMGETBADBLOCK, &offs);
 	if (ret < 0)
-		dev_err(circ->dev, "Failed to use ioctl to check for bad block at offset %ld, %d\n",
-			offs, ret);
+		dev_err(circ->dev, "Failed to use ioctl to check for bad block at offset %lld, %d\n",
+			(long long) offs, ret);
 
 	return ret;
 }
diff --git a/common/state/backend_bucket_direct.c b/common/state/backend_bucket_direct.c
index 9d6a337e66..1f00b0fb2f 100644
--- a/common/state/backend_bucket_direct.c
+++ b/common/state/backend_bucket_direct.c
@@ -110,6 +110,9 @@ static int state_backend_bucket_direct_write(struct state_backend_storage_bucket
 	int ret;
 	struct state_backend_storage_bucket_direct_meta meta;
 
+	if (len > direct->max_size - sizeof(meta))
+		return -E2BIG;
+
 	ret = lseek(direct->fd, direct->offset, SEEK_SET);
 	if (ret < 0) {
 		dev_err(direct->dev, "Failed to seek file, %d\n", ret);
diff --git a/common/state/backend_format_raw.c b/common/state/backend_format_raw.c
index 4369f76e37..5a71149d34 100644
--- a/common/state/backend_format_raw.c
+++ b/common/state/backend_format_raw.c
@@ -127,7 +127,7 @@ static int backend_format_raw_verify(struct state_backend_format *format,
 		return -EINVAL;
 	}
 
-	if (backend_raw->algo && !(flags & STATE_FLAG_NO_AUTHENTIFICATION)) {
+	if (backend_raw->algo && !(flags & STATE_FLAG_NO_AUTHENTICATION)) {
 		ret = backend_raw_digest_init(backend_raw);
 		if (ret)
 			return ret;
@@ -153,7 +153,7 @@ static int backend_format_raw_verify(struct state_backend_format *format,
 
 	*lenp = header->data_len + sizeof(*header);
 
-	if (backend_raw->algo && !(flags & STATE_FLAG_NO_AUTHENTIFICATION)) {
+	if (backend_raw->algo && !(flags & STATE_FLAG_NO_AUTHENTICATION)) {
 		const void *hmac = data + header->data_len;
 
 		/* hmac over header and data */
diff --git a/common/state/backend_storage.c b/common/state/backend_storage.c
index c6ebe86244..fca887e93f 100644
--- a/common/state/backend_storage.c
+++ b/common/state/backend_storage.c
@@ -19,7 +19,7 @@
 #include <linux/kernel.h>
 #include <linux/list.h>
 #include <linux/mtd/mtd-abi.h>
-#include <linux/stat.h>
+#include <sys/stat.h>
 #include <malloc.h>
 #include <printk.h>
 
@@ -111,11 +111,11 @@ refresh:
 	ret = bucket->write(bucket, buf, len);
 
 	if (ret) {
-		dev_warn(storage->dev, "Failed to restore bucket %d@0x%08lx\n",
-			 bucket->num, bucket->offset);
+		dev_warn(storage->dev, "Failed to restore bucket %d@0x%08llx\n",
+			 bucket->num, (long long) bucket->offset);
 	} else {
-		dev_info(storage->dev, "restored bucket %d@0x%08lx\n",
-			 bucket->num, bucket->offset);
+		dev_info(storage->dev, "restored bucket %d@0x%08llx\n",
+			 bucket->num, (long long) bucket->offset);
 		bucket->needs_refresh = 0;
 	}
 
@@ -166,7 +166,7 @@ int state_storage_read(struct state_backend_storage *storage,
 		if (!ret && !bucket_used)
 			bucket_used = bucket;
 		if (ret)
-			dev_info(storage->dev, "Ignoring broken bucket %d@0x%08lx...\n", bucket->num, bucket->offset);
+			dev_info(storage->dev, "Ignoring broken bucket %d@0x%08llx...\n", bucket->num, (long long) bucket->offset);
 	}
 
 	dev_dbg(storage->dev, "Checking redundant buckets finished.\n");
@@ -177,7 +177,7 @@ int state_storage_read(struct state_backend_storage *storage,
 		return -ENOENT;
 	}
 
-	dev_info(storage->dev, "Using bucket %d@0x%08lx\n", bucket_used->num, bucket_used->offset);
+	dev_info(storage->dev, "Using bucket %d@0x%08llx\n", bucket_used->num, (long long) bucket_used->offset);
 
 	/*
 	 * Restore/refresh all buckets except the one we currently use (in case
@@ -252,8 +252,8 @@ static int state_storage_mtd_buckets_init(struct state_backend_storage *storage,
 		end = meminfo->size;
 
 	if (!IS_ALIGNED(storage->offset, meminfo->erasesize)) {
-		dev_err(storage->dev, "Offset within the device is not aligned to eraseblocks. Offset is %ld, erasesize %u\n",
-			storage->offset, meminfo->erasesize);
+		dev_err(storage->dev, "Offset within the device is not aligned to eraseblocks. Offset is %lld, erasesize %u\n",
+			(long long) storage->offset, meminfo->erasesize);
 		return -EINVAL;
 	}
 
@@ -326,8 +326,8 @@ static int state_storage_file_buckets_init(struct state_backend_storage *storage
 							 &bucket, offset,
 							 stridesize);
 		if (ret) {
-			dev_warn(storage->dev, "Failed to create direct bucket at '%s' offset %ld\n",
-				 storage->path, offset);
+			dev_warn(storage->dev, "Failed to create direct bucket at '%s' offset %lld\n",
+				 storage->path, (long long) offset);
 			continue;
 		}
 
diff --git a/common/state/state.c b/common/state/state.c
index 54c57232e6..3f5d43ecbf 100644
--- a/common/state/state.c
+++ b/common/state/state.c
@@ -128,7 +128,7 @@ int state_load(struct state *state)
 
 int state_load_no_auth(struct state *state)
 {
-	return state_do_load(state, STATE_FLAG_NO_AUTHENTIFICATION);
+	return state_do_load(state, STATE_FLAG_NO_AUTHENTICATION);
 }
 
 static int state_format_init(struct state *state, const char *backend_format,
@@ -596,6 +596,8 @@ struct state *state_new_from_node(struct device_node *node, bool readonly)
 	const char *alias;
 	uint32_t stridesize;
 	struct device_node *partition_node;
+	off_t offset = 0;
+	size_t size = 0;
 
 	alias = of_alias_get(node);
 	if (!alias) {
@@ -614,7 +616,11 @@ struct state *state_new_from_node(struct device_node *node, bool readonly)
 		goto out_release_state;
 	}
 
+#ifdef __BAREBOX__
 	ret = of_find_path_by_node(partition_node, &state->backend_path, 0);
+#else
+	ret = of_get_devicepath(partition_node, &state->backend_path, &offset, &size);
+#endif
 	if (ret) {
 		if (ret != -EPROBE_DEFER)
 			dev_err(&state->dev, "state failed to parse path to backend: %s\n",
@@ -645,8 +651,8 @@ struct state *state_new_from_node(struct device_node *node, bool readonly)
 	if (ret)
 		goto out_release_state;
 
-	ret = state_storage_init(state, state->backend_path, 0,
-				 0, stridesize, storage_type);
+	ret = state_storage_init(state, state->backend_path, offset,
+				 size, stridesize, storage_type);
 	if (ret)
 		goto out_release_state;
 
diff --git a/common/state/state.h b/common/state/state.h
index 3a0662fd25..912d6d4848 100644
--- a/common/state/state.h
+++ b/common/state/state.h
@@ -6,7 +6,7 @@ struct state;
 struct mtd_info_user;
 
 enum state_flags {
-	STATE_FLAG_NO_AUTHENTIFICATION = (1 << 0),
+	STATE_FLAG_NO_AUTHENTICATION = (1 << 0),
 };
 
 enum state_variable_type {
diff --git a/common/ubiformat.c b/common/ubiformat.c
index 0811525bd2..1968bd98f8 100644
--- a/common/ubiformat.c
+++ b/common/ubiformat.c
@@ -710,8 +710,10 @@ int ubiformat_write(struct mtd_info *mtd, const void *buf, size_t count,
 		    loff_t offset)
 {
 	int writesize = mtd->writesize >> mtd->subpage_sft;
-	size_t retlen;
 	int ret;
+	int peb = 0;
+	int n_skip_blocks = mtd_div_by_eb(offset, mtd);
+	int offset_in_peb = mtd_mod_by_eb(offset, mtd);
 
 	if (offset & (mtd->writesize - 1))
 		return -EINVAL;
@@ -719,27 +721,44 @@ int ubiformat_write(struct mtd_info *mtd, const void *buf, size_t count,
 	if (count & (mtd->writesize - 1))
 		return -EINVAL;
 
+	/* Seek forward to the first PEB we actually want to write */
+	while (1) {
+		ret = mtd_skip_bad(mtd, &peb);
+		if (ret)
+			return ret;
+
+		if (!n_skip_blocks)
+			break;
+
+		peb++;
+		n_skip_blocks--;
+	}
+
 	while (count) {
-		size_t now;
+		size_t now = mtd->erasesize - offset_in_peb;
 
-		now = ALIGN(offset, mtd->erasesize) - offset;
 		if (now > count)
 			now = count;
 
-		if (!now) {
+		if (!offset_in_peb) {
 			const struct ubi_ec_hdr *ec = buf;
 			const struct ubi_vid_hdr *vid;
 
+			ret = mtd_skip_bad(mtd, &peb);
+			if (ret)
+				return ret;
+
 			if (be32_to_cpu(ec->magic) != UBI_EC_HDR_MAGIC) {
 				pr_err("bad UBI magic %#08x, should be %#08x",
 					be32_to_cpu(ec->magic), UBI_EC_HDR_MAGIC);
 				return -EINVAL;
 			}
 
-			/* skip ec header */
-			offset += writesize;
+			/* skip ec header in both flash and image */
+			offset_in_peb = writesize;
 			buf += writesize;
 			count -= writesize;
+			now -= writesize;
 
 			if (!count)
 				break;
@@ -750,19 +769,16 @@ int ubiformat_write(struct mtd_info *mtd, const void *buf, size_t count,
 				       be32_to_cpu(vid->magic), UBI_VID_HDR_MAGIC);
 				return -EINVAL;
 			}
-
-			continue;
 		}
 
-		ret = mtd_write(mtd, offset, now, &retlen, buf);
+		ret = mtd_peb_write(mtd, buf, peb, offset_in_peb, now);
 		if (ret < 0)
 			return ret;
-		if (retlen != now)
-			return -EIO;
 
 		buf += now;
 		count -= now;
-		offset += now;
+		offset_in_peb = 0;
+		peb++;
 	}
 
 	return 0;