uboot-mtk-20220606: port mtd_*_skip_bad func from uboot-2023

Signed-off-by: Tianling Shen <cnsztl@immortalwrt.org>

uboot-mtk-20220606/board/mediatek/common/boot_helper.c

@@ -120,6 +120,8 @@ int boot_from_mmc_partition(u32 dev, int part, const char *part_name)
 #ifdef CONFIG_MTD
 int boot_from_mtd(struct mtd_info *mtd, u64 offset)
 {
+	u32 fit_hdrsize = sizeof(struct fdt_header);
+	u32 legacy_hdrsize = image_get_header_size();
 	ulong data_load_addr;
 	u32 size;
 	int ret;
@@ -131,8 +133,10 @@ int boot_from_mtd(struct mtd_info *mtd, u64 offset)
 	data_load_addr = CONFIG_LOADADDR;
 #endif
 
-	ret = mtd_read_generic(mtd, offset, (void *)data_load_addr,
-			       mtd->writesize);
+	size = max3(fit_hdrsize, legacy_hdrsize, mtd->writesize);
+
+	ret = mtd_read_skip_bad(mtd, offset, size, mtd->size, NULL,
+			        (void *)data_load_addr);
 	if (ret)
 		return ret;
 
@@ -147,8 +151,8 @@ int boot_from_mtd(struct mtd_info *mtd, u64 offset)
 		size = fit_get_totalsize((const void *)data_load_addr);
 		if (size <= 0x2000) {
 			/* Load FDT header into memory */
-			ret = mtd_read_generic(mtd, offset, (void *)data_load_addr,
-					       mtd->writesize);
+			ret = mtd_read_skip_bad(mtd, offset, size, mtd->size, NULL,
+						(void *)data_load_addr);
 			if (ret)
 				return ret;
 
@@ -163,60 +167,8 @@ int boot_from_mtd(struct mtd_info *mtd, u64 offset)
 		return -EINVAL;
 	}
 
-	ret = mtd_read_generic(mtd, offset, (void *)data_load_addr, size);
-	if (ret)
-		return ret;
-
-	return boot_from_mem(data_load_addr);
-}
-#endif
-
-#ifdef CONFIG_DM_SPI_FLASH
-int boot_from_snor(struct spi_flash *snor, u32 offset)
-{
-	ulong data_load_addr;
-	u32 size;
-	int ret;
-
-	/* Set load address */
-#if defined(CONFIG_SYS_LOAD_ADDR)
-	data_load_addr = CONFIG_SYS_LOAD_ADDR;
-#elif defined(CONFIG_LOADADDR)
-	data_load_addr = CONFIG_LOADADDR;
-#endif
-
-	ret = snor_read_generic(snor, offset, (void *)data_load_addr,
-				snor->page_size);
-	if (ret)
-		return ret;
-
-	switch (genimg_get_format((const void *)data_load_addr)) {
-#if defined(CONFIG_LEGACY_IMAGE_FORMAT)
-	case IMAGE_FORMAT_LEGACY:
-		size = image_get_image_size((const void *)data_load_addr);
-		break;
-#endif
-#if defined(CONFIG_FIT)
-	case IMAGE_FORMAT_FIT:
-		size = fit_get_totalsize((const void *)data_load_addr);
-		if (size <= 0x2000) {
-			/* Load FDT header into memory */
-			ret = snor_read_generic(snor, offset, (void *)data_load_addr,
-					       size);
-			if (ret)
-				return ret;
-
-			/* Read whole FIT image */
-			size = fit_get_totalsize((const void *)data_load_addr);
-		}
-		break;
-#endif
-	default:
-		printf("Error: no Image found in SPI-NOR at 0x%x\n", offset);
-		return -EINVAL;
-	}
-
-	ret = snor_read_generic(snor, offset, (void *)data_load_addr, size);
+	ret = mtd_read_skip_bad(mtd, offset, size, mtd->size, NULL,
+				(void *)data_load_addr);
 	if (ret)
 		return ret;
 

uboot-mtk-20220606/board/mediatek/common/boot_helper.h

@@ -29,11 +29,6 @@ int boot_from_mmc_partition(u32 dev, int part, const char *part_name);
 int boot_from_mtd(struct mtd_info *mtd, u64 offset);
 #endif
 
-#ifdef CONFIG_DM_SPI_FLASH
-#include <spi_flash.h>
-int boot_from_snor(struct spi_flash *snor, u32 offset);
-#endif
-
 struct bootarg {
 	const char *key;
 	const char *value;

uboot-mtk-20220606/board/mediatek/common/ubi_helper.c

@@ -23,50 +23,6 @@
 #include "colored_print.h"
 #include "untar.h"
 
-void ubi_probe_mtd_devices(void)
-{
-#ifdef CONFIG_MEDIATEK_UBI_FIXED_MTDPARTS
-	const char *mtdids = NULL, *mtdparts = NULL;
-
-#if defined(CONFIG_SYS_MTDPARTS_RUNTIME)
-	board_mtdparts_default(&mtdids, &mtdparts);
-#else
-#if defined(MTDIDS_DEFAULT)
-	mtdids = MTDIDS_DEFAULT;
-#elif defined(CONFIG_MTDIDS_DEFAULT)
-	mtdids = CONFIG_MTDIDS_DEFAULT;
-#endif
-
-#if defined(MTDPARTS_DEFAULT)
-	mtdparts = MTDPARTS_DEFAULT;
-#elif defined(CONFIG_MTDPARTS_DEFAULT)
-	mtdparts = CONFIG_MTDPARTS_DEFAULT;
-#endif
-#endif
-
-	if (mtdids)
-		env_set("mtdids", mtdids);
-
-	if (mtdparts)
-		env_set("mtdparts", mtdparts);
-#endif
-
-	mtd_probe_devices();
-}
-
-static int mtd_update_generic(struct mtd_info *mtd, const void *data,
-			      size_t size)
-{
-	int ret;
-
-	/* Write ubi part to kernel MTD partition */
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
-		return ret;
-
-	return mtd_write_generic(mtd, 0, 0, data, size, true);
-}
-
 static int write_ubi1_image(const void *data, size_t size,
 			    struct mtd_info *mtd_kernel,
 			    struct mtd_info *mtd_ubi,
@@ -80,13 +36,13 @@ static int write_ubi1_image(const void *data, size_t size,
 	}
 
 	/* Write kernel part to kernel MTD partition */
-	ret = mtd_update_generic(mtd_kernel, data, mtd_kernel->size);
+	ret = mtd_update_generic(mtd_kernel, data, mtd_kernel->size, true);
 	if (ret)
 		return ret;
 
 	/* Write ubi part to kernel MTD partition */
 	return mtd_update_generic(mtd_ubi, data + mtd_kernel->size,
-				  ii->ubi_size + ii->marker_size);
+				  ii->ubi_size + ii->marker_size, true);
 }
 
 static int mount_ubi(struct mtd_info *mtd)
@@ -109,7 +65,8 @@ static int mount_ubi(struct mtd_info *mtd)
 
 		ubi_mtd_param_parse(mtd->name, NULL);
 
-		ret = mtd_erase_generic(mtd, 0, mtd->size);
+		ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+					 false);
 		if (ret)
 			return ret;
 
@@ -410,7 +367,7 @@ static int write_ubi1_tar_image(const void *data, size_t size,
 		return ret;
 
 	/* Write kernel part to kernel MTD partition */
-	ret = mtd_update_generic(mtd_kernel, kernel_data, kernel_size);
+	ret = mtd_update_generic(mtd_kernel, kernel_data, kernel_size, true);
 	if (ret)
 		return ret;
 
@@ -523,7 +480,7 @@ static int write_ubi_fit_image(const void *data, size_t size,
 	if (!find_ubi_volume("fit")) {
 		/* ubi is dirty, erase ubi and recreate volumes */
 		umount_ubi();
-		ret = mtd_erase_generic(mtd, 0, mtd->size);
+		ret = mtd_erase_skip_bad(mtd, 0, size, mtd->size, NULL, NULL, true);
 		if (ret)
 			return ret;
 
@@ -595,7 +552,7 @@ int ubi_upgrade_image(const void *data, size_t size)
 	const char *ubi_rootfs_part;
 #endif
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 	mtd_kernel = get_mtd_device_nm("kernel");
 	if (!IS_ERR_OR_NULL(mtd_kernel))
@@ -630,7 +587,7 @@ int ubi_upgrade_image(const void *data, size_t size)
 				return write_ubi_fit_image(data, size, mtd);
 
 			if (!ret && ii.type == IMAGE_UBI2)
-				return mtd_update_generic(mtd, data, size);
+				return mtd_update_generic(mtd, data, size, true);
 
 			if (!ret && ii.type == IMAGE_TAR) {
 #ifdef CONFIG_MEDIATEK_MULTI_MTD_LAYOUT
@@ -659,7 +616,7 @@ int ubi_upgrade_image(const void *data, size_t size)
 
 		ret = parse_image_ram(data, size, mtd->erasesize, &ii);
 		if (!ret && (ii.type == IMAGE_RAW || ii.type == IMAGE_UBI1))
-			return mtd_update_generic(mtd, data, size);
+			return mtd_update_generic(mtd, data, size, true);
 	}
 
 	return -ENOTSUPP;
@@ -670,7 +627,7 @@ int ubi_boot_image(void)
 	struct mtd_info *mtd, *mtd_kernel;
 	const char *ubi_boot_part = "ubi";
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 	mtd_kernel = get_mtd_device_nm("kernel");
 	if (!IS_ERR_OR_NULL(mtd_kernel))

uboot-mtk-20220606/board/mediatek/common/ubi_helper.h

@@ -10,7 +10,6 @@
 #ifndef _UBI_HELPER_H_
 #define _UBI_HELPER_H_
 
-void ubi_probe_mtd_devices(void);
 int ubi_upgrade_image(const void *data, size_t size);
 int ubi_boot_image(void);
 

uboot-mtk-20220606/board/mediatek/common/upgrade_helper.c

@@ -13,6 +13,7 @@
 #include <part.h>
 #include <part_efi.h>
 #include <spi.h>
+#include <mtd.h>
 #include <linux/bitops.h>
 #include <linux/sizes.h>
 #include <linux/types.h>
@@ -28,7 +29,7 @@
 /* 16K for max NAND page size + oob size */
 static u8 scratch_buffer[SZ_16K];
 
-static int check_data_size(u64 total_size, u64 offset, size_t max_size,
+int check_data_size(u64 total_size, u64 offset, size_t max_size,
 			   size_t size, bool write)
 {
 	if (offset + size > total_size) {
@@ -51,7 +52,7 @@ abort:
 	return -EINVAL;
 }
 
-static bool verify_data(const u8 *orig, const u8 *rdback, u64 offset,
+bool verify_data(const u8 *orig, const u8 *rdback, u64 offset,
 			size_t size)
 {
 	bool passed = true;
@@ -563,11 +564,43 @@ int mmc_write_gpt(u32 dev, int part, size_t max_size, const void *data,
 #endif /* CONFIG_GENERIC_MMC */
 
 #ifdef CONFIG_MTD
-static int _mtd_erase_generic(struct mtd_info *mtd, u64 offset, u64 size,
-			      const char *name)
+void gen_mtd_probe_devices(void)
 {
+#ifdef CONFIG_MEDIATEK_UBI_FIXED_MTDPARTS
+	const char *mtdids = NULL, *mtdparts = NULL;
+
+#if defined(CONFIG_SYS_MTDPARTS_RUNTIME)
+	board_mtdparts_default(&mtdids, &mtdparts);
+#else
+#if defined(MTDIDS_DEFAULT)
+	mtdids = MTDIDS_DEFAULT;
+#elif defined(CONFIG_MTDIDS_DEFAULT)
+	mtdids = CONFIG_MTDIDS_DEFAULT;
+#endif
+
+#if defined(MTDPARTS_DEFAULT)
+	mtdparts = MTDPARTS_DEFAULT;
+#elif defined(CONFIG_MTDPARTS_DEFAULT)
+	mtdparts = CONFIG_MTDPARTS_DEFAULT;
+#endif
+#endif
+
+	if (mtdids)
+		env_set("mtdids", mtdids);
+
+	if (mtdparts)
+		env_set("mtdparts", mtdparts);
+#endif
+
+	mtd_probe_devices();
+}
+
+int mtd_erase_skip_bad(struct mtd_info *mtd, u64 offset, u64 size,
+		       u64 maxsize, u64 *erasedsize, const char *name,
+		       bool fullerase)
+{
+	u64 start, end, len, limit;
 	struct erase_info ei;
-	u64 start, end;
 	int ret;
 
 	if (!mtd)
@@ -579,32 +612,73 @@ static int _mtd_erase_generic(struct mtd_info *mtd, u64 offset, u64 size,
 	if (offset >= mtd->size) {
 		printf("\n");
 		cprintln(ERROR, "*** Erase offset pasts flash size! ***");
-		return -EINVAL;
+		return -ERANGE;
 	}
 
-	if (offset + size > mtd->size) {
+	if (size > mtd->size - offset) {
 		printf("\n");
 		cprintln(ERROR, "*** Erase size pasts flash size! ***");
-		return -EINVAL;
+		return -ERANGE;
 	}
 
-	start = offset & (~(u64)mtd->erasesize_mask);
-	end = (start + size + mtd->erasesize_mask) &
-	      (~(u64)mtd->erasesize_mask);
+	if (maxsize < size)
+		maxsize = size;
 
-	printf("Erasing%s%s from 0x%llx to 0x%llx, size 0x%llx ... ",
-	       name ? " " : "", name ? name : "", start, end - 1, end - start);
+	if (maxsize > mtd->size - offset)
+		maxsize = mtd->size - offset;
+
+	limit = (offset + maxsize + mtd->erasesize_mask) &
+		(~(u64)mtd->erasesize_mask);
+	end = (offset + size + mtd->erasesize_mask) &
+	      (~(u64)mtd->erasesize_mask);
+	start = offset & (~(u64)mtd->erasesize_mask);
+	len = end - start;
+
+	printf("Erasing '%s' from 0x%llx, size 0x%llx ... ",
+	       name ? name : mtd->name, mtd->offset + start, len);
 
 	memset(&ei, 0, sizeof(ei));
 
 	ei.mtd = mtd;
-	ei.addr = start;
-	ei.len = end - start;
 
-	ret = mtd_erase(mtd, &ei);
-	if (ret) {
-		printf("Fail\n");
-		return ret;
+	if (erasedsize)
+		*erasedsize = 0;
+
+	while (len && start < limit) {
+		ret = mtd_block_isbad(mtd, start);
+		if (ret < 0) {
+			printf("Failed to check bad block at 0x%llx\n",
+			       mtd->offset + start);
+			return ret;
+		}
+
+		if (ret) {
+			printf("Skipped bad block at 0x%llx\n",
+			       mtd->offset + start);
+			start += mtd->erasesize;
+			continue;
+		}
+
+		ei.addr = start;
+		ei.len = mtd->erasesize;
+
+		ret = mtd_erase(mtd, &ei);
+		if (ret) {
+			printf("Failed at 0x%llx, err = %d\n",
+			       mtd->offset + start, ret);
+			return ret;
+		}
+
+		len -= mtd->erasesize;
+		start += mtd->erasesize;
+
+		if (erasedsize)
+			*erasedsize += mtd->erasesize;
+	}
+
+	if (len && fullerase) {
+		printf("No enough blocks erased\n");
+		return -ENODATA;
 	}
 
 	printf("OK\n");
@@ -612,91 +686,79 @@ static int _mtd_erase_generic(struct mtd_info *mtd, u64 offset, u64 size,
 	return 0;
 }
 
-int mtd_erase_generic(struct mtd_info *mtd, u64 offset, u64 size)
+static int mtd_validate_block(struct mtd_info *mtd, u64 addr, size_t size,
+			      const void *data)
 {
-	return _mtd_erase_generic(mtd, offset, size, NULL);
-}
-
-int mtd_write_generic(struct mtd_info *mtd, u64 offset, u64 max_size,
-		      const void *data, size_t size, bool verify)
-{
-	size_t size_left, chksz;
 	struct mtd_oob_ops ops;
-	u64 verify_offset;
+	size_t col, chksz;
 	int ret;
 
-	if (!mtd)
-		return -EINVAL;
+	static u8 *vbuff_ptr = NULL;
+	static size_t vbuff_sz = 0;
 
-	if (!size)
-		return 0;
+	if (mtd->writesize > vbuff_sz) {
+		if (!vbuff_ptr) {
+			vbuff_ptr = malloc(mtd->writesize);
+		} else {
+			u8 *tptr = realloc(vbuff_ptr, mtd->writesize);
+			if (tptr) {
+				vbuff_ptr = tptr;
+			} else {
+				free(vbuff_ptr);
+				vbuff_ptr = NULL;
+			}
+		}
 
-	if (check_data_size(mtd->size, offset, max_size, size, true))
-		return -EINVAL;
+		if (!vbuff_ptr) {
+			cprintln(ERROR,
+				 "*** Insufficient memory for verifying! ***");
+			return -ENOMEM;
+		}
 
-	printf("Writing from 0x%lx to 0x%llx, size 0x%zx ... ", (ulong)data,
-	       offset, size);
+		vbuff_sz = mtd->writesize;
+	}
 
 	memset(&ops, 0, sizeof(ops));
 
-	ops.mode = MTD_OPS_AUTO_OOB;
-	ops.datbuf = (void *)data;
-	ops.len = size;
+	while (size) {
+		col = addr & mtd->writesize_mask;
 
-	ret = mtd_write_oob(mtd, offset, &ops);
-	if (ret) {
-		printf("Fail\n");
-		return ret;
-	}
-
-	printf("OK\n");
-
-	if (!verify)
-		return 0;
-
-	printf("Verifying from 0x%llx to 0x%llx, size 0x%zx ... ", offset,
-	       offset + size - 1, size);
-
-	size_left = size;
-	verify_offset = 0;
-
-	while (size_left) {
-		chksz = min(sizeof(scratch_buffer), size_left);
-
-		if (chksz > mtd->writesize)
-			chksz &= ~(size_t)(mtd->writesize - 1);
-
-		memset(&ops, 0, sizeof(ops));
+		chksz = mtd->writesize - col;
+		if (chksz > size)
+			chksz = size;
 
 		ops.mode = MTD_OPS_AUTO_OOB;
-		ops.datbuf = scratch_buffer;
+		ops.datbuf = vbuff_ptr;
 		ops.len = chksz;
 		ops.retlen = 0;
 
-		ret = mtd_read_oob(mtd, offset + verify_offset, &ops);
-		if (ret < 0 && ret != -EUCLEAN) {
-			printf("Fail (ret = %d)\n", ret);
-			cprintln(ERROR, "*** Only 0x%zx read! ***",
-				 size - size_left + ops.retlen);
-			return -EIO;
+		ret = mtd_read_oob(mtd, addr, &ops);
+		if (ret && ret != -EUCLEAN) {
+			printf("Failed to read at 0x%llx, err = %d\n",
+			       mtd->offset + addr, ret);
+			return ret;
 		}
 
-		if (!verify_data(data + verify_offset, scratch_buffer,
-				 offset + verify_offset, chksz))
-			return -EIO;
+		if (!verify_data(data, vbuff_ptr, mtd->offset + addr, chksz))
+			return -EBADMSG;
 
-		verify_offset += chksz;
-		size_left -= chksz;
+		size -= chksz;
+		addr += chksz;
+		data = (char *)data + chksz;
 	}
 
-	printf("OK\n");
-
 	return 0;
 }
 
-int mtd_read_generic(struct mtd_info *mtd, u64 offset, void *data, size_t size)
+int mtd_write_skip_bad(struct mtd_info *mtd, u64 offset, size_t size,
+		       u64 maxsize, size_t *writtensize, const void *data,
+		       bool verify)
 {
 	struct mtd_oob_ops ops;
+	bool checkbad = true;
+	size_t len, chksz;
+	u64 addr, limit;
+	u32 blockoff;
 	int ret;
 
 	if (!mtd)
@@ -705,158 +767,91 @@ int mtd_read_generic(struct mtd_info *mtd, u64 offset, void *data, size_t size)
 	if (!size)
 		return 0;
 
-	if (check_data_size(mtd->size, offset, 0, size, false))
+	if (offset >= mtd->size) {
+		printf("\n");
+		cprintln(ERROR, "*** Write offset pasts flash size! ***");
+		return -ERANGE;
+	}
+
+	if (maxsize < size)
+		maxsize = size;
+
+	if (maxsize > mtd->size - offset)
+		maxsize = mtd->size - offset;
+
+	if (check_data_size(mtd->size, offset, maxsize, size, true))
 		return -EINVAL;
 
-	printf("Reading from 0x%llx to 0x%lx, size 0x%zx ... ", offset,
-	       (ulong)data, size);
+	printf("Writing '%s' from 0x%lx to 0x%llx, size 0x%zx ... ", mtd->name,
+	       (ulong)data, mtd->offset + offset, size);
+
+	limit = offset + maxsize;
+	addr = offset;
+	len = size;
 
 	memset(&ops, 0, sizeof(ops));
 
-	ops.mode = MTD_OPS_AUTO_OOB;
-	ops.datbuf = (void *)data;
-	ops.len = size;
+	if (writtensize)
+		*writtensize = 0;
 
-	ret = mtd_read_oob(mtd, offset, &ops);
-	if (ret < 0 && ret != -EUCLEAN) {
-		printf("Fail (ret = %d)\n", ret);
-		return ret;
-	}
+	while (len && addr < limit) {
+		if (checkbad || !(addr & mtd->erasesize_mask)) {
+			ret = mtd_block_isbad(mtd, addr);
+			if (ret < 0) {
+				printf("Failed to check bad block at 0x%llx\n",
+				       mtd->offset + addr);
+				return ret;
+			}
 
-	printf("OK\n");
+			if (ret) {
+				printf("Skipped bad block at 0x%llx\n",
+				       mtd->offset + addr);
+				addr += mtd->erasesize;
+				checkbad = true;
+				continue;
+			}
 
-	return 0;
-}
-
-int mtd_erase_env(struct mtd_info *mtd, u64 offset, u64 size)
-{
-	return _mtd_erase_generic(mtd, offset, size, "environment");
-}
-#endif
-
-#ifdef CONFIG_DM_SPI_FLASH
-struct spi_flash *snor_get_dev(int bus, int cs)
-{
-	struct udevice *new, *bus_dev;
-	int ret;
-
-	/* Remove the old device, otherwise probe will just be a nop */
-	ret = spi_find_bus_and_cs(bus, cs, &bus_dev, &new);
-	if (!ret)
-		device_remove(new, DM_REMOVE_NORMAL);
-
-	ret = spi_flash_probe_bus_cs(bus, cs, &new);
-	if (!ret)
-		return dev_get_uclass_priv(new);
-
-	cprintln(ERROR,
-		 "*** Failed to initialize SPI-NOR at %u:%u (error %d) ***",
-		 bus, cs, ret);
-
-	return NULL;
-}
-
-static int _snor_erase_generic(struct spi_flash *snor, u32 offset, u32 size,
-			       const char *name)
-{
-	u32 start, end, erasesize_mask;
-	int ret;
-
-	if (!snor)
-		return -EINVAL;
-
-	if (!size)
-		return 0;
-
-	if (offset >= snor->size) {
-		printf("\n");
-		cprintln(ERROR, "*** Erase offset pasts flash size! ***");
-		return -EINVAL;
-	}
-
-	if (offset + size > snor->size) {
-		printf("\n");
-		cprintln(ERROR, "*** Erase size pasts flash size! ***");
-		return -EINVAL;
-	}
-
-	erasesize_mask = (1 << (ffs(snor->erase_size) - 1)) - 1;
-
-	start = offset & (~erasesize_mask);
-	end = (start + size + erasesize_mask) & (~erasesize_mask);
-
-	printf("Erasing%s%s from 0x%x to 0x%x, size 0x%x ... ",
-	       name ? " " : "", name ? name : "", start, end - 1, end - start);
-
-	ret = spi_flash_erase(snor, start, end - start);
-	if (ret) {
-		printf("Fail\n");
-		return ret;
-	}
-
-	printf("OK\n");
-
-	return 0;
-}
-
-int snor_erase_generic(struct spi_flash *snor, u32 offset, u32 size)
-{
-	return _snor_erase_generic(snor, offset, size, NULL);
-}
-
-int snor_write_generic(struct spi_flash *snor, u32 offset, u32 max_size,
-		       const void *data, size_t size, bool verify)
-{
-	size_t size_left, chksz;
-	u32 verify_offset;
-	int ret;
-
-	if (!snor)
-		return -EINVAL;
-
-	if (!size)
-		return 0;
-
-	if (check_data_size(snor->size, offset, max_size, size, true))
-		return -EINVAL;
-
-	printf("Writing from 0x%lx to 0x%x, size 0x%zx ... ", (ulong)data,
-	       offset, size);
-
-	ret = spi_flash_write(snor, offset, size, data);
-	if (ret) {
-		printf("Fail\n");
-		return ret;
-	}
-
-	printf("OK\n");
-
-	if (!verify)
-		return 0;
-
-	printf("Verifying from 0x%x to 0x%zx, size 0x%zx ... ", offset,
-	       offset + size - 1, size);
-
-	size_left = size;
-	verify_offset = 0;
-
-	while (size_left) {
-		chksz = min(sizeof(scratch_buffer), size_left);
-
-		ret = spi_flash_read(snor, offset + verify_offset, chksz,
-				     scratch_buffer);
-		if (ret) {
-			printf("Fail\n");
-			cprintln(ERROR, "*** Failed to read SPI-NOR! ***");
-			return -EIO;
+			checkbad = false;
 		}
 
-		if (!verify_data(data + verify_offset, scratch_buffer,
-				 offset + verify_offset, chksz))
-			return -EIO;
+		blockoff = addr & mtd->erasesize_mask;
 
-		verify_offset += chksz;
-		size_left -= chksz;
+		chksz = mtd->erasesize - blockoff;
+		if (chksz > len)
+			chksz = len;
+
+		if (addr + chksz > limit)
+			chksz = limit - addr;
+
+		ops.mode = MTD_OPS_AUTO_OOB;
+		ops.datbuf = (void *)data;
+		ops.len = chksz;
+		ops.retlen = 0;
+
+		ret = mtd_write_oob(mtd, addr, &ops);
+		if (ret) {
+			printf("Failed at 0x%llx, err = %d\n",
+			       mtd->offset + addr, ret);
+			return ret;
+		}
+
+		if (verify) {
+			ret = mtd_validate_block(mtd, addr, chksz, data);
+			if (ret)
+				return ret;
+		}
+
+		addr += ops.retlen;
+		len -= ops.retlen;
+		data += ops.retlen;
+
+		if (writtensize)
+			*writtensize += ops.retlen;
+	}
+
+	if (len) {
+		printf("Incomplete write\n");
+		return -ENODATA;
 	}
 
 	printf("OK\n");
@@ -864,27 +859,101 @@ int snor_write_generic(struct spi_flash *snor, u32 offset, u32 max_size,
 	return 0;
 }
 
-int snor_read_generic(struct spi_flash *snor, u32 offset, void *data,
-		      size_t size)
+int mtd_read_skip_bad(struct mtd_info *mtd, u64 offset, size_t size,
+		      u64 maxsize, size_t *readsize, void *data)
 {
+	struct mtd_oob_ops ops;
+	bool checkbad = true;
+	size_t len, chksz;
+	u64 addr, limit;
+	u32 blockoff;
 	int ret;
 
-	if (!snor)
+	if (!mtd)
 		return -EINVAL;
 
 	if (!size)
 		return 0;
 
-	if (check_data_size(snor->size, offset, 0, size, false))
+	if (offset >= mtd->size) {
+		printf("\n");
+		cprintln(ERROR, "*** Read offset pasts flash size! ***");
+		return -ERANGE;
+	}
+
+	if (maxsize < size)
+		maxsize = size;
+
+	if (maxsize > mtd->size - offset)
+		maxsize = mtd->size - offset;
+
+	if (check_data_size(mtd->size, offset, maxsize, size, false))
 		return -EINVAL;
 
-	printf("Reading from 0x%x to 0x%lx, size 0x%zx ... ", offset,
-	       (ulong)data, size);
+	printf("Reading '%s' from 0x%llx to 0x%lx, size 0x%zx ... ", mtd->name,
+	       mtd->offset + offset, (ulong)data, size);
 
-	ret = spi_flash_read(snor, offset, size, data);
-	if (ret) {
-		printf("Fail\n");
-		return ret;
+	limit = offset + maxsize;
+	addr = offset;
+	len = size;
+
+	memset(&ops, 0, sizeof(ops));
+
+	if (readsize)
+		*readsize = 0;
+
+	while (len && addr < limit) {
+		if (checkbad || !(addr & mtd->erasesize_mask)) {
+			ret = mtd_block_isbad(mtd, addr);
+			if (ret < 0) {
+				printf("Failed to check bad block at 0x%llx\n",
+				       mtd->offset + addr);
+				return ret;
+			}
+
+			if (ret) {
+				printf("Skipped bad block at 0x%llx\n",
+				       mtd->offset + addr);
+				addr += mtd->erasesize;
+				checkbad = true;
+				continue;
+			}
+
+			checkbad = false;
+		}
+
+		blockoff = addr & mtd->erasesize_mask;
+
+		chksz = mtd->erasesize - blockoff;
+		if (chksz > len)
+			chksz = len;
+
+		if (addr + chksz > limit)
+			chksz = limit - addr;
+
+		ops.mode = MTD_OPS_AUTO_OOB;
+		ops.datbuf = data;
+		ops.len = chksz;
+		ops.retlen = 0;
+
+		ret = mtd_read_oob(mtd, addr, &ops);
+		if (ret && ret != -EUCLEAN) {
+			printf("Failed at 0x%llx, err = %d\n",
+			       mtd->offset + addr, ret);
+			return ret;
+		}
+
+		addr += ops.retlen;
+		len -= ops.retlen;
+		data += ops.retlen;
+
+		if (readsize)
+			*readsize += ops.retlen;
+	}
+
+	if (len) {
+		printf("Incomplete read\n");
+		return -ENODATA;
 	}
 
 	printf("OK\n");
@@ -892,8 +961,15 @@ int snor_read_generic(struct spi_flash *snor, u32 offset, void *data,
 	return 0;
 }
 
-int snor_erase_env(struct spi_flash *snor, u32 offset, u32 size)
+int mtd_update_generic(struct mtd_info *mtd, const void *data, size_t size,
+		       bool verify)
 {
-	return _snor_erase_generic(snor, offset, size, "environment");
+	int ret;
+
+	ret = mtd_erase_skip_bad(mtd, 0, size, mtd->size, NULL, NULL, true);
+	if (ret)
+		return ret;
+
+	return mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, verify);
 }
 #endif

uboot-mtk-20220606/board/mediatek/common/upgrade_helper.h

@@ -72,22 +72,21 @@ int mmc_write_gpt(u32 dev, int part, size_t max_size, const void *data,
 
 #ifdef CONFIG_MTD
 #include <linux/mtd/mtd.h>
-int mtd_erase_generic(struct mtd_info *mtd, u64 offset, u64 size);
-int mtd_write_generic(struct mtd_info *mtd, u64 offset, u64 max_size,
-		      const void *data, size_t size, bool verify);
-int mtd_read_generic(struct mtd_info *mtd, u64 offset, void *data, size_t size);
-int mtd_erase_env(struct mtd_info *mtd, u64 offset, u64 size);
+void gen_mtd_probe_devices(void);
+int mtd_erase_skip_bad(struct mtd_info *mtd, u64 offset, u64 size,
+		       u64 maxsize, u64 *erasedsize, const char *name,
+		       bool fullerase);
+int mtd_write_skip_bad(struct mtd_info *mtd, u64 offset, size_t size,
+		       u64 maxsize, size_t *writtensize, const void *data,
+		       bool verify);
+int mtd_read_skip_bad(struct mtd_info *mtd, u64 offset, size_t size,
+		      u64 maxsize, size_t *readsize, void *data);
+int mtd_update_generic(struct mtd_info *mtd, const void *data, size_t size,
+		       bool verify);
 #endif
 
-#ifdef CONFIG_DM_SPI_FLASH
-#include <spi_flash.h>
-struct spi_flash *snor_get_dev(int bus, int cs);
-int snor_erase_generic(struct spi_flash *snor, u32 offset, u32 size);
-int snor_write_generic(struct spi_flash *snor, u32 offset, u32 max_size,
-		       const void *data, size_t size, bool verify);
-int snor_read_generic(struct spi_flash *snor, u32 offset, void *data,
-		      size_t size);
-int snor_erase_env(struct spi_flash *snor, u32 offset, u32 size);
-#endif
+int check_data_size(u64 total_size, u64 offset, size_t max_size, size_t size,
+		    bool write);
+bool verify_data(const u8 *orig, const u8 *rdback, u64 offset, size_t size);
 
 #endif /* _UPGRADE_HELPER_H_ */

uboot-mtk-20220606/board/mediatek/common/verify_helper.c

@@ -660,61 +660,3 @@ int mmc_boot_verify(struct mmc_verify_data *vd)
 	return 0;
 }
 #endif /* CONFIG_MMC */
-
-#ifdef CONFIG_DM_SPI_FLASH
-struct snor_image_read_priv {
-	struct image_read_priv p;
-	struct spi_flash *snor;
-};
-
-static int snor_image_read(struct image_read_priv *rpriv, void *buff, u64 addr,
-			   size_t size)
-{
-	struct snor_image_read_priv *priv =
-		container_of(rpriv, struct snor_image_read_priv, p);
-
-	return snor_read_generic(priv->snor, (u32)addr, buff, size);
-}
-
-int snor_boot_verify(struct snor_verify_data *vd)
-{
-	struct snor_image_read_priv read_priv;
-	void *rootfs_data;
-	bool ret;
-
-	read_priv.snor = snor_get_dev(vd->bus, vd->cs);
-	if (!read_priv.snor)
-		return -ENODEV;
-
-	read_priv.p.page_size = read_priv.snor->page_size;
-	read_priv.p.block_size = read_priv.snor->sector_size;
-	read_priv.p.read = snor_image_read;
-
-	/* Verify kernel first */
-	ret = read_verify_kernel(&read_priv.p, vd->load_ptr,
-				 vd->firmware_offset, vd->max_size,
-				 &vd->kernel_size, &vd->kernel_hashes);
-	if (!ret) {
-		printf("Error: kernel verification failed\n");
-		return -EBADMSG;
-	}
-
-	/* Verify rootfs, requiring valid kernel FIT image */
-	rootfs_data = vd->load_ptr + vd->kernel_size;
-	ret = read_verify_rootfs(&read_priv.p, vd->load_ptr, rootfs_data,
-				 vd->firmware_offset + vd->kernel_size,
-				 vd->max_size - vd->kernel_size,
-				 &vd->rootfs_size, true, &vd->padding_size,
-				 &vd->rootfs_hashes);
-	if (!ret) {
-		printf("Error: rootfs verification failed\n");
-		return -EBADMSG;
-	}
-
-	/* Read padding area */
-	snor_read_generic(read_priv.snor, vd->firmware_offset + vd->kernel_size,
-			  rootfs_data, vd->padding_size);
-
-	return 0;
-}
-#endif /* CONFIG_DM_SPI_FLASH */

uboot-mtk-20220606/board/mediatek/common/verify_helper.h

@@ -67,23 +67,4 @@ struct mmc_verify_data {
 int mmc_boot_verify(struct mmc_verify_data *vd);
 #endif /* CONFIG_MMC */
 
-#ifdef CONFIG_DM_SPI_FLASH
-struct snor_verify_data {
-	int bus;
-	int cs;
-	u32 firmware_offset;
-	u32 max_size;
-	void *load_ptr;
-
-	size_t kernel_size;
-	size_t padding_size;
-	size_t rootfs_size;
-
-	struct fit_hashes kernel_hashes;
-	struct fit_hashes rootfs_hashes;
-};
-
-int snor_boot_verify(struct snor_verify_data *vd);
-#endif /* CONFIG_DM_SPI_FLASH */
-
 #endif /* _VERIFY_HELPER_H_ */

uboot-mtk-20220606/board/mediatek/mt7622/bootmenu_snand.c

@@ -7,117 +7,78 @@
 
 #include <linux/mtd/mtd.h>
 #include <linux/sizes.h>
+#include <mtd.h>
 #include "upgrade_helper.h"
 #include "boot_helper.h"
 #include "autoboot_helper.h"
 #include "colored_print.h"
 
-#include <nmbm/nmbm.h>
-#include <nmbm/nmbm-mtd.h>
-
-enum upgrade_part_type {
-	UPGRADE_PART_BL2,
-	UPGRADE_PART_FIP,
-	UPGRADE_PART_FW,
-};
-
-struct upgrade_part {
-	u64 offset;
-	size_t size;
-};
-
-static const struct upgrade_part upgrade_parts[] = {
-	[UPGRADE_PART_BL2] = {
-		.offset = 0,
-		.size = 0x80000,
-	},
-	[UPGRADE_PART_FIP] = {
-		.offset = 0x80000,
-		.size = 0x200000,
-	},
-	[UPGRADE_PART_FW] = {
-		.offset = 0x400000,
-		.size = 0x3000000,
-	},
-};
-
-static struct mtd_info *board_get_mtd_device(void)
-{
-	struct mtd_info *mtd;
-
-#ifdef CONFIG_ENABLE_NAND_NMBM
-	mtd = nmbm_mtd_get_upper_by_index(0);
-
-	if (mtd)
-		mtd = get_mtd_device(mtd, -1);
-
-	if (!mtd)
-		cprintln(ERROR, "*** NMBM MTD device %u not found! ***", 0);
-#else
-	mtd = get_mtd_device(NULL, 0);
-
-	if (!mtd)
-		cprintln(ERROR, "*** NAND MTD device %u not found! ***", 0);
-#endif
-
-	return mtd;
-}
-
-static int write_part(enum upgrade_part_type pt, const void *data, size_t size,
+static int write_part(const char *partname, const void *data, size_t size,
 		      bool verify)
 {
-	const struct upgrade_part *part = &upgrade_parts[pt];
 	struct mtd_info *mtd;
 	int ret;
 
-	mtd = board_get_mtd_device();
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_generic(mtd, part->offset, size);
-	if (ret)
-		return ret;
+	mtd = get_mtd_device_nm(partname);
+	if (IS_ERR(mtd)) {
+		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
+			 partname);
+		put_mtd_device(mtd);
+		return -PTR_ERR(mtd);
+	}
 
-	ret = mtd_write_generic(mtd, part->offset, part->size, data, size,
-				verify);
+	ret = mtd_update_generic(mtd, data, size, verify);
 
 	put_mtd_device(mtd);
 
 	return ret;
 }
 
-static int write_bl2(void *priv, const struct data_part_entry *dpe,
+int write_bl2(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_BL2, data, size, true);
+	return write_part("bl2", data, size, true);
 }
 
-static int write_fip(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FIP, data, size, true);
+	return write_part("fip", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FW, data, size, true);
+	return write_part("firmware", data, size, true);
 }
 
 static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 			     const void *data, size_t size)
 {
-	struct mtd_info *mtd = board_get_mtd_device();
+	struct mtd_info *mtd;
 	int ret;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
+#ifdef CONFIG_ENABLE_NAND_NMBM
+	mtd = get_mtd_device_nm("nmbm0");
+#else
+	mtd = get_mtd_device(NULL, 0);
+#endif
+
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				false);
+	if (ret) {
+		put_mtd_device(mtd);
 		return ret;
+	}
 
-	ret = mtd_write_generic(mtd, 0, 0, data, size, true);
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
 
 	put_mtd_device(mtd);
 
@@ -127,21 +88,23 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-#if !defined(CONFIG_MTK_SECURE_BOOT) && !defined(CONFIG_ENV_IS_NOWHERE)
-	struct mtd_info *mtd = board_get_mtd_device();
-	int ret;
+	int ret = 0;
+#if !defined(CONFIG_MTK_SECURE_BOOT) && defined(CONFIG_ENV_IS_IN_MTD)
+	struct mtd_info *mtd;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_env(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
+	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
 
 	put_mtd_device(mtd);
+#endif
 
 	return ret;
-#else
-	return 0;
-#endif
 }
 
 static const struct data_part_entry snand_parts[] = {
@@ -157,7 +120,7 @@ static const struct data_part_entry snand_parts[] = {
 		.env_name = "bootfile.fip",
 		.write = write_fip,
 		.post_action = UPGRADE_ACTION_CUSTOM,
-		.do_post_action = erase_env,
+		//.do_post_action = erase_env,
 	},
 	{
 		.name = "Firmware",
@@ -182,18 +145,18 @@ void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 
 int board_boot_default(void)
 {
-	const struct upgrade_part *part = &upgrade_parts[UPGRADE_PART_FW];
-	struct mtd_info *mtd = board_get_mtd_device();
-	int ret;
+	struct mtd_info *mtd;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = boot_from_mtd(mtd, part->offset);
+	mtd = get_mtd_device_nm("firmware");
+	if (!IS_ERR_OR_NULL(mtd)) {
+		put_mtd_device(mtd);
 
-	put_mtd_device(mtd);
+		return boot_from_mtd(mtd, 0);
+	}
 
-	return ret;
+	return -ENODEV;
 }
 
 static const struct bootmenu_entry snand_bootmenu_entries[] = {

uboot-mtk-20220606/board/mediatek/mt7622/bootmenu_snor.c

@@ -5,118 +5,99 @@
  * Author: Weijie Gao <weijie.gao@mediatek.com>
  */
 
+#include <linux/mtd/mtd.h>
 #include <command.h>
 #include <linux/sizes.h>
+#include <mtd.h>
 #include "upgrade_helper.h"
 #include "boot_helper.h"
 #include "autoboot_helper.h"
 #include "colored_print.h"
 
-enum upgrade_part_type {
-	UPGRADE_PART_BL2,
-	UPGRADE_PART_FIP,
-	UPGRADE_PART_FW,
-};
-
-struct upgrade_part {
-	u32 offset;
-	size_t size;
-};
-
-static const struct upgrade_part upgrade_parts[] = {
-	[UPGRADE_PART_BL2] = {
-		.offset = 0,
-		.size = 0x20000,
-	},
-	[UPGRADE_PART_FIP] = {
-		.offset = 0x20000,
-		.size = 0x200000,
-	},
-	[UPGRADE_PART_FW] = {
-		.offset = 0x270000,
-		.size = 0xd900000,
-	},
-};
-
-/*
- * layout: 128k(bl2),2048k(fip),64k(config),256k(factory),-(firmware)
- */
-
-static inline struct spi_flash *board_get_snor_dev(void)
-{
-	return snor_get_dev(CONFIG_SF_DEFAULT_BUS, CONFIG_SF_DEFAULT_CS);
-}
-
-static int write_part(enum upgrade_part_type pt, const void *data, size_t size,
+static int write_part(const char *partname, const void *data, size_t size,
 		      bool verify)
 {
-	const struct upgrade_part *part = &upgrade_parts[pt];
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 	int ret;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = snor_erase_generic(snor, part->offset, size);
-	if (ret)
-		return ret;
+	mtd = get_mtd_device_nm(partname);
+	if (IS_ERR(mtd)) {
+		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
+			 partname);
+		return -PTR_ERR(mtd);
+	}
 
-	ret = snor_write_generic(snor, part->offset, part->size, data, size,
-				 verify);
+	ret = mtd_update_generic(mtd, data, size, verify);
+
+	put_mtd_device(mtd);
 
 	return ret;
 }
 
-static int write_bl2(void *priv, const struct data_part_entry *dpe,
+int write_bl2(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_BL2, data, size, true);
+	return write_part("bl2", data, size, true);
 }
 
-static int write_fip(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FIP, data, size, true);
+	return write_part("fip", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FW, data, size, false);
+	return write_part("firmware", data, size, false);
 }
 
 static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 			     const void *data, size_t size)
 {
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 	int ret;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	mtd = get_mtd_device(NULL, 0);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
 
-	ret = snor_erase_generic(snor, 0, size);
-	if (ret)
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				 false);
+	if (ret) {
+		put_mtd_device(mtd);
 		return ret;
+	}
 
-	return snor_write_generic(snor, 0, 0, data, size, false);
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
+
+	put_mtd_device(mtd);
+
+	return ret;
 }
 
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-#if !defined(CONFIG_ENV_IS_NOWHERE)
-	struct spi_flash *snor;
+	int ret = 0;
+#if defined(CONFIG_ENV_IS_IN_MTD)
+	struct mtd_info *mtd;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	return snor_erase_env(snor, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
-#else
-	return 0;
+	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
+
+	put_mtd_device(mtd);
 #endif
+
+	return ret;
 }
 
 static const struct data_part_entry snor_parts[] = {
@@ -132,7 +113,7 @@ static const struct data_part_entry snor_parts[] = {
 		.env_name = "bootfile.fip",
 		.write = write_fip,
 		.post_action = UPGRADE_ACTION_CUSTOM,
-		.do_post_action = erase_env,
+		//.do_post_action = erase_env,
 	},
 	{
 		.name = "Firmware",
@@ -157,12 +138,18 @@ void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 
 int board_boot_default(void)
 {
-	const struct upgrade_part *part = &upgrade_parts[UPGRADE_PART_FW];
-	char cmd[64];
+	struct mtd_info *mtd;
 
-	sprintf(cmd, "bootm 0x%lx", 0x30000000 + (ulong)part->offset);
+	gen_mtd_probe_devices();
 
-	return run_command(cmd, 0);
+	mtd = get_mtd_device_nm("firmware");
+	if (!IS_ERR_OR_NULL(mtd)) {
+		put_mtd_device(mtd);
+
+		return boot_from_mtd(mtd, 0);
+	}
+
+	return -ENODEV;
 }
 
 static const struct bootmenu_entry snor_bootmenu_entries[] = {

uboot-mtk-20220606/board/mediatek/mt7622/bootmenu_ubi.c

@@ -19,37 +19,36 @@ static int write_part(const char *partname, const void *data, size_t size,
 	struct mtd_info *mtd;
 	int ret;
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 	mtd = get_mtd_device_nm(partname);
 	if (IS_ERR(mtd)) {
 		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
 			 partname);
+		put_mtd_device(mtd);
 		return -PTR_ERR(mtd);
 	}
 
+	ret = mtd_update_generic(mtd, data, size, verify);
+
 	put_mtd_device(mtd);
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
-		return ret;
-
-	return mtd_write_generic(mtd, 0, 0, data, size, verify);
+	return ret;
 }
 
-static int write_bl2(void *priv, const struct data_part_entry *dpe,
+int write_bl2(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
 	return write_part("bl2", data, size, true);
 }
 
-static int write_fip(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
 	return write_part("fip", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
 	int ret;
@@ -68,7 +67,7 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 	struct mtd_info *mtd;
 	int ret;
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 #ifdef CONFIG_ENABLE_NAND_NMBM
 	mtd = get_mtd_device_nm("nmbm0");
@@ -79,37 +78,43 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 	if (IS_ERR(mtd))
 		return -PTR_ERR(mtd);
 
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				false);
+	if (ret) {
+		put_mtd_device(mtd);
+		return ret;
+	}
+
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
+
 	put_mtd_device(mtd);
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
-		return ret;
-
-	return mtd_write_generic(mtd, 0, 0, data, size, true);
+	return ret;
 }
 
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-#if !defined(CONFIG_ENV_IS_NOWHERE)
+	int ret = 0;
+#if !defined(CONFIG_MTK_SECURE_BOOT) && defined(CONFIG_ENV_IS_IN_MTD)
 	struct mtd_info *mtd;
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
-
 	if (IS_ERR(mtd))
 		return -PTR_ERR(mtd);
 
-	put_mtd_device(mtd);
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
 
-	return mtd_erase_env(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
-#else
-	return 0;
+	put_mtd_device(mtd);
 #endif
+
+	return ret;
 }
 
-static const struct data_part_entry snand_parts[] = {
+static const struct data_part_entry ubi_parts[] = {
 	{
 		.name = "ATF BL2",
 		.abbr = "bl2",
@@ -141,8 +146,8 @@ static const struct data_part_entry snand_parts[] = {
 
 void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 {
-	*dpes = snand_parts;
-	*count = ARRAY_SIZE(snand_parts);
+	*dpes = ubi_parts;
+	*count = ARRAY_SIZE(ubi_parts);
 }
 
 int board_boot_default(void)
@@ -150,7 +155,7 @@ int board_boot_default(void)
 	return ubi_boot_image();
 }
 
-static const struct bootmenu_entry snand_bootmenu_entries[] = {
+static const struct bootmenu_entry ubi_bootmenu_entries[] = {
 	{
 		.desc = "Startup system (Default)",
 		.cmd = "mtkboardboot"
@@ -179,6 +184,6 @@ static const struct bootmenu_entry snand_bootmenu_entries[] = {
 
 void board_bootmenu_entries(const struct bootmenu_entry **menu, u32 *count)
 {
-	*menu = snand_bootmenu_entries;
-	*count = ARRAY_SIZE(snand_bootmenu_entries);
+	*menu = ubi_bootmenu_entries;
+	*count = ARRAY_SIZE(ubi_bootmenu_entries);
 }

uboot-mtk-20220606/board/mediatek/mt7629/bootmenu_snand.c

@@ -7,106 +7,72 @@
 
 #include <linux/mtd/mtd.h>
 #include <linux/sizes.h>
+#include <mtd.h>
 #include "upgrade_helper.h"
 #include "boot_helper.h"
 #include "autoboot_helper.h"
 #include "colored_print.h"
 
-#include <nmbm/nmbm.h>
-#include <nmbm/nmbm-mtd.h>
-
-enum upgrade_part_type {
-	UPGRADE_PART_BL,
-	UPGRADE_PART_FW,
-};
-
-struct upgrade_part {
-	u64 offset;
-	size_t size;
-};
-
-static const struct upgrade_part upgrade_parts[] = {
-	[UPGRADE_PART_BL] = {
-		.offset = 0,
-		.size = 0x100000,
-	},
-	[UPGRADE_PART_FW] = {
-		.offset = 0x1c0000,
-		.size = 0x2000000,
-	},
-};
-
-static struct mtd_info *board_get_mtd_device(void)
-{
-	struct mtd_info *mtd;
-
-#ifdef CONFIG_ENABLE_NAND_NMBM
-	mtd = nmbm_mtd_get_upper_by_index(0);
-
-	if (mtd)
-		mtd = get_mtd_device(mtd, -1);
-
-	if (!mtd)
-		cprintln(ERROR, "*** NMBM MTD device %u not found! ***", 0);
-#else
-	mtd = get_mtd_device(NULL, 0);
-
-	if (!mtd)
-		cprintln(ERROR, "*** NAND MTD device %u not found! ***", 0);
-#endif
-
-	return mtd;
-}
-
-static int write_part(enum upgrade_part_type pt, const void *data, size_t size,
+static int write_part(const char *partname, const void *data, size_t size,
 		      bool verify)
 {
-	const struct upgrade_part *part = &upgrade_parts[pt];
 	struct mtd_info *mtd;
 	int ret;
 
-	mtd = board_get_mtd_device();
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_generic(mtd, part->offset, size);
-	if (ret)
-		return ret;
+	mtd = get_mtd_device_nm(partname);
+	if (IS_ERR(mtd)) {
+		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
+			 partname);
+		put_mtd_device(mtd);
+		return -PTR_ERR(mtd);
+	}
 
-	ret = mtd_write_generic(mtd, part->offset, part->size, data, size,
-				verify);
+	ret = mtd_update_generic(mtd, data, size, verify);
 
 	put_mtd_device(mtd);
 
 	return ret;
 }
 
-static int write_bl(void *priv, const struct data_part_entry *dpe,
-		    const void *data, size_t size)
+int write_bl(void *priv, const struct data_part_entry *dpe,
+		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_BL, data, size, true);
+	return write_part("bootloader", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FW, data, size, true);
+	return write_part("firmware", data, size, true);
 }
 
 static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 			     const void *data, size_t size)
 {
-	struct mtd_info *mtd = board_get_mtd_device();
+	struct mtd_info *mtd;
 	int ret;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
+#ifdef CONFIG_ENABLE_NAND_NMBM
+	mtd = get_mtd_device_nm("nmbm0");
+#else
+	mtd = get_mtd_device(NULL, 0);
+#endif
+
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				false);
+	if (ret) {
+		put_mtd_device(mtd);
 		return ret;
+	}
 
-	ret = mtd_write_generic(mtd, 0, 0, data, size, true);
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
 
 	put_mtd_device(mtd);
 
@@ -116,21 +82,23 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-#if !defined(CONFIG_ENV_IS_NOWHERE)
-	struct mtd_info *mtd = board_get_mtd_device();
-	int ret;
+	int ret = 0;
+#if !defined(CONFIG_MTK_SECURE_BOOT) && defined(CONFIG_ENV_IS_IN_MTD)
+	struct mtd_info *mtd;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_env(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
+	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
 
 	put_mtd_device(mtd);
+#endif
 
 	return ret;
-#else
-	return 0;
-#endif
 }
 
 static const struct data_part_entry snand_parts[] = {
@@ -140,7 +108,7 @@ static const struct data_part_entry snand_parts[] = {
 		.env_name = "bootfile.bl",
 		.write = write_bl,
 		.post_action = UPGRADE_ACTION_CUSTOM,
-		.do_post_action = erase_env,
+		//.do_post_action = erase_env,
 	},
 	{
 		.name = "Firmware",
@@ -165,18 +133,18 @@ void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 
 int board_boot_default(void)
 {
-	const struct upgrade_part *part = &upgrade_parts[UPGRADE_PART_FW];
-	struct mtd_info *mtd = board_get_mtd_device();
-	int ret;
+	struct mtd_info *mtd;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = boot_from_mtd(mtd, part->offset);
+	mtd = get_mtd_device_nm("firmware");
+	if (!IS_ERR_OR_NULL(mtd)) {
+		put_mtd_device(mtd);
 
-	put_mtd_device(mtd);
+		return boot_from_mtd(mtd, 0);
+	}
 
-	return ret;
+	return -ENODEV;
 }
 
 static const struct bootmenu_entry snand_bootmenu_entries[] = {

uboot-mtk-20220606/board/mediatek/mt7629/bootmenu_snand_sb.c

@@ -7,121 +7,79 @@
 
 #include <linux/mtd/mtd.h>
 #include <linux/sizes.h>
+#include <mtd.h>
 #include "upgrade_helper.h"
 #include "boot_helper.h"
 #include "autoboot_helper.h"
 #include "colored_print.h"
 
-#include <nmbm/nmbm.h>
-#include <nmbm/nmbm-mtd.h>
 
-enum upgrade_part_type {
-	UPGRADE_PART_BL2,
-	UPGRADE_PART_FIP,
-	UPGRADE_PART_FW,
-};
-
-struct upgrade_part {
-	u64 offset;
-	size_t size;
-};
-
-static const struct upgrade_part upgrade_parts[] = {
-	[UPGRADE_PART_BL2] = {
-		.offset = 0,
-		.size = 0x80000,
-	},
-	[UPGRADE_PART_FIP] = {
-		.offset = 0x80000,
-		.size = 0x200000,
-	},
-	[UPGRADE_PART_FW] = {
-		.offset = 0x400000,
-		.size = 0x1400000,
-	},
-};
-
-/*
- * layout: 512k(bl2),2048k(fip),512k(config),1024k(factory),-(firmware)
- */
-
-static struct mtd_info *board_get_mtd_device(void)
-{
-	struct mtd_info *mtd;
-
-#ifdef CONFIG_ENABLE_NAND_NMBM
-	mtd = nmbm_mtd_get_upper_by_index(0);
-
-	if (mtd)
-		mtd = get_mtd_device(mtd, -1);
-
-	if (!mtd)
-		cprintln(ERROR, "*** NMBM MTD device %u not found! ***", 0);
-#else
-	mtd = get_mtd_device(NULL, 0);
-
-	if (!mtd)
-		cprintln(ERROR, "*** NAND MTD device %u not found! ***", 0);
-#endif
-
-	return mtd;
-}
-
-static int write_part(enum upgrade_part_type pt, const void *data, size_t size,
+static int write_part(const char *partname, const void *data, size_t size,
 		      bool verify)
 {
-	const struct upgrade_part *part = &upgrade_parts[pt];
 	struct mtd_info *mtd;
 	int ret;
 
-	mtd = board_get_mtd_device();
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_generic(mtd, part->offset, size);
-	if (ret)
-		return ret;
+	mtd = get_mtd_device_nm(partname);
+	if (IS_ERR(mtd)) {
+		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
+			 partname);
+		put_mtd_device(mtd);
+		return -PTR_ERR(mtd);
+	}
 
-	ret = mtd_write_generic(mtd, part->offset, part->size, data, size,
-				verify);
+	ret = mtd_update_generic(mtd, data, size, verify);
 
 	put_mtd_device(mtd);
 
 	return ret;
 }
 
-static int write_bl2(void *priv, const struct data_part_entry *dpe,
+int write_bl2(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_BL2, data, size, true);
+	return write_part("bl2", data, size, true);
 }
 
-static int write_fip(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FIP, data, size, true);
+	return write_part("fip", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FW, data, size, true);
+	return write_part("firmware", data, size, true);
 }
 
 static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 			     const void *data, size_t size)
 {
-	struct mtd_info *mtd = board_get_mtd_device();
+	struct mtd_info *mtd;
 	int ret;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
+#ifdef CONFIG_ENABLE_NAND_NMBM
+	mtd = get_mtd_device_nm("nmbm0");
+#else
+	mtd = get_mtd_device(NULL, 0);
+#endif
+
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				false);
+	if (ret) {
+		put_mtd_device(mtd);
 		return ret;
+	}
 
-	ret = mtd_write_generic(mtd, 0, 0, data, size, true);
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
 
 	put_mtd_device(mtd);
 
@@ -165,18 +123,18 @@ void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 
 int board_boot_default(void)
 {
-	const struct upgrade_part *part = &upgrade_parts[UPGRADE_PART_FW];
-	struct mtd_info *mtd = board_get_mtd_device();
-	int ret;
+	struct mtd_info *mtd;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = boot_from_mtd(mtd, part->offset);
+	mtd = get_mtd_device_nm("firmware");
+	if (!IS_ERR_OR_NULL(mtd)) {
+		put_mtd_device(mtd);
 
-	put_mtd_device(mtd);
+		return boot_from_mtd(mtd, 0);
+	}
 
-	return ret;
+	return -ENODEV;
 }
 
 static const struct bootmenu_entry snand_bootmenu_entries[] = {

uboot-mtk-20220606/board/mediatek/mt7629/bootmenu_snor.c

@@ -5,113 +5,115 @@
  * Author: Weijie Gao <weijie.gao@mediatek.com>
  */
 
+#include <linux/mtd/mtd.h>
 #include <command.h>
 #include <linux/sizes.h>
+#include <mtd.h>
 #include "upgrade_helper.h"
 #include "boot_helper.h"
 #include "autoboot_helper.h"
 #include "colored_print.h"
 
-enum upgrade_part_type {
-	UPGRADE_PART_BL,
-	UPGRADE_PART_FW,
-};
-
-struct upgrade_part {
-	u32 offset;
-	size_t size;
-};
-
-static const struct upgrade_part upgrade_parts[] = {
-	[UPGRADE_PART_BL] = {
-		.offset = 0,
-		.size = 0x60000,
-	},
-	[UPGRADE_PART_FW] = {
-		.offset = 0xb0000,
-		.size = 0xf20000,
-	},
-};
-
-static inline struct spi_flash *board_get_snor_dev(void)
-{
-	return snor_get_dev(CONFIG_SF_DEFAULT_BUS, CONFIG_SF_DEFAULT_CS);
-}
-
-static int write_part(enum upgrade_part_type pt, const void *data, size_t size,
+static int write_part(const char *partname, const void *data, size_t size,
 		      bool verify)
 {
-	const struct upgrade_part *part = &upgrade_parts[pt];
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 	int ret;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = snor_erase_generic(snor, part->offset, size);
-	if (ret)
-		return ret;
+	mtd = get_mtd_device_nm(partname);
+	if (IS_ERR(mtd)) {
+		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
+			 partname);
+		return -PTR_ERR(mtd);
+	}
 
-	ret = snor_write_generic(snor, part->offset, part->size, data, size,
-				 verify);
+	ret = mtd_update_generic(mtd, data, size, verify);
+
+	put_mtd_device(mtd);
 
 	return ret;
 }
 
-static int write_bl(void *priv, const struct data_part_entry *dpe,
-		    const void *data, size_t size)
+int write_bl2(void *priv, const struct data_part_entry *dpe,
+		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_BL, data, size, true);
+	return write_part("bl2", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
+		     const void *data, size_t size)
+{
+	return write_part("fip", data, size, true);
+}
+
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FW, data, size, false);
+	return write_part("firmware", data, size, false);
 }
 
 static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 			     const void *data, size_t size)
 {
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 	int ret;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	mtd = get_mtd_device(NULL, 0);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
 
-	ret = snor_erase_generic(snor, 0, size);
-	if (ret)
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				 false);
+	if (ret) {
+		put_mtd_device(mtd);
 		return ret;
+	}
 
-	return snor_write_generic(snor, 0, 0, data, size, false);
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
+
+	put_mtd_device(mtd);
+
+	return ret;
 }
 
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-#if !defined(CONFIG_ENV_IS_NOWHERE)
-	struct spi_flash *snor;
+	int ret = 0;
+#if defined(CONFIG_ENV_IS_IN_MTD)
+	struct mtd_info *mtd;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	return snor_erase_env(snor, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
-#else
-	return 0;
+	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
+
+	put_mtd_device(mtd);
 #endif
+
+	return ret;
 }
 
 static const struct data_part_entry snor_parts[] = {
 	{
-		.name = "Bootloader",
-		.abbr = "bl",
-		.env_name = "bootfile.bl",
-		.write = write_bl,
+		.name = "ATF BL2",
+		.abbr = "bl2",
+		.env_name = "bootfile.bl2",
+		.write = write_bl2,
+	},
+	{
+		.name = "ATF FIP",
+		.abbr = "fip",
+		.env_name = "bootfile.fip",
+		.write = write_fip,
 		.post_action = UPGRADE_ACTION_CUSTOM,
-		.do_post_action = erase_env,
+		//.do_post_action = erase_env,
 	},
 	{
 		.name = "Firmware",
@@ -136,12 +138,18 @@ void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 
 int board_boot_default(void)
 {
-	const struct upgrade_part *part = &upgrade_parts[UPGRADE_PART_FW];
-	char cmd[64];
+	struct mtd_info *mtd;
 
-	sprintf(cmd, "bootm 0x%lx", 0x30000000 + (ulong)part->offset);
+	gen_mtd_probe_devices();
 
-	return run_command(cmd, 0);
+	mtd = get_mtd_device_nm("firmware");
+	if (!IS_ERR_OR_NULL(mtd)) {
+		put_mtd_device(mtd);
+
+		return boot_from_mtd(mtd, 0);
+	}
+
+	return -ENODEV;
 }
 
 static const struct bootmenu_entry snor_bootmenu_entries[] = {
@@ -154,8 +162,12 @@ static const struct bootmenu_entry snor_bootmenu_entries[] = {
 		.cmd = "mtkupgrade fw"
 	},
 	{
-		.desc = "Upgrade bootloader",
-		.cmd = "mtkupgrade bl"
+		.desc = "Upgrade ATF BL2",
+		.cmd = "mtkupgrade bl2"
+	},
+	{
+		.desc = "Upgrade ATF FIP",
+		.cmd = "mtkupgrade fip"
 	},
 	{
 		.desc = "Upgrade single image",

uboot-mtk-20220606/board/mediatek/mt7629/bootmenu_snor_sb.c

@@ -5,114 +5,97 @@
  * Author: Weijie Gao <weijie.gao@mediatek.com>
  */
 
+#include <linux/mtd/mtd.h>
 #include <command.h>
 #include <linux/sizes.h>
+#include <mtd.h>
 #include "upgrade_helper.h"
 #include "boot_helper.h"
 #include "autoboot_helper.h"
 #include "colored_print.h"
 
-enum upgrade_part_type {
-	UPGRADE_PART_BL2,
-	UPGRADE_PART_FIP,
-	UPGRADE_PART_FW,
-};
-
-struct upgrade_part {
-	u32 offset;
-	size_t size;
-};
-
-static const struct upgrade_part upgrade_parts[] = {
-	[UPGRADE_PART_BL2] = {
-		.offset = 0,
-		.size = 0x20000,
-	},
-	[UPGRADE_PART_FIP] = {
-		.offset = 0x20000,
-		.size = 0x200000,
-	},
-	[UPGRADE_PART_FW] = {
-		.offset = 0x270000,
-		.size = 0xd900000,
-	},
-};
-
-/*
- * layout: 128k(bl2),2048k(fip),64k(config),256k(factory),-(firmware)
- */
-
-static inline struct spi_flash *board_get_snor_dev(void)
-{
-	return snor_get_dev(CONFIG_SF_DEFAULT_BUS, CONFIG_SF_DEFAULT_CS);
-}
-
-static int write_part(enum upgrade_part_type pt, const void *data, size_t size,
+static int write_part(const char *partname, const void *data, size_t size,
 		      bool verify)
 {
-	const struct upgrade_part *part = &upgrade_parts[pt];
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 	int ret;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = snor_erase_generic(snor, part->offset, size);
-	if (ret)
-		return ret;
+	mtd = get_mtd_device_nm(partname);
+	if (IS_ERR(mtd)) {
+		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
+			 partname);
+		return -PTR_ERR(mtd);
+	}
 
-	ret = snor_write_generic(snor, part->offset, part->size, data, size,
-				 verify);
+	ret = mtd_update_generic(mtd, data, size, verify);
+
+	put_mtd_device(mtd);
 
 	return ret;
 }
 
-static int write_bl2(void *priv, const struct data_part_entry *dpe,
+int write_bl2(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_BL2, data, size, true);
+	return write_part("bl2", data, size, true);
 }
 
-static int write_fip(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FIP, data, size, true);
+	return write_part("fip", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FW, data, size, false);
+	return write_part("firmware", data, size, false);
 }
 
 static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 			     const void *data, size_t size)
 {
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 	int ret;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	mtd = get_mtd_device(NULL, 0);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
 
-	ret = snor_erase_generic(snor, 0, size);
-	if (ret)
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				 false);
+	if (ret) {
+		put_mtd_device(mtd);
 		return ret;
+	}
 
-	return snor_write_generic(snor, 0, 0, data, size, false);
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
+
+	put_mtd_device(mtd);
+
+	return ret;
 }
 
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	struct spi_flash *snor;
+	int ret;
+	struct mtd_info *mtd;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	return snor_erase_env(snor, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
+	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
+
+	put_mtd_device(mtd);
+
+	return ret;
 }
 
 static const struct data_part_entry snor_parts[] = {
@@ -128,7 +111,7 @@ static const struct data_part_entry snor_parts[] = {
 		.env_name = "bootfile.fip",
 		.write = write_fip,
 		.post_action = UPGRADE_ACTION_CUSTOM,
-		.do_post_action = erase_env,
+		//.do_post_action = erase_env,
 	},
 	{
 		.name = "Firmware",
@@ -153,12 +136,18 @@ void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 
 int board_boot_default(void)
 {
-	const struct upgrade_part *part = &upgrade_parts[UPGRADE_PART_FW];
-	char cmd[64];
+	struct mtd_info *mtd;
 
-	sprintf(cmd, "bootm 0x%lx", 0x30000000 + (ulong)part->offset);
+	gen_mtd_probe_devices();
 
-	return run_command(cmd, 0);
+	mtd = get_mtd_device_nm("firmware");
+	if (!IS_ERR_OR_NULL(mtd)) {
+		put_mtd_device(mtd);
+
+		return boot_from_mtd(mtd, 0);
+	}
+
+	return -ENODEV;
 }
 
 static const struct bootmenu_entry snor_bootmenu_entries[] = {

uboot-mtk-20220606/board/mediatek/mt7981/bootmenu_snand.c

@@ -7,117 +7,78 @@
 
 #include <linux/mtd/mtd.h>
 #include <linux/sizes.h>
+#include <mtd.h>
 #include "upgrade_helper.h"
 #include "boot_helper.h"
 #include "autoboot_helper.h"
 #include "colored_print.h"
 
-#include <nmbm/nmbm.h>
-#include <nmbm/nmbm-mtd.h>
-
-enum upgrade_part_type {
-	UPGRADE_PART_BL2,
-	UPGRADE_PART_FIP,
-	UPGRADE_PART_FW,
-};
-
-struct upgrade_part {
-	u64 offset;
-	size_t size;
-};
-
-static const struct upgrade_part upgrade_parts[] = {
-	[UPGRADE_PART_BL2] = {
-		.offset = 0,
-		.size = 0x100000,
-	},
-	[UPGRADE_PART_FIP] = {
-		.offset = 0x380000,
-		.size = 0x200000,
-	},
-	[UPGRADE_PART_FW] = {
-		.offset = 0x580000,
-		.size = 0x2800000,
-	},
-};
-
-static struct mtd_info *board_get_mtd_device(void)
-{
-	struct mtd_info *mtd;
-
-#ifdef CONFIG_ENABLE_NAND_NMBM
-	mtd = nmbm_mtd_get_upper_by_index(0);
-
-	if (mtd)
-		mtd = get_mtd_device(mtd, -1);
-
-	if (!mtd)
-		cprintln(ERROR, "*** NMBM MTD device %u not found! ***", 0);
-#else
-	mtd = get_mtd_device(NULL, 0);
-
-	if (!mtd)
-		cprintln(ERROR, "*** NAND MTD device %u not found! ***", 0);
-#endif
-
-	return mtd;
-}
-
-static int write_part(enum upgrade_part_type pt, const void *data, size_t size,
+static int write_part(const char *partname, const void *data, size_t size,
 		      bool verify)
 {
-	const struct upgrade_part *part = &upgrade_parts[pt];
 	struct mtd_info *mtd;
 	int ret;
 
-	mtd = board_get_mtd_device();
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_generic(mtd, part->offset, size);
-	if (ret)
-		return ret;
+	mtd = get_mtd_device_nm(partname);
+	if (IS_ERR(mtd)) {
+		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
+			 partname);
+		put_mtd_device(mtd);
+		return -PTR_ERR(mtd);
+	}
 
-	ret = mtd_write_generic(mtd, part->offset, part->size, data, size,
-				verify);
+	ret = mtd_update_generic(mtd, data, size, verify);
 
 	put_mtd_device(mtd);
 
 	return ret;
 }
 
-static int write_bl2(void *priv, const struct data_part_entry *dpe,
+int write_bl2(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_BL2, data, size, true);
+	return write_part("bl2", data, size, true);
 }
 
-static int write_fip(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FIP, data, size, true);
+	return write_part("fip", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FW, data, size, true);
+	return write_part("firmware", data, size, true);
 }
 
 static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 			     const void *data, size_t size)
 {
-	struct mtd_info *mtd = board_get_mtd_device();
+	struct mtd_info *mtd;
 	int ret;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
+#ifdef CONFIG_ENABLE_NAND_NMBM
+	mtd = get_mtd_device_nm("nmbm0");
+#else
+	mtd = get_mtd_device(NULL, 0);
+#endif
+
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				false);
+	if (ret) {
+		put_mtd_device(mtd);
 		return ret;
+	}
 
-	ret = mtd_write_generic(mtd, 0, 0, data, size, true);
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
 
 	put_mtd_device(mtd);
 
@@ -127,21 +88,23 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-#if !defined(CONFIG_MTK_SECURE_BOOT) && !defined(CONFIG_ENV_IS_NOWHERE)
-	struct mtd_info *mtd = board_get_mtd_device();
-	int ret;
+	int ret = 0;
+#if !defined(CONFIG_MTK_SECURE_BOOT) && defined(CONFIG_ENV_IS_IN_MTD)
+	struct mtd_info *mtd;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_env(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
+	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
 
 	put_mtd_device(mtd);
+#endif
 
 	return ret;
-#else
-	return 0;
-#endif
 }
 
 static const struct data_part_entry snand_parts[] = {
@@ -157,7 +120,7 @@ static const struct data_part_entry snand_parts[] = {
 		.env_name = "bootfile.fip",
 		.write = write_fip,
 		.post_action = UPGRADE_ACTION_CUSTOM,
-		.do_post_action = erase_env,
+		//.do_post_action = erase_env,
 	},
 	{
 		.name = "Firmware",
@@ -182,18 +145,18 @@ void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 
 int board_boot_default(void)
 {
-	const struct upgrade_part *part = &upgrade_parts[UPGRADE_PART_FW];
-	struct mtd_info *mtd = board_get_mtd_device();
-	int ret;
+	struct mtd_info *mtd;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = boot_from_mtd(mtd, part->offset);
+	mtd = get_mtd_device_nm("firmware");
+	if (!IS_ERR_OR_NULL(mtd)) {
+		put_mtd_device(mtd);
 
-	put_mtd_device(mtd);
+		return boot_from_mtd(mtd, 0);
+	}
 
-	return ret;
+	return -ENODEV;
 }
 
 static const struct bootmenu_entry snand_bootmenu_entries[] = {

uboot-mtk-20220606/board/mediatek/mt7981/bootmenu_snor.c

@@ -5,118 +5,99 @@
  * Author: Weijie Gao <weijie.gao@mediatek.com>
  */
 
+#include <linux/mtd/mtd.h>
 #include <command.h>
 #include <linux/sizes.h>
+#include <mtd.h>
 #include "upgrade_helper.h"
 #include "boot_helper.h"
 #include "autoboot_helper.h"
 #include "colored_print.h"
 
-enum upgrade_part_type {
-	UPGRADE_PART_BL2,
-	UPGRADE_PART_FIP,
-	UPGRADE_PART_FW,
-};
-
-struct upgrade_part {
-	u32 offset;
-	size_t size;
-};
-
-static const struct upgrade_part upgrade_parts[] = {
-	[UPGRADE_PART_BL2] = {
-		.offset = 0,
-		.size = 0x40000,
-	},
-	[UPGRADE_PART_FIP] = {
-		.offset = 0x100000,
-		.size = 0x80000,
-	},
-	[UPGRADE_PART_FW] = {
-		.offset = 0x180000,
-		.size = 0xd900000,
-	},
-};
-
-/*
- * layout: 128k(bl2),2048k(fip),64k(config),256k(factory),-(firmware)
- */
-
-static inline struct spi_flash *board_get_snor_dev(void)
-{
-	return snor_get_dev(CONFIG_SF_DEFAULT_BUS, CONFIG_SF_DEFAULT_CS);
-}
-
-static int write_part(enum upgrade_part_type pt, const void *data, size_t size,
+static int write_part(const char *partname, const void *data, size_t size,
 		      bool verify)
 {
-	const struct upgrade_part *part = &upgrade_parts[pt];
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 	int ret;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = snor_erase_generic(snor, part->offset, size);
-	if (ret)
-		return ret;
+	mtd = get_mtd_device_nm(partname);
+	if (IS_ERR(mtd)) {
+		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
+			 partname);
+		return -PTR_ERR(mtd);
+	}
 
-	ret = snor_write_generic(snor, part->offset, part->size, data, size,
-				 verify);
+	ret = mtd_update_generic(mtd, data, size, verify);
+
+	put_mtd_device(mtd);
 
 	return ret;
 }
 
-static int write_bl2(void *priv, const struct data_part_entry *dpe,
+int write_bl2(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_BL2, data, size, true);
+	return write_part("bl2", data, size, true);
 }
 
-static int write_fip(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FIP, data, size, true);
+	return write_part("fip", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FW, data, size, false);
+	return write_part("firmware", data, size, false);
 }
 
 static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 			     const void *data, size_t size)
 {
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 	int ret;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	mtd = get_mtd_device(NULL, 0);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
 
-	ret = snor_erase_generic(snor, 0, size);
-	if (ret)
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				 false);
+	if (ret) {
+		put_mtd_device(mtd);
 		return ret;
+	}
 
-	return snor_write_generic(snor, 0, 0, data, size, false);
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
+
+	put_mtd_device(mtd);
+
+	return ret;
 }
 
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-#if !defined(CONFIG_ENV_IS_NOWHERE)
-	struct spi_flash *snor;
+	int ret = 0;
+#if defined(CONFIG_ENV_IS_IN_MTD)
+	struct mtd_info *mtd;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	return snor_erase_env(snor, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
-#else
-	return 0;
+	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
+
+	put_mtd_device(mtd);
 #endif
+
+	return ret;
 }
 
 static const struct data_part_entry snor_parts[] = {
@@ -132,7 +113,7 @@ static const struct data_part_entry snor_parts[] = {
 		.env_name = "bootfile.fip",
 		.write = write_fip,
 		.post_action = UPGRADE_ACTION_CUSTOM,
-		.do_post_action = erase_env,
+		//.do_post_action = erase_env,
 	},
 	{
 		.name = "Firmware",
@@ -157,14 +138,18 @@ void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 
 int board_boot_default(void)
 {
-	const struct upgrade_part *part = &upgrade_parts[UPGRADE_PART_FW];
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	return boot_from_snor(snor, part->offset);
+	mtd = get_mtd_device_nm("firmware");
+	if (!IS_ERR_OR_NULL(mtd)) {
+		put_mtd_device(mtd);
+
+		return boot_from_mtd(mtd, 0);
+	}
+
+	return -ENODEV;
 }
 
 static const struct bootmenu_entry snor_bootmenu_entries[] = {

uboot-mtk-20220606/board/mediatek/mt7981/bootmenu_ubi.c

@@ -19,22 +19,21 @@ static int write_part(const char *partname, const void *data, size_t size,
 	struct mtd_info *mtd;
 	int ret;
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 	mtd = get_mtd_device_nm(partname);
 	if (IS_ERR(mtd)) {
 		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
 			 partname);
+		put_mtd_device(mtd);
 		return -PTR_ERR(mtd);
 	}
 
+	ret = mtd_update_generic(mtd, data, size, verify);
+
 	put_mtd_device(mtd);
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
-		return ret;
-
-	return mtd_write_generic(mtd, 0, 0, data, size, verify);
+	return ret;
 }
 
 int write_bl2(void *priv, const struct data_part_entry *dpe,
@@ -68,7 +67,7 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 	struct mtd_info *mtd;
 	int ret;
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 #ifdef CONFIG_ENABLE_NAND_NMBM
 	mtd = get_mtd_device_nm("nmbm0");
@@ -79,34 +78,40 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 	if (IS_ERR(mtd))
 		return -PTR_ERR(mtd);
 
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				false);
+	if (ret) {
+		put_mtd_device(mtd);
+		return ret;
+	}
+
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
+
 	put_mtd_device(mtd);
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
-		return ret;
-
-	return mtd_write_generic(mtd, 0, 0, data, size, true);
+	return ret;
 }
 
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
+	int ret = 0;
 #if !defined(CONFIG_MTK_SECURE_BOOT) && defined(CONFIG_ENV_IS_IN_MTD)
 	struct mtd_info *mtd;
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
-
 	if (IS_ERR(mtd))
 		return -PTR_ERR(mtd);
 
-	put_mtd_device(mtd);
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
 
-	return mtd_erase_env(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
-#else
-	return 0;
+	put_mtd_device(mtd);
 #endif
+
+	return ret;
 }
 
 static const struct data_part_entry ubi_parts[] = {

uboot-mtk-20220606/board/mediatek/mt7986/bootmenu_snand.c

@@ -7,117 +7,78 @@
 
 #include <linux/mtd/mtd.h>
 #include <linux/sizes.h>
+#include <mtd.h>
 #include "upgrade_helper.h"
 #include "boot_helper.h"
 #include "autoboot_helper.h"
 #include "colored_print.h"
 
-#include <nmbm/nmbm.h>
-#include <nmbm/nmbm-mtd.h>
-
-enum upgrade_part_type {
-	UPGRADE_PART_BL2,
-	UPGRADE_PART_FIP,
-	UPGRADE_PART_FW,
-};
-
-struct upgrade_part {
-	u64 offset;
-	size_t size;
-};
-
-static const struct upgrade_part upgrade_parts[] = {
-	[UPGRADE_PART_BL2] = {
-		.offset = 0,
-		.size = 0x100000,
-	},
-	[UPGRADE_PART_FIP] = {
-		.offset = 0x380000,
-		.size = 0x200000,
-	},
-	[UPGRADE_PART_FW] = {
-		.offset = 0x580000,
-		.size = 0x2800000,
-	},
-};
-
-static struct mtd_info *board_get_mtd_device(void)
-{
-	struct mtd_info *mtd;
-
-#ifdef CONFIG_ENABLE_NAND_NMBM
-	mtd = nmbm_mtd_get_upper_by_index(0);
-
-	if (mtd)
-		mtd = get_mtd_device(mtd, -1);
-
-	if (!mtd)
-		cprintln(ERROR, "*** NMBM MTD device %u not found! ***", 0);
-#else
-	mtd = get_mtd_device(NULL, 0);
-
-	if (!mtd)
-		cprintln(ERROR, "*** NAND MTD device %u not found! ***", 0);
-#endif
-
-	return mtd;
-}
-
-static int write_part(enum upgrade_part_type pt, const void *data, size_t size,
+static int write_part(const char *partname, const void *data, size_t size,
 		      bool verify)
 {
-	const struct upgrade_part *part = &upgrade_parts[pt];
 	struct mtd_info *mtd;
 	int ret;
 
-	mtd = board_get_mtd_device();
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_generic(mtd, part->offset, size);
-	if (ret)
-		return ret;
+	mtd = get_mtd_device_nm(partname);
+	if (IS_ERR(mtd)) {
+		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
+			 partname);
+		put_mtd_device(mtd);
+		return -PTR_ERR(mtd);
+	}
 
-	ret = mtd_write_generic(mtd, part->offset, part->size, data, size,
-				verify);
+	ret = mtd_update_generic(mtd, data, size, verify);
 
 	put_mtd_device(mtd);
 
 	return ret;
 }
 
-static int write_bl2(void *priv, const struct data_part_entry *dpe,
+int write_bl2(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_BL2, data, size, true);
+	return write_part("bl2", data, size, true);
 }
 
-static int write_fip(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FIP, data, size, true);
+	return write_part("fip", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FW, data, size, true);
+	return write_part("firmware", data, size, true);
 }
 
 static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 			     const void *data, size_t size)
 {
-	struct mtd_info *mtd = board_get_mtd_device();
+	struct mtd_info *mtd;
 	int ret;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
+#ifdef CONFIG_ENABLE_NAND_NMBM
+	mtd = get_mtd_device_nm("nmbm0");
+#else
+	mtd = get_mtd_device(NULL, 0);
+#endif
+
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				false);
+	if (ret) {
+		put_mtd_device(mtd);
 		return ret;
+	}
 
-	ret = mtd_write_generic(mtd, 0, 0, data, size, true);
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
 
 	put_mtd_device(mtd);
 
@@ -127,21 +88,23 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-#if !defined(CONFIG_MTK_SECURE_BOOT) && !defined(CONFIG_ENV_IS_NOWHERE)
-	struct mtd_info *mtd = board_get_mtd_device();
-	int ret;
+	int ret = 0;
+#if !defined(CONFIG_MTK_SECURE_BOOT) && defined(CONFIG_ENV_IS_IN_MTD)
+	struct mtd_info *mtd;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_env(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
+	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
 
 	put_mtd_device(mtd);
+#endif
 
 	return ret;
-#else
-	return 0;
-#endif
 }
 
 static const struct data_part_entry snand_parts[] = {
@@ -157,7 +120,7 @@ static const struct data_part_entry snand_parts[] = {
 		.env_name = "bootfile.fip",
 		.write = write_fip,
 		.post_action = UPGRADE_ACTION_CUSTOM,
-		.do_post_action = erase_env,
+		//.do_post_action = erase_env,
 	},
 	{
 		.name = "Firmware",
@@ -182,18 +145,18 @@ void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 
 int board_boot_default(void)
 {
-	const struct upgrade_part *part = &upgrade_parts[UPGRADE_PART_FW];
-	struct mtd_info *mtd = board_get_mtd_device();
-	int ret;
+	struct mtd_info *mtd;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = boot_from_mtd(mtd, part->offset);
+	mtd = get_mtd_device_nm("firmware");
+	if (!IS_ERR_OR_NULL(mtd)) {
+		put_mtd_device(mtd);
 
-	put_mtd_device(mtd);
+		return boot_from_mtd(mtd, 0);
+	}
 
-	return ret;
+	return -ENODEV;
 }
 
 static const struct bootmenu_entry snand_bootmenu_entries[] = {

uboot-mtk-20220606/board/mediatek/mt7986/bootmenu_snor.c

@@ -5,118 +5,100 @@
  * Author: Weijie Gao <weijie.gao@mediatek.com>
  */
 
+#include <linux/mtd/mtd.h>
 #include <command.h>
 #include <linux/sizes.h>
+#include <mtd.h>
 #include "upgrade_helper.h"
 #include "boot_helper.h"
 #include "autoboot_helper.h"
 #include "colored_print.h"
 
-enum upgrade_part_type {
-	UPGRADE_PART_BL2,
-	UPGRADE_PART_FIP,
-	UPGRADE_PART_FW,
-};
-
-struct upgrade_part {
-	u32 offset;
-	size_t size;
-};
-
-static const struct upgrade_part upgrade_parts[] = {
-	[UPGRADE_PART_BL2] = {
-		.offset = 0,
-		.size = 0x40000,
-	},
-	[UPGRADE_PART_FIP] = {
-		.offset = 0x100000,
-		.size = 0x80000,
-	},
-	[UPGRADE_PART_FW] = {
-		.offset = 0x180000,
-		.size = 0xd900000,
-	},
-};
-
-/*
- * layout: 128k(bl2),2048k(fip),64k(config),256k(factory),-(firmware)
- */
-
-static inline struct spi_flash *board_get_snor_dev(void)
-{
-	return snor_get_dev(CONFIG_SF_DEFAULT_BUS, CONFIG_SF_DEFAULT_CS);
-}
-
-static int write_part(enum upgrade_part_type pt, const void *data, size_t size,
+static int write_part(const char *partname, const void *data, size_t size,
 		      bool verify)
 {
-	const struct upgrade_part *part = &upgrade_parts[pt];
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 	int ret;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = snor_erase_generic(snor, part->offset, size);
-	if (ret)
-		return ret;
+	mtd = get_mtd_device_nm(partname);
+	if (IS_ERR(mtd)) {
+		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
+			 partname);
+		return -PTR_ERR(mtd);
+	}
 
-	ret = snor_write_generic(snor, part->offset, part->size, data, size,
-				 verify);
+	ret = mtd_update_generic(mtd, data, size, verify);
+
+	put_mtd_device(mtd);
 
 	return ret;
 }
 
-static int write_bl2(void *priv, const struct data_part_entry *dpe,
+int write_bl2(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_BL2, data, size, true);
+	return write_part("bl2", data, size, true);
 }
 
-static int write_fip(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FIP, data, size, true);
+	return write_part("fip", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FW, data, size, false);
+	return write_part("firmware", data, size, false);
 }
 
 static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 			     const void *data, size_t size)
 {
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 	int ret;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	mtd = get_mtd_device(NULL, 0);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
 
-	ret = snor_erase_generic(snor, 0, size);
-	if (ret)
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				 false);
+	if (ret) {
+		put_mtd_device(mtd);
 		return ret;
+	}
 
-	return snor_write_generic(snor, 0, 0, data, size, false);
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
+
+	put_mtd_device(mtd);
+
+	return ret;
 }
 
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-#if !defined(CONFIG_ENV_IS_NOWHERE)
-	struct spi_flash *snor;
+	int ret = 0;
+#if defined(CONFIG_ENV_IS_IN_MTD)
+	struct mtd_info *mtd;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	return snor_erase_env(snor, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
-#else
-	return 0;
+	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
+
+	put_mtd_device(mtd);
 #endif
+
+	return ret;
 }
 
 static const struct data_part_entry snor_parts[] = {
@@ -132,7 +114,7 @@ static const struct data_part_entry snor_parts[] = {
 		.env_name = "bootfile.fip",
 		.write = write_fip,
 		.post_action = UPGRADE_ACTION_CUSTOM,
-		.do_post_action = erase_env,
+		//.do_post_action = erase_env,
 	},
 	{
 		.name = "Firmware",
@@ -157,14 +139,18 @@ void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 
 int board_boot_default(void)
 {
-	const struct upgrade_part *part = &upgrade_parts[UPGRADE_PART_FW];
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	return boot_from_snor(snor, part->offset);
+	mtd = get_mtd_device_nm("firmware");
+	if (!IS_ERR_OR_NULL(mtd)) {
+		put_mtd_device(mtd);
+
+		return boot_from_mtd(mtd, 0);
+	}
+
+	return -ENODEV;
 }
 
 static const struct bootmenu_entry snor_bootmenu_entries[] = {

uboot-mtk-20220606/board/mediatek/mt7986/bootmenu_ubi.c

@@ -19,22 +19,21 @@ static int write_part(const char *partname, const void *data, size_t size,
 	struct mtd_info *mtd;
 	int ret;
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 	mtd = get_mtd_device_nm(partname);
 	if (IS_ERR(mtd)) {
 		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
 			 partname);
+		put_mtd_device(mtd);
 		return -PTR_ERR(mtd);
 	}
 
+	ret = mtd_update_generic(mtd, data, size, verify);
+
 	put_mtd_device(mtd);
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
-		return ret;
-
-	return mtd_write_generic(mtd, 0, 0, data, size, verify);
+	return ret;
 }
 
 int write_bl2(void *priv, const struct data_part_entry *dpe,
@@ -68,7 +67,7 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 	struct mtd_info *mtd;
 	int ret;
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 #ifdef CONFIG_ENABLE_NAND_NMBM
 	mtd = get_mtd_device_nm("nmbm0");
@@ -79,34 +78,40 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 	if (IS_ERR(mtd))
 		return -PTR_ERR(mtd);
 
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				false);
+	if (ret) {
+		put_mtd_device(mtd);
+		return ret;
+	}
+
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
+
 	put_mtd_device(mtd);
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
-		return ret;
-
-	return mtd_write_generic(mtd, 0, 0, data, size, true);
+	return ret;
 }
 
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
+	int ret = 0;
 #if !defined(CONFIG_MTK_SECURE_BOOT) && defined(CONFIG_ENV_IS_IN_MTD)
 	struct mtd_info *mtd;
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
-
 	if (IS_ERR(mtd))
 		return -PTR_ERR(mtd);
 
-	put_mtd_device(mtd);
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
 
-	return mtd_erase_env(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
-#else
-	return 0;
+	put_mtd_device(mtd);
 #endif
+
+	return ret;
 }
 
 static const struct data_part_entry ubi_parts[] = {

uboot-mtk-20220606/board/mediatek/mt7988/bootmenu_snand.c

@@ -7,117 +7,78 @@
 
 #include <linux/mtd/mtd.h>
 #include <linux/sizes.h>
+#include <mtd.h>
 #include "upgrade_helper.h"
 #include "boot_helper.h"
 #include "autoboot_helper.h"
 #include "colored_print.h"
 
-#include <nmbm/nmbm.h>
-#include <nmbm/nmbm-mtd.h>
-
-enum upgrade_part_type {
-	UPGRADE_PART_BL2,
-	UPGRADE_PART_FIP,
-	UPGRADE_PART_FW,
-};
-
-struct upgrade_part {
-	u64 offset;
-	size_t size;
-};
-
-static const struct upgrade_part upgrade_parts[] = {
-	[UPGRADE_PART_BL2] = {
-		.offset = 0,
-		.size = 0x100000,
-	},
-	[UPGRADE_PART_FIP] = {
-		.offset = 0x380000,
-		.size = 0x200000,
-	},
-	[UPGRADE_PART_FW] = {
-		.offset = 0x580000,
-		.size = 0x2800000,
-	},
-};
-
-static struct mtd_info *board_get_mtd_device(void)
-{
-	struct mtd_info *mtd;
-
-#ifdef CONFIG_ENABLE_NAND_NMBM
-	mtd = nmbm_mtd_get_upper_by_index(0);
-
-	if (mtd)
-		mtd = get_mtd_device(mtd, -1);
-
-	if (!mtd)
-		cprintln(ERROR, "*** NMBM MTD device %u not found! ***", 0);
-#else
-	mtd = get_mtd_device(NULL, 0);
-
-	if (!mtd)
-		cprintln(ERROR, "*** NAND MTD device %u not found! ***", 0);
-#endif
-
-	return mtd;
-}
-
-static int write_part(enum upgrade_part_type pt, const void *data, size_t size,
+static int write_part(const char *partname, const void *data, size_t size,
 		      bool verify)
 {
-	const struct upgrade_part *part = &upgrade_parts[pt];
 	struct mtd_info *mtd;
 	int ret;
 
-	mtd = board_get_mtd_device();
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_generic(mtd, part->offset, size);
-	if (ret)
-		return ret;
+	mtd = get_mtd_device_nm(partname);
+	if (IS_ERR(mtd)) {
+		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
+			 partname);
+		put_mtd_device(mtd);
+		return -PTR_ERR(mtd);
+	}
 
-	ret = mtd_write_generic(mtd, part->offset, part->size, data, size,
-				verify);
+	ret = mtd_update_generic(mtd, data, size, verify);
 
 	put_mtd_device(mtd);
 
 	return ret;
 }
 
-static int write_bl2(void *priv, const struct data_part_entry *dpe,
+int write_bl2(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_BL2, data, size, true);
+	return write_part("bl2", data, size, true);
 }
 
-static int write_fip(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FIP, data, size, true);
+	return write_part("fip", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FW, data, size, true);
+	return write_part("firmware", data, size, true);
 }
 
 static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 			     const void *data, size_t size)
 {
-	struct mtd_info *mtd = board_get_mtd_device();
+	struct mtd_info *mtd;
 	int ret;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
+#ifdef CONFIG_ENABLE_NAND_NMBM
+	mtd = get_mtd_device_nm("nmbm0");
+#else
+	mtd = get_mtd_device(NULL, 0);
+#endif
+
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				false);
+	if (ret) {
+		put_mtd_device(mtd);
 		return ret;
+	}
 
-	ret = mtd_write_generic(mtd, 0, 0, data, size, true);
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
 
 	put_mtd_device(mtd);
 
@@ -127,21 +88,23 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-#if !defined(CONFIG_MTK_SECURE_BOOT) && !defined(CONFIG_ENV_IS_NOWHERE)
-	struct mtd_info *mtd = board_get_mtd_device();
-	int ret;
+	int ret = 0;
+#if !defined(CONFIG_MTK_SECURE_BOOT) && defined(CONFIG_ENV_IS_IN_MTD)
+	struct mtd_info *mtd;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = mtd_erase_env(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
+	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
 
 	put_mtd_device(mtd);
+#endif
 
 	return ret;
-#else
-	return 0;
-#endif
 }
 
 static const struct data_part_entry snand_parts[] = {
@@ -157,7 +120,7 @@ static const struct data_part_entry snand_parts[] = {
 		.env_name = "bootfile.fip",
 		.write = write_fip,
 		.post_action = UPGRADE_ACTION_CUSTOM,
-		.do_post_action = erase_env,
+		//.do_post_action = erase_env,
 	},
 	{
 		.name = "Firmware",
@@ -182,18 +145,18 @@ void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 
 int board_boot_default(void)
 {
-	const struct upgrade_part *part = &upgrade_parts[UPGRADE_PART_FW];
-	struct mtd_info *mtd = board_get_mtd_device();
-	int ret;
+	struct mtd_info *mtd;
 
-	if (!mtd)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = boot_from_mtd(mtd, part->offset);
+	mtd = get_mtd_device_nm("firmware");
+	if (!IS_ERR_OR_NULL(mtd)) {
+		put_mtd_device(mtd);
 
-	put_mtd_device(mtd);
+		return boot_from_mtd(mtd, 0);
+	}
 
-	return ret;
+	return -ENODEV;
 }
 
 static const struct bootmenu_entry snand_bootmenu_entries[] = {

uboot-mtk-20220606/board/mediatek/mt7988/bootmenu_snor.c

@@ -5,118 +5,99 @@
  * Author: Weijie Gao <weijie.gao@mediatek.com>
  */
 
+#include <linux/mtd/mtd.h>
 #include <command.h>
 #include <linux/sizes.h>
+#include <mtd.h>
 #include "upgrade_helper.h"
 #include "boot_helper.h"
 #include "autoboot_helper.h"
 #include "colored_print.h"
 
-enum upgrade_part_type {
-	UPGRADE_PART_BL2,
-	UPGRADE_PART_FIP,
-	UPGRADE_PART_FW,
-};
-
-struct upgrade_part {
-	u32 offset;
-	size_t size;
-};
-
-static const struct upgrade_part upgrade_parts[] = {
-	[UPGRADE_PART_BL2] = {
-		.offset = 0,
-		.size = 0x40000,
-	},
-	[UPGRADE_PART_FIP] = {
-		.offset = 0x100000,
-		.size = 0x80000,
-	},
-	[UPGRADE_PART_FW] = {
-		.offset = 0x180000,
-		.size = 0xd900000,
-	},
-};
-
-/*
- * layout: 128k(bl2),2048k(fip),64k(config),256k(factory),-(firmware)
- */
-
-static inline struct spi_flash *board_get_snor_dev(void)
-{
-	return snor_get_dev(CONFIG_SF_DEFAULT_BUS, CONFIG_SF_DEFAULT_CS);
-}
-
-static int write_part(enum upgrade_part_type pt, const void *data, size_t size,
+static int write_part(const char *partname, const void *data, size_t size,
 		      bool verify)
 {
-	const struct upgrade_part *part = &upgrade_parts[pt];
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 	int ret;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	ret = snor_erase_generic(snor, part->offset, size);
-	if (ret)
-		return ret;
+	mtd = get_mtd_device_nm(partname);
+	if (IS_ERR(mtd)) {
+		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
+			 partname);
+		return -PTR_ERR(mtd);
+	}
 
-	ret = snor_write_generic(snor, part->offset, part->size, data, size,
-				 verify);
+	ret = mtd_update_generic(mtd, data, size, verify);
+
+	put_mtd_device(mtd);
 
 	return ret;
 }
 
-static int write_bl2(void *priv, const struct data_part_entry *dpe,
+int write_bl2(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_BL2, data, size, true);
+	return write_part("bl2", data, size, true);
 }
 
-static int write_fip(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FIP, data, size, true);
+	return write_part("fip", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
-	return write_part(UPGRADE_PART_FW, data, size, false);
+	return write_part("firmware", data, size, false);
 }
 
 static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 			     const void *data, size_t size)
 {
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 	int ret;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	mtd = get_mtd_device(NULL, 0);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
 
-	ret = snor_erase_generic(snor, 0, size);
-	if (ret)
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				 false);
+	if (ret) {
+		put_mtd_device(mtd);
 		return ret;
+	}
 
-	return snor_write_generic(snor, 0, 0, data, size, false);
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
+
+	put_mtd_device(mtd);
+
+	return ret;
 }
 
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-#if !defined(CONFIG_ENV_IS_NOWHERE)
-	struct spi_flash *snor;
+	int ret = 0;
+#if defined(CONFIG_ENV_IS_IN_MTD)
+	struct mtd_info *mtd;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	return snor_erase_env(snor, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
-#else
-	return 0;
+	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
+	if (IS_ERR(mtd))
+		return -PTR_ERR(mtd);
+
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
+
+	put_mtd_device(mtd);
 #endif
+
+	return ret;
 }
 
 static const struct data_part_entry snor_parts[] = {
@@ -132,7 +113,7 @@ static const struct data_part_entry snor_parts[] = {
 		.env_name = "bootfile.fip",
 		.write = write_fip,
 		.post_action = UPGRADE_ACTION_CUSTOM,
-		.do_post_action = erase_env,
+		//.do_post_action = erase_env,
 	},
 	{
 		.name = "Firmware",
@@ -157,14 +138,18 @@ void board_upgrade_data_parts(const struct data_part_entry **dpes, u32 *count)
 
 int board_boot_default(void)
 {
-	const struct upgrade_part *part = &upgrade_parts[UPGRADE_PART_FW];
-	struct spi_flash *snor;
+	struct mtd_info *mtd;
 
-	snor = board_get_snor_dev();
-	if (!snor)
-		return -ENODEV;
+	gen_mtd_probe_devices();
 
-	return boot_from_snor(snor, part->offset);
+	mtd = get_mtd_device_nm("firmware");
+	if (!IS_ERR_OR_NULL(mtd)) {
+		put_mtd_device(mtd);
+
+		return boot_from_mtd(mtd, 0);
+	}
+
+	return -ENODEV;
 }
 
 static const struct bootmenu_entry snor_bootmenu_entries[] = {

uboot-mtk-20220606/board/mediatek/mt7988/bootmenu_ubi.c

@@ -19,37 +19,36 @@ static int write_part(const char *partname, const void *data, size_t size,
 	struct mtd_info *mtd;
 	int ret;
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 	mtd = get_mtd_device_nm(partname);
 	if (IS_ERR(mtd)) {
 		cprintln(ERROR, "*** MTD partition '%s' not found! ***",
 			 partname);
+		put_mtd_device(mtd);
 		return -PTR_ERR(mtd);
 	}
 
+	ret = mtd_update_generic(mtd, data, size, verify);
+
 	put_mtd_device(mtd);
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
-		return ret;
-
-	return mtd_write_generic(mtd, 0, 0, data, size, verify);
+	return ret;
 }
 
-static int write_bl2(void *priv, const struct data_part_entry *dpe,
+int write_bl2(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
 	return write_part("bl2", data, size, true);
 }
 
-static int write_fip(void *priv, const struct data_part_entry *dpe,
+int write_fip(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
 	return write_part("fip", data, size, true);
 }
 
-static int write_firmware(void *priv, const struct data_part_entry *dpe,
+int write_firmware(void *priv, const struct data_part_entry *dpe,
 			  const void *data, size_t size)
 {
 	int ret;
@@ -68,7 +67,7 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 	struct mtd_info *mtd;
 	int ret;
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 #ifdef CONFIG_ENABLE_NAND_NMBM
 	mtd = get_mtd_device_nm("nmbm0");
@@ -79,34 +78,40 @@ static int write_flash_image(void *priv, const struct data_part_entry *dpe,
 	if (IS_ERR(mtd))
 		return -PTR_ERR(mtd);
 
+	ret = mtd_erase_skip_bad(mtd, 0, mtd->size, mtd->size, NULL, NULL,
+				false);
+	if (ret) {
+		put_mtd_device(mtd);
+		return ret;
+	}
+
+	ret = mtd_write_skip_bad(mtd, 0, size, mtd->size, NULL, data, true);
+
 	put_mtd_device(mtd);
 
-	ret = mtd_erase_generic(mtd, 0, size);
-	if (ret)
-		return ret;
-
-	return mtd_write_generic(mtd, 0, 0, data, size, true);
+	return ret;
 }
 
 static int erase_env(void *priv, const struct data_part_entry *dpe,
 		     const void *data, size_t size)
 {
-#if !defined(CONFIG_MTK_SECURE_BOOT) && !defined(CONFIG_ENV_IS_NOWHERE)
+	int ret = 0;
+#if !defined(CONFIG_MTK_SECURE_BOOT) && defined(CONFIG_ENV_IS_IN_MTD)
 	struct mtd_info *mtd;
 
-	ubi_probe_mtd_devices();
+	gen_mtd_probe_devices();
 
 	mtd = get_mtd_device_nm(CONFIG_ENV_MTD_NAME);
-
 	if (IS_ERR(mtd))
 		return -PTR_ERR(mtd);
 
-	put_mtd_device(mtd);
+	ret = mtd_erase_skip_bad(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE,
+				mtd->size, NULL, "environment", false);
 
-	return mtd_erase_env(mtd, CONFIG_ENV_OFFSET, CONFIG_ENV_SIZE);
-#else
-	return 0;
+	put_mtd_device(mtd);
 #endif
+
+	return ret;
 }
 
 static const struct data_part_entry ubi_parts[] = {