[U-Boot] [PATCH v3 10/11] mtd: sf: Unregister the MTD device prior to removing the spi_flash obj

[Boris Brezillon <boris.brezillon@bootlin.com>]

The DM implementation of spi_flash_free() does not unregister the MTD
device before removing the spi dev object. This leads to a use-after-free
bug when the MTD device is later accessed by a MTD user (observed when
attaching the device to UBI after env_sf_load() has called
spi_flash_free()).

Implement ->remove() and call spi_flash_mtd_unregister() from there.

Fixes: 9fe6d8716e09 ("mtd, spi: Add MTD layer driver")
Signed-off-by: Boris Brezillon <boris.brezillon@bootlin.com>
---
Changes in v3:
- New patch
---
 drivers/mtd/spi/sf_probe.c | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/drivers/mtd/spi/sf_probe.c b/drivers/mtd/spi/sf_probe.c
index 94fde2ae7a36..4d7320fe8c14 100644
--- a/drivers/mtd/spi/sf_probe.c
+++ b/drivers/mtd/spi/sf_probe.c
@@ -137,6 +137,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,
@@ -153,6 +161,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,
 };
-- 
2.17.1