6 files changed, 339 insertions, 81 deletions

diff --git a/drivers/mtd/mtd_uboot.c b/drivers/mtd/mtd_uboot.c
index 5ca560c9687..d638f700d04 100644
--- a/drivers/mtd/mtd_uboot.c
+++ b/drivers/mtd/mtd_uboot.c
@@ -13,6 +13,29 @@
 
 #define MTD_NAME_MAX_LEN 20
 
+void board_mtdparts_default(const char **mtdids, const char **mtdparts);
+
+static const char *get_mtdids(void)
+{
+	__maybe_unused const char *mtdparts = NULL;
+	const char *mtdids = env_get("mtdids");
+
+	if (mtdids)
+		return mtdids;
+
+#if defined(CONFIG_SYS_MTDPARTS_RUNTIME)
+	board_mtdparts_default(&mtdids, &mtdparts);
+#elif defined(MTDIDS_DEFAULT)
+	mtdids = MTDIDS_DEFAULT;
+#elif defined(CONFIG_MTDIDS_DEFAULT)
+	mtdids = CONFIG_MTDIDS_DEFAULT;
+#endif
+
+	if (mtdids)
+		env_set("mtdids", mtdids);
+
+	return mtdids;
+}
 
 /**
  * mtd_search_alternate_name - Search an alternate name for @mtdname thanks to
@@ -34,7 +57,7 @@ int mtd_search_alternate_name(const char *mtdname, char *altname,
 	const char *mtdids, *equal, *comma, *dev_id, *mtd_id;
 	int dev_id_len, mtd_id_len;
 
-	mtdids = env_get("mtdids");
+	mtdids = get_mtdids();
 	if (!mtdids)
 		return -EINVAL;
 
@@ -92,30 +115,6 @@ static void mtd_probe_uclass_mtd_devs(void) { }
 #endif
 
 #if defined(CONFIG_MTD_PARTITIONS)
-extern void board_mtdparts_default(const char **mtdids,
-				   const char **mtdparts);
-
-static const char *get_mtdids(void)
-{
-	__maybe_unused const char *mtdparts = NULL;
-	const char *mtdids = env_get("mtdids");
-
-	if (mtdids)
-		return mtdids;
-
-#if defined(CONFIG_SYS_MTDPARTS_RUNTIME)
-	board_mtdparts_default(&mtdids, &mtdparts);
-#elif defined(MTDIDS_DEFAULT)
-	mtdids = MTDIDS_DEFAULT;
-#elif defined(CONFIG_MTDIDS_DEFAULT)
-	mtdids = CONFIG_MTDIDS_DEFAULT;
-#endif
-
-	if (mtdids)
-		env_set("mtdids", mtdids);
-
-	return mtdids;
-}
 
 #define MTDPARTS_MAXLEN         512
 
@@ -150,20 +149,74 @@ static const char *get_mtdparts(void)
 	return mtdparts;
 }
 
+static int mtd_del_parts(struct mtd_info *mtd, bool quiet)
+{
+	int ret;
+
+	if (!mtd_has_partitions(mtd))
+		return 0;
+
+	/* do not delete partitions if they are in use. */
+	if (mtd_partitions_used(mtd)) {
+		if (!quiet)
+			printf("\"%s\" partitions still in use, can't delete them\n",
+			       mtd->name);
+		return -EACCES;
+	}
+
+	ret = del_mtd_partitions(mtd);
+	if (ret)
+		return ret;
+
+	return 1;
+}
+
+static bool mtd_del_all_parts_failed;
+
+static void mtd_del_all_parts(void)
+{
+	struct mtd_info *mtd;
+	int ret = 0;
+
+	mtd_del_all_parts_failed = false;
+
+	/*
+	 * It is not safe to remove entries from the mtd_for_each_device loop
+	 * as it uses idr indexes and the partitions removal is done in bulk
+	 * (all partitions of one device at the same time), so break and
+	 * iterate from start each time a new partition is found and deleted.
+	 */
+	do {
+		mtd_for_each_device(mtd) {
+			ret = mtd_del_parts(mtd, false);
+			if (ret > 0)
+				break;
+			else if (ret < 0)
+				mtd_del_all_parts_failed = true;
+		}
+	} while (ret > 0);
+}
+
 int mtd_probe_devices(void)
 {
 	static char *old_mtdparts;
 	static char *old_mtdids;
 	const char *mtdparts = get_mtdparts();
 	const char *mtdids = get_mtdids();
-	bool remaining_partitions = true;
+	const char *mtdparts_next = mtdparts;
 	struct mtd_info *mtd;
 
 	mtd_probe_uclass_mtd_devs();
 
-	/* Check if mtdparts/mtdids changed since last call, otherwise: exit */
+	/*
+	 * Check if mtdparts/mtdids changed, if the MTD dev list was updated
+	 * or if our previous attempt to delete existing partititions failed.
+	 * In any of these cases we want to update the partitions, otherwise,
+	 * everything is up-to-date and we can return 0 directly.
+	 */
 	if ((!mtdparts && !old_mtdparts && !mtdids && !old_mtdids) ||
 	    (mtdparts && old_mtdparts && mtdids && old_mtdids &&
+	     !mtd_dev_list_updated() && !mtd_del_all_parts_failed &&
 	     !strcmp(mtdparts, old_mtdparts) &&
 	     !strcmp(mtdids, old_mtdids)))
 		return 0;
@@ -174,55 +227,55 @@ int mtd_probe_devices(void)
 	old_mtdparts = strdup(mtdparts);
 	old_mtdids = strdup(mtdids);
 
-	/* If at least one partition is still in use, do not delete anything */
-	mtd_for_each_device(mtd) {
-		if (mtd->usecount) {
-			printf("Partition \"%s\" already in use, aborting\n",
-			       mtd->name);
-			return -EACCES;
-		}
-	}
+	/*
+	 * Remove all old parts. Note that partition removal can fail in case
+	 * one of the partition is still being used by an MTD user, so this
+	 * does not guarantee that all old partitions are gone.
+	 */
+	mtd_del_all_parts();
 
 	/*
-	 * Everything looks clear, remove all partitions. It is not safe to
-	 * remove entries from the mtd_for_each_device loop as it uses idr
-	 * indexes and the partitions removal is done in bulk (all partitions of
-	 * one device at the same time), so break and iterate from start each
-	 * time a new partition is found and deleted.
+	 * Call mtd_dev_list_updated() to clear updates generated by our own
+	 * parts removal loop.
 	 */
-	while (remaining_partitions) {
-		remaining_partitions = false;
-		mtd_for_each_device(mtd) {
-			if (!mtd_is_partition(mtd) && mtd_has_partitions(mtd)) {
-				del_mtd_partitions(mtd);
-				remaining_partitions = true;
-				break;
-			}
-		}
-	}
+	mtd_dev_list_updated();
 
 	/* If either mtdparts or mtdids is empty, then exit */
 	if (!mtdparts || !mtdids)
 		return 0;
 
 	/* Start the parsing by ignoring the extra 'mtdparts=' prefix, if any */
-	if (strstr(mtdparts, "mtdparts="))
+	if (!strncmp(mtdparts, "mtdparts=", sizeof("mtdparts=") - 1))
 		mtdparts += 9;
 
 	/* For each MTD device in mtdparts */
-	while (mtdparts[0] != '\0') {
+	for (; mtdparts[0] != '\0'; mtdparts = mtdparts_next) {
 		char mtd_name[MTD_NAME_MAX_LEN], *colon;
 		struct mtd_partition *parts;
-		int mtd_name_len, nparts;
-		int ret;
+		unsigned int mtd_name_len;
+		int nparts, ret;
+
+		mtdparts_next = strchr(mtdparts, ';');
+		if (!mtdparts_next)
+			mtdparts_next = mtdparts + strlen(mtdparts);
+		else
+			mtdparts_next++;
 
 		colon = strchr(mtdparts, ':');
+		if (colon > mtdparts_next)
+			colon = NULL;
+
 		if (!colon) {
 			printf("Wrong mtdparts: %s\n", mtdparts);
 			return -EINVAL;
 		}
 
-		mtd_name_len = colon - mtdparts;
+		mtd_name_len = (unsigned int)(colon - mtdparts);
+		if (mtd_name_len + 1 > sizeof(mtd_name)) {
+			printf("MTD name too long: %s\n", mtdparts);
+			return -EINVAL;
+		}
+
 		strncpy(mtd_name, mtdparts, mtd_name_len);
 		mtd_name[mtd_name_len] = '\0';
 		/* Move the pointer forward (including the ':') */
@@ -249,15 +302,23 @@ int mtd_probe_devices(void)
 			if (ret || IS_ERR_OR_NULL(mtd)) {
 				printf("Could not find a valid device for %s\n",
 				       mtd_name);
-				mtdparts = strchr(mtdparts, ';');
-				if (mtdparts)
-					mtdparts++;
-
+				mtdparts = mtdparts_next;
 				continue;
 			}
 		}
 
 		/*
+		 * Call mtd_del_parts() again, even if it's already been called
+		 * in mtd_del_all_parts(). We need to know if old partitions are
+		 * still around (because they are still being used by someone),
+		 * and if they are, we shouldn't create new partitions, so just
+		 * skip this MTD device and try the next one.
+		 */
+		ret = mtd_del_parts(mtd, true);
+		if (ret < 0)
+			continue;
+
+		/*
 		 * Parse the MTD device partitions. It will update the mtdparts
 		 * pointer, create an array of parts (that must be freed), and
 		 * return the number of partition structures in the array.
@@ -281,6 +342,12 @@ int mtd_probe_devices(void)
 		put_mtd_device(mtd);
 	}
 
+	/*
+	 * Call mtd_dev_list_updated() to clear updates generated by our own
+	 * parts registration loop.
+	 */
+	mtd_dev_list_updated();
+
 	return 0;
 }
 #else
diff --git a/drivers/mtd/mtdcore.c b/drivers/mtd/mtdcore.c
index fb6c779abbf..cb7ca38d074 100644
--- a/drivers/mtd/mtdcore.c
+++ b/drivers/mtd/mtdcore.c
@@ -87,14 +87,17 @@ struct idr_layer {
 
 struct idr {
 	struct idr_layer id[MAX_IDR_ID];
+	bool updated;
 };
 
 #define DEFINE_IDR(name)	struct idr name;
 
 void idr_remove(struct idr *idp, int id)
 {
-	if (idp->id[id].used)
+	if (idp->id[id].used) {
 		idp->id[id].used = 0;
+		idp->updated = true;
+	}
 
 	return;
 }
@@ -134,6 +137,7 @@ int idr_alloc(struct idr *idp, void *ptr, int start, int end, gfp_t gfp_mask)
 		if (idl->used == 0) {
 			idl->used = 1;
 			idl->ptr = ptr;
+			idp->updated = true;
 			return i;
 		}
 		i++;
@@ -155,6 +159,16 @@ struct mtd_info *__mtd_next_device(int i)
 }
 EXPORT_SYMBOL_GPL(__mtd_next_device);
 
+bool mtd_dev_list_updated(void)
+{
+	if (mtd_idr.updated) {
+		mtd_idr.updated = false;
+		return true;
+	}
+
+	return false;
+}
+
 #ifndef __UBOOT__
 static LIST_HEAD(mtd_notifiers);
 
@@ -514,6 +528,13 @@ int del_mtd_device(struct mtd_info *mtd)
 	struct mtd_notifier *not;
 #endif
 
+	ret = del_mtd_partitions(mtd);
+	if (ret) {
+		debug("Failed to delete MTD partitions attached to %s (err %d)\n",
+		      mtd->name, ret);
+		return ret;
+	}
+
 	mutex_lock(&mtd_table_mutex);
 
 	if (idr_find(&mtd_idr, mtd->index) != mtd) {
diff --git a/drivers/mtd/mtdpart.c b/drivers/mtd/mtdpart.c
index 4d2ac8107f0..fd8d8e5ea72 100644
--- a/drivers/mtd/mtdpart.c
+++ b/drivers/mtd/mtdpart.c
@@ -63,6 +63,18 @@ char *kstrdup(const char *s, gfp_t gfp)
 #define MTD_SIZE_REMAINING		(~0LLU)
 #define MTD_OFFSET_NOT_SPECIFIED	(~0LLU)
 
+bool mtd_partitions_used(struct mtd_info *master)
+{
+	struct mtd_info *slave;
+
+	list_for_each_entry(slave, &master->partitions, node) {
+		if (slave->usecount)
+			return true;
+	}
+
+	return false;
+}
+
 /**
  * mtd_parse_partition - Parse @mtdparts partition definition, fill @partition
  *                       with it and update the @mtdparts string pointer.
diff --git a/drivers/mtd/nand/raw/pxa3xx_nand.c b/drivers/mtd/nand/raw/pxa3xx_nand.c
index 4c783f1e1e4..4d2712df4c7 100644
--- a/drivers/mtd/nand/raw/pxa3xx_nand.c
+++ b/drivers/mtd/nand/raw/pxa3xx_nand.c
@@ -195,6 +195,7 @@ struct pxa3xx_nand_info {
 
 	int			cs;
 	int			use_ecc;	/* use HW ECC ? */
+	int			force_raw;	/* prevent use_ecc to be set */
 	int			ecc_bch;	/* using BCH ECC? */
 	int			use_spare;	/* use spare ? */
 	int			need_wait;
@@ -326,14 +327,14 @@ static struct nand_ecclayout ecc_layout_2KB_bch4bit = {
 static struct nand_ecclayout ecc_layout_2KB_bch8bit = {
 	.eccbytes = 64,
 	.eccpos = {
-		64,  65,  66,  67,  68,  69,  70,  71,
-		72,  73,  74,  75,  76,  77,  78,  79,
-		80,  81,  82,  83,  84,  85,  86,  87,
-		88,  89,  90,  91,  92,  93,  94,  95,
-		96,  97,  98,  99,  100, 101, 102, 103,
-		104, 105, 106, 107, 108, 109, 110, 111,
-		112, 113, 114, 115, 116, 117, 118, 119,
-		120, 121, 122, 123, 124, 125, 126, 127},
+		32, 33, 34, 35, 36, 37, 38, 39,
+		40, 41, 42, 43, 44, 45, 46, 47,
+		48, 49, 50, 51, 52, 53, 54, 55,
+		56, 57, 58, 59, 60, 61, 62, 63,
+		64, 65, 66, 67, 68, 69, 70, 71,
+		72, 73, 74, 75, 76, 77, 78, 79,
+		80, 81, 82, 83, 84, 85, 86, 87,
+		88, 89, 90, 91, 92, 93, 94, 95},
 	.oobfree = { {1, 4}, {6, 26} }
 };
 
@@ -579,7 +580,7 @@ static void disable_int(struct pxa3xx_nand_info *info, uint32_t int_mask)
 
 static void drain_fifo(struct pxa3xx_nand_info *info, void *data, int len)
 {
-	if (info->ecc_bch) {
+	if (info->ecc_bch && !info->force_raw) {
 		u32 ts;
 
 		/*
@@ -612,12 +613,22 @@ static void drain_fifo(struct pxa3xx_nand_info *info, void *data, int len)
 
 static void handle_data_pio(struct pxa3xx_nand_info *info)
 {
+	int data_len = info->step_chunk_size;
+
+	/*
+	 * In raw mode, include the spare area and the ECC bytes that are not
+	 * consumed by the controller in the data section. Do not reorganize
+	 * here, do it in the ->read_page_raw() handler instead.
+	 */
+	if (info->force_raw)
+		data_len += info->step_spare_size + info->ecc_size;
+
 	switch (info->state) {
 	case STATE_PIO_WRITING:
 		if (info->step_chunk_size)
 			writesl(info->mmio_base + NDDB,
 				info->data_buff + info->data_buff_pos,
-				DIV_ROUND_UP(info->step_chunk_size, 4));
+				DIV_ROUND_UP(data_len, 4));
 
 		if (info->step_spare_size)
 			writesl(info->mmio_base + NDDB,
@@ -628,7 +639,10 @@ static void handle_data_pio(struct pxa3xx_nand_info *info)
 		if (info->step_chunk_size)
 			drain_fifo(info,
 				   info->data_buff + info->data_buff_pos,
-				   DIV_ROUND_UP(info->step_chunk_size, 4));
+				   DIV_ROUND_UP(data_len, 4));
+
+		if (info->force_raw)
+			break;
 
 		if (info->step_spare_size)
 			drain_fifo(info,
@@ -642,7 +656,7 @@ static void handle_data_pio(struct pxa3xx_nand_info *info)
 	}
 
 	/* Update buffer pointers for multi-page read/write */
-	info->data_buff_pos += info->step_chunk_size;
+	info->data_buff_pos += data_len;
 	info->oob_buff_pos += info->step_spare_size;
 }
 
@@ -796,7 +810,8 @@ static void prepare_start_command(struct pxa3xx_nand_info *info, int command)
 	case NAND_CMD_READ0:
 	case NAND_CMD_READOOB:
 	case NAND_CMD_PAGEPROG:
-		info->use_ecc = 1;
+		if (!info->force_raw)
+			info->use_ecc = 1;
 		break;
 	case NAND_CMD_PARAM:
 		info->use_spare = 0;
@@ -866,7 +881,13 @@ static int prepare_set_command(struct pxa3xx_nand_info *info, int command,
 		 * which is either naked-read or last-read according to the
 		 * state.
 		 */
-		if (mtd->writesize == info->chunk_size) {
+		if (info->force_raw) {
+			info->ndcb0 |= NDCB0_DBC | (NAND_CMD_READSTART << 8) |
+				       NDCB0_LEN_OVRD |
+				       NDCB0_EXT_CMD_TYPE(ext_cmd_type);
+			info->ndcb3 = info->step_chunk_size +
+				      info->step_spare_size + info->ecc_size;
+		} else if (mtd->writesize == info->chunk_size) {
 			info->ndcb0 |= NDCB0_DBC | (NAND_CMD_READSTART << 8);
 		} else if (mtd->writesize > info->chunk_size) {
 			info->ndcb0 |= NDCB0_DBC | (NAND_CMD_READSTART << 8)
@@ -1216,6 +1237,7 @@ static int pxa3xx_nand_read_page_hwecc(struct mtd_info *mtd,
 {
 	struct pxa3xx_nand_host *host = nand_get_controller_data(chip);
 	struct pxa3xx_nand_info *info = host->info_data;
+	int bf;
 
 	chip->read_buf(mtd, buf, mtd->writesize);
 	chip->read_buf(mtd, chip->oob_poi, mtd->oobsize);
@@ -1223,12 +1245,30 @@ static int pxa3xx_nand_read_page_hwecc(struct mtd_info *mtd,
 	if (info->retcode == ERR_CORERR && info->use_ecc) {
 		mtd->ecc_stats.corrected += info->ecc_err_cnt;
 
-	} else if (info->retcode == ERR_UNCORERR) {
+	} else if (info->retcode == ERR_UNCORERR && info->ecc_bch) {
 		/*
-		 * for blank page (all 0xff), HW will calculate its ECC as
-		 * 0, which is different from the ECC information within
-		 * OOB, ignore such uncorrectable errors
+		 * Empty pages will trigger uncorrectable errors. Re-read the
+		 * entire page in raw mode and check for bits not being "1".
+		 * If there are more than the supported strength, then it means
+		 * this is an actual uncorrectable error.
 		 */
+		chip->ecc.read_page_raw(mtd, chip, buf, oob_required, page);
+		bf = nand_check_erased_ecc_chunk(buf, mtd->writesize,
+						 chip->oob_poi, mtd->oobsize,
+						 NULL, 0, chip->ecc.strength);
+		if (bf < 0) {
+			mtd->ecc_stats.failed++;
+		} else if (bf) {
+			mtd->ecc_stats.corrected += bf;
+			info->max_bitflips = max_t(unsigned int,
+						   info->max_bitflips, bf);
+			info->retcode = ERR_CORERR;
+		} else {
+			info->retcode = ERR_NONE;
+		}
+
+	} else if (info->retcode == ERR_UNCORERR && !info->ecc_bch) {
+		/* Raw read is not supported with Hamming ECC engine */
 		if (is_buf_blank(buf, mtd->writesize))
 			info->retcode = ERR_NONE;
 		else
@@ -1238,6 +1278,69 @@ static int pxa3xx_nand_read_page_hwecc(struct mtd_info *mtd,
 	return info->max_bitflips;
 }
 
+static int pxa3xx_nand_read_page_raw(struct mtd_info *mtd,
+				     struct nand_chip *chip, uint8_t *buf,
+				     int oob_required, int page)
+{
+	struct pxa3xx_nand_host *host = chip->priv;
+	struct pxa3xx_nand_info *info = host->info_data;
+	int chunk, ecc_off_buf;
+
+	if (!info->ecc_bch)
+		return -ENOTSUPP;
+
+	/*
+	 * Set the force_raw boolean, then re-call ->cmdfunc() that will run
+	 * pxa3xx_nand_start(), which will actually disable the ECC engine.
+	 */
+	info->force_raw = true;
+	chip->cmdfunc(mtd, NAND_CMD_READ0, 0x00, page);
+
+	ecc_off_buf = (info->nfullchunks * info->spare_size) +
+		      info->last_spare_size;
+	for (chunk = 0; chunk < info->nfullchunks; chunk++) {
+		chip->read_buf(mtd,
+			       buf + (chunk * info->chunk_size),
+			       info->chunk_size);
+		chip->read_buf(mtd,
+			       chip->oob_poi +
+			       (chunk * (info->spare_size)),
+			       info->spare_size);
+		chip->read_buf(mtd,
+			       chip->oob_poi + ecc_off_buf +
+			       (chunk * (info->ecc_size)),
+			       info->ecc_size - 2);
+	}
+
+	if (info->ntotalchunks > info->nfullchunks) {
+		chip->read_buf(mtd,
+			       buf + (info->nfullchunks * info->chunk_size),
+			       info->last_chunk_size);
+		chip->read_buf(mtd,
+			       chip->oob_poi +
+			       (info->nfullchunks * (info->spare_size)),
+			       info->last_spare_size);
+		chip->read_buf(mtd,
+			       chip->oob_poi + ecc_off_buf +
+			       (info->nfullchunks * (info->ecc_size)),
+			       info->ecc_size - 2);
+	}
+
+	info->force_raw = false;
+
+	return 0;
+}
+
+static int pxa3xx_nand_read_oob_raw(struct mtd_info *mtd,
+				    struct nand_chip *chip, int page)
+{
+	/* Invalidate page cache */
+	chip->pagebuf = -1;
+
+	return chip->ecc.read_page_raw(mtd, chip, chip->buffers->databuf, true,
+				       page);
+}
+
 static uint8_t pxa3xx_nand_read_byte(struct mtd_info *mtd)
 {
 	struct nand_chip *chip = mtd_to_nand(mtd);
@@ -1488,7 +1591,7 @@ static int pxa_ecc_init(struct pxa3xx_nand_info *info,
 		info->chunk_size = 1024;
 		info->spare_size = 0;
 		info->last_chunk_size = 1024;
-		info->last_spare_size = 64;
+		info->last_spare_size = 32;
 		info->ecc_size = 32;
 		ecc->mode = NAND_ECC_HW;
 		ecc->size = info->chunk_size;
@@ -1669,6 +1772,8 @@ static int alloc_nand_resource(struct pxa3xx_nand_info *info)
 
 		nand_set_controller_data(chip, host);
 		chip->ecc.read_page	= pxa3xx_nand_read_page_hwecc;
+		chip->ecc.read_page_raw	= pxa3xx_nand_read_page_raw;
+		chip->ecc.read_oob_raw	= pxa3xx_nand_read_oob_raw;
 		chip->ecc.write_page	= pxa3xx_nand_write_page_hwecc;
 		chip->controller        = &info->controller;
 		chip->waitfunc		= pxa3xx_nand_waitfunc;
diff --git a/drivers/mtd/spi/sf_mtd.c b/drivers/mtd/spi/sf_mtd.c
index 58d7e443990..68c36002bee 100644
--- a/drivers/mtd/spi/sf_mtd.c
+++ b/drivers/mtd/spi/sf_mtd.c
@@ -10,6 +10,7 @@
 #include <spi_flash.h>
 
 static struct mtd_info sf_mtd_info;
+static bool sf_mtd_registered;
 static char sf_mtd_name[8];
 
 static int spi_flash_mtd_erase(struct mtd_info *mtd, struct erase_info *instr)
@@ -17,6 +18,9 @@ static int spi_flash_mtd_erase(struct mtd_info *mtd, struct erase_info *instr)
 	struct spi_flash *flash = mtd->priv;
 	int err;
 
+	if (!flash)
+		return -ENODEV;
+
 	instr->state = MTD_ERASING;
 
 	err = spi_flash_erase(flash, instr->addr, instr->len);
@@ -38,6 +42,9 @@ static int spi_flash_mtd_read(struct mtd_info *mtd, loff_t from, size_t len,
 	struct spi_flash *flash = mtd->priv;
 	int err;
 
+	if (!flash)
+		return -ENODEV;
+
 	err = spi_flash_read(flash, from, len, buf);
 	if (!err)
 		*retlen = len;
@@ -51,6 +58,9 @@ static int spi_flash_mtd_write(struct mtd_info *mtd, loff_t to, size_t len,
 	struct spi_flash *flash = mtd->priv;
 	int err;
 
+	if (!flash)
+		return -ENODEV;
+
 	err = spi_flash_write(flash, to, len, buf);
 	if (!err)
 		*retlen = len;
@@ -73,6 +83,17 @@ static int spi_flash_mtd_number(void)
 
 int spi_flash_mtd_register(struct spi_flash *flash)
 {
+	int ret;
+
+	if (sf_mtd_registered) {
+		ret = del_mtd_device(&sf_mtd_info);
+		if (ret)
+			return ret;
+
+		sf_mtd_registered = false;
+	}
+
+	sf_mtd_registered = false;
 	memset(&sf_mtd_info, 0, sizeof(sf_mtd_info));
 	sprintf(sf_mtd_name, "nor%d", spi_flash_mtd_number());
 
@@ -94,10 +115,33 @@ int spi_flash_mtd_register(struct spi_flash *flash)
 	sf_mtd_info.numeraseregions = 0;
 	sf_mtd_info.erasesize = flash->sector_size;
 
-	return add_mtd_device(&sf_mtd_info);
+	ret = add_mtd_device(&sf_mtd_info);
+	if (!ret)
+		sf_mtd_registered = true;
+
+	return ret;
 }
 
 void spi_flash_mtd_unregister(void)
 {
-	del_mtd_device(&sf_mtd_info);
+	int ret;
+
+	if (!sf_mtd_registered)
+		return;
+
+	ret = del_mtd_device(&sf_mtd_info);
+	if (!ret) {
+		sf_mtd_registered = false;
+		return;
+	}
+
+	/*
+	 * Setting mtd->priv to NULL is the best we can do. Thanks to that,
+	 * the MTD layer can still call mtd hooks without risking a
+	 * use-after-free bug. Still, things should be fixed to prevent the
+	 * spi_flash object from being destroyed when del_mtd_device() fails.
+	 */
+	sf_mtd_info.priv = NULL;
+	printf("Failed to unregister MTD %s and the spi_flash object is going away: you're in deep trouble!",
+	       sf_mtd_info.name);
 }
diff --git a/drivers/mtd/spi/sf_probe.c b/drivers/mtd/spi/sf_probe.c
index 5a2e932de8f..00f8558e701 100644
--- a/drivers/mtd/spi/sf_probe.c
+++ b/drivers/mtd/spi/sf_probe.c
@@ -144,6 +144,14 @@ static int spi_flash_std_probe(struct udevice *dev)
 	return spi_flash_probe_slave(flash);
 }
 
+static int spi_flash_std_remove(struct udevice *dev)
+{
+#ifdef CONFIG_SPI_FLASH_MTD
+	spi_flash_mtd_unregister();
+#endif
+	return 0;
+}
+
 static const struct dm_spi_flash_ops spi_flash_std_ops = {
 	.read = spi_flash_std_read,
 	.write = spi_flash_std_write,
@@ -161,6 +169,7 @@ U_BOOT_DRIVER(spi_flash_std) = {
 	.id		= UCLASS_SPI_FLASH,
 	.of_match	= spi_flash_std_ids,
 	.probe		= spi_flash_std_probe,
+	.remove		= spi_flash_std_remove,
 	.priv_auto_alloc_size = sizeof(struct spi_flash),
 	.ops		= &spi_flash_std_ops,
 };