1 files changed, 28 insertions, 12 deletions

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;