2 files changed, 55 insertions, 2 deletions

diff --git a/drivers/mtd/spi/sf_mtd.c b/drivers/mtd/spi/sf_mtd.c
index 58d7e44399..68c36002be 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 5a2e932de8..00f8558e70 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,
 };