[PATCH v3 2/2] mtd: spi-nor: fix use-after-free on unbind with open handles

[tze.yee.ng@altera.com]

From: Tze Yee Ng <tze.yee.ng@altera.com>

Unbind frees spi_nor and its buffers (devm) while /dev/mtdX may still be
open, so a later read, ioctl, close or mtdblock flush uses freed memory.

On the spi-mem probe path, allocate them with kmalloc/kzalloc and
refcount spi_nor: one reference for registration plus one per open
handle. The last put frees them.

In remove(), unregister debugfs (it reads devres data) and the MTD
device first, so mtdblock can still flush. Then set nor->removed under
the nor lock; prep_and_lock*() return -ENODEV once it is set. Cache the
controller module at probe, as ->_put_device() now runs after unbind.

Legacy controllers keep devres (nor->refcounted is unset).

Signed-off-by: Tze Yee Ng <tze.yee.ng@altera.com>
---
Changes in v3:
- Use a spi-nor kref again (nor->refcnt, nor->refcounted, nor->removed,
  nor->controller_module) and drop the core hooks from v2. Patch 1 is now
  only the module_put() reorder it needs.
- Re-check nor->removed under nor->lock (or the RWW exclusive section) in
  prep_and_lock*(), instead of a lockless check in the core. A check done
  before taking the lock can pass just before remove() and then operate on
  freed dirmaps.
- Cache the controller module at probe: ->_put_device() runs after unbind
  and dereferenced the freed nor->spimem and a NULL dev->driver.
- In remove(), unregister debugfs first (it reads devres data), then call
  mtd_device_unregister() before fencing, so that the mtdblock release
  notifier can still flush its dirty cache. Then set nor->removed, restore
  and drop the registration reference. The debugfs devres action is gone.
- Assign mtd->eraseregions right after allocation so the error path
  in spi_nor_set_mtd_eraseregions() doesn't leak it.
- Free through the kref on probe errors. On a failed mtd_device_register(),
  fence ops first, as a handle may already be open.
- Move the unlock_and_unprep_pe/rd() prototypes to the top of the file.

Changes in v2:
- Drop the spi-nor kref, nor->refcounted, and controller_module.
  Use mtd->_free / mtd->removed from patch 1.
- Drop the nor->removed checks in prep_and_lock*(). The core now
  fences hw ops, new openers, and suspend/resume.
- Keep eraseregions on the same lifetime as params/bouncebuf
  (plain kcalloc when mtd->_free is set).
- Free params/bouncebuf/eraseregions/nor on the probe error path.
  Unregister debugfs from _free instead of a devm action.
- Legacy controllers still use devm (no mtd->_free).
---
 drivers/mtd/spi-nor/core.c    | 173 +++++++++++++++++++++++++++-------
 drivers/mtd/spi-nor/core.h    |   2 +
 drivers/mtd/spi-nor/debugfs.c |   9 +-
 include/linux/mtd/spi-nor.h   |   9 ++
 4 files changed, 153 insertions(+), 40 deletions(-)

diff --git a/drivers/mtd/spi-nor/core.c b/drivers/mtd/spi-nor/core.c
index 8bc117b46e02..76928bd8413e 100644
--- a/drivers/mtd/spi-nor/core.c
+++ b/drivers/mtd/spi-nor/core.c
@@ -45,6 +45,9 @@
 #define SPI_NOR_SRST_SLEEP_MIN 200
 #define SPI_NOR_SRST_SLEEP_MAX 400
 
+static void spi_nor_unlock_and_unprep_pe(struct spi_nor *nor, loff_t start, size_t len);
+static void spi_nor_unlock_and_unprep_rd(struct spi_nor *nor, loff_t start, size_t len);
+
 /**
  * spi_nor_get_cmd_ext() - Get the command opcode extension based on the
  *			   extension type.
@@ -1345,8 +1348,16 @@ int spi_nor_prep_and_lock(struct spi_nor *nor)
 	else
 		ret = wait_event_killable(nor->rww.wait,
 					  spi_nor_rww_start_exclusive(nor));
+	if (ret)
+		return ret;
 
-	return ret;
+	/* remove() sets @removed under the lock, so this check is race-free. */
+	if (nor->removed) {
+		spi_nor_unlock_and_unprep(nor);
+		return -ENODEV;
+	}
+
+	return 0;
 }
 
 void spi_nor_unlock_and_unprep(struct spi_nor *nor)
@@ -1416,8 +1427,15 @@ static int spi_nor_prep_and_lock_pe(struct spi_nor *nor, loff_t start, size_t le
 	else
 		ret = wait_event_killable(nor->rww.wait,
 					  spi_nor_rww_start_pe(nor, start, len));
+	if (ret)
+		return ret;
 
-	return ret;
+	if (nor->removed) {
+		spi_nor_unlock_and_unprep_pe(nor, start, len);
+		return -ENODEV;
+	}
+
+	return 0;
 }
 
 static void spi_nor_unlock_and_unprep_pe(struct spi_nor *nor, loff_t start, size_t len)
@@ -1489,8 +1507,15 @@ static int spi_nor_prep_and_lock_rd(struct spi_nor *nor, loff_t start, size_t le
 	else
 		ret = wait_event_killable(nor->rww.wait,
 					  spi_nor_rww_start_rd(nor, start, len));
+	if (ret)
+		return ret;
 
-	return ret;
+	if (nor->removed) {
+		spi_nor_unlock_and_unprep_rd(nor, start, len);
+		return -ENODEV;
+	}
+
+	return 0;
 }
 
 static void spi_nor_unlock_and_unprep_rd(struct spi_nor *nor, loff_t start, size_t len)
@@ -3186,7 +3211,11 @@ static int spi_nor_init_params(struct spi_nor *nor)
 {
 	int ret;
 
-	nor->params = devm_kzalloc(nor->dev, sizeof(*nor->params), GFP_KERNEL);
+	if (nor->refcounted)
+		nor->params = kzalloc_obj(*nor->params, GFP_KERNEL);
+	else
+		nor->params = devm_kzalloc(nor->dev, sizeof(*nor->params),
+					   GFP_KERNEL);
 	if (!nor->params)
 		return -ENOMEM;
 
@@ -3420,20 +3449,48 @@ static void spi_nor_resume(struct mtd_info *mtd)
 		dev_err(dev, "resume() failed\n");
 }
 
-static int spi_nor_get_device(struct mtd_info *mtd)
+/* Free spi_nor and its buffers; safe on partial init. */
+static void spi_nor_free(struct spi_nor *nor)
+{
+	spi_nor_debugfs_unregister(nor);
+	kfree(nor->mtd.eraseregions);
+	kfree(nor->bouncebuf);
+	kfree(nor->params);
+	kfree(nor);
+}
+
+static void spi_nor_release(struct kref *kref)
+{
+	spi_nor_free(container_of(kref, struct spi_nor, refcnt));
+}
+
+static struct module *spi_nor_controller_owner(struct spi_nor *nor)
 {
-	struct mtd_info *master = mtd_get_master(mtd);
-	struct spi_nor *nor = mtd_to_spi_nor(master);
 	struct device *dev;
 
+	/* The controller may be gone on the last put; use the cached owner. */
+	if (nor->refcounted)
+		return nor->controller_module;
+
 	if (nor->spimem)
 		dev = nor->spimem->spi->controller->dev.parent;
 	else
 		dev = nor->dev;
 
-	if (!try_module_get(dev->driver->owner))
+	return dev->driver->owner;
+}
+
+static int spi_nor_get_device(struct mtd_info *mtd)
+{
+	struct mtd_info *master = mtd_get_master(mtd);
+	struct spi_nor *nor = mtd_to_spi_nor(master);
+
+	if (!try_module_get(spi_nor_controller_owner(nor)))
 		return -ENODEV;
 
+	if (nor->refcounted)
+		kref_get(&nor->refcnt);
+
 	return 0;
 }
 
@@ -3441,14 +3498,12 @@ static void spi_nor_put_device(struct mtd_info *mtd)
 {
 	struct mtd_info *master = mtd_get_master(mtd);
 	struct spi_nor *nor = mtd_to_spi_nor(master);
-	struct device *dev;
 
-	if (nor->spimem)
-		dev = nor->spimem->spi->controller->dev.parent;
-	else
-		dev = nor->dev;
+	module_put(spi_nor_controller_owner(nor));
 
-	module_put(dev->driver->owner);
+	/* May free @nor. */
+	if (nor->refcounted)
+		kref_put(&nor->refcnt, spi_nor_release);
 }
 
 static void spi_nor_restore(struct spi_nor *nor)
@@ -3549,11 +3604,18 @@ static int spi_nor_set_mtd_eraseregions(struct spi_nor *nor)
 	struct mtd_info *mtd = &nor->mtd;
 	u32 erasesize, i;
 
-	mtd_region = devm_kcalloc(nor->dev, map->n_regions, sizeof(*mtd_region),
-				  GFP_KERNEL);
+	if (nor->refcounted)
+		mtd_region = kcalloc(map->n_regions, sizeof(*mtd_region),
+				     GFP_KERNEL);
+	else
+		mtd_region = devm_kcalloc(nor->dev, map->n_regions,
+					  sizeof(*mtd_region), GFP_KERNEL);
 	if (!mtd_region)
 		return -ENOMEM;
 
+	/* Let spi_nor_free() release it on the error path below. */
+	mtd->eraseregions = mtd_region;
+
 	for (i = 0; i < map->n_regions; i++) {
 		erasesize = spi_nor_get_region_erasesize(&region[i],
 							 map->erase_type);
@@ -3566,7 +3628,6 @@ static int spi_nor_set_mtd_eraseregions(struct spi_nor *nor)
 	}
 
 	mtd->numeraseregions = map->n_regions;
-	mtd->eraseregions = mtd_region;
 
 	return 0;
 }
@@ -3655,8 +3716,11 @@ int spi_nor_scan(struct spi_nor *nor, const char *name,
 	 * than 1KB) after spi_nor_scan() returns.
 	 */
 	nor->bouncebuf_size = PAGE_SIZE;
-	nor->bouncebuf = devm_kmalloc(dev, nor->bouncebuf_size,
-				      GFP_KERNEL);
+	if (nor->refcounted)
+		nor->bouncebuf = kmalloc(nor->bouncebuf_size, GFP_KERNEL);
+	else
+		nor->bouncebuf = devm_kmalloc(dev, nor->bouncebuf_size,
+					      GFP_KERNEL);
 	if (!nor->bouncebuf)
 		return -ENOMEM;
 
@@ -3770,6 +3834,24 @@ static int spi_nor_create_write_dirmap(struct spi_nor *nor)
 	return PTR_ERR_OR_ZERO(nor->dirmap.wdesc);
 }
 
+/* Drain in-flight ops, then fence later ones with -ENODEV. */
+static void spi_nor_set_removed(struct spi_nor *nor)
+{
+	if (!spi_nor_use_parallel_locking(nor))
+		mutex_lock(&nor->lock);
+	else
+		wait_event(nor->rww.wait, spi_nor_rww_start_exclusive(nor));
+
+	nor->removed = true;
+
+	if (!spi_nor_use_parallel_locking(nor)) {
+		mutex_unlock(&nor->lock);
+	} else {
+		spi_nor_rww_end_exclusive(nor);
+		wake_up(&nor->rww.wait);
+	}
+}
+
 static int spi_nor_probe(struct spi_mem *spimem)
 {
 	struct spi_device *spi = spimem->spi;
@@ -3788,12 +3870,15 @@ static int spi_nor_probe(struct spi_mem *spimem)
 	if (ret)
 		return ret;
 
-	nor = devm_kzalloc(dev, sizeof(*nor), GFP_KERNEL);
+	nor = kzalloc_obj(*nor, GFP_KERNEL);
 	if (!nor)
 		return -ENOMEM;
 
 	nor->spimem = spimem;
 	nor->dev = dev;
+	nor->controller_module = spi->controller->dev.parent->driver->owner;
+	nor->refcounted = true;
+	kref_init(&nor->refcnt);
 	spi_nor_set_flash_node(nor, dev->of_node);
 
 	spi_mem_set_drvdata(spimem, nor);
@@ -3819,7 +3904,7 @@ static int spi_nor_probe(struct spi_mem *spimem)
 
 	ret = spi_nor_scan(nor, flash_name, &hwcaps);
 	if (ret)
-		return ret;
+		goto err_free;
 
 	spi_nor_debugfs_register(nor);
 
@@ -3830,33 +3915,57 @@ static int spi_nor_probe(struct spi_mem *spimem)
 	 */
 	if (nor->params->page_size > PAGE_SIZE) {
 		nor->bouncebuf_size = nor->params->page_size;
-		devm_kfree(dev, nor->bouncebuf);
-		nor->bouncebuf = devm_kmalloc(dev, nor->bouncebuf_size,
-					      GFP_KERNEL);
-		if (!nor->bouncebuf)
-			return -ENOMEM;
+		kfree(nor->bouncebuf);
+		nor->bouncebuf = kmalloc(nor->bouncebuf_size, GFP_KERNEL);
+		if (!nor->bouncebuf) {
+			ret = -ENOMEM;
+			goto err_free;
+		}
 	}
 
 	ret = spi_nor_create_read_dirmap(nor);
 	if (ret)
-		return ret;
+		goto err_free;
 
 	ret = spi_nor_create_write_dirmap(nor);
 	if (ret)
-		return ret;
+		goto err_free;
 
-	return mtd_device_register(&nor->mtd, data ? data->parts : NULL,
-				   data ? data->nr_parts : 0);
+	ret = mtd_device_register(&nor->mtd, data ? data->parts : NULL,
+				  data ? data->nr_parts : 0);
+	if (ret) {
+		/* A handle may be open on a partition that was already visible. */
+		spi_nor_set_removed(nor);
+		goto err_free;
+	}
+
+	return 0;
+
+err_free:
+	spi_nor_debugfs_unregister(nor);
+	kref_put(&nor->refcnt, spi_nor_release);
+	return ret;
 }
 
 static int spi_nor_remove(struct spi_mem *spimem)
 {
 	struct spi_nor *nor = spi_mem_get_drvdata(spimem);
+	int ret;
+
+	/* debugfs reads devres data (nor->id, erase map) freed after remove() */
+	spi_nor_debugfs_unregister(nor);
+
+	/* Unregister first: notifiers (e.g. mtdblock) flush dirty data here. */
+	ret = mtd_device_unregister(&nor->mtd);
+
+	spi_nor_set_removed(nor);
 
 	spi_nor_restore(nor);
 
-	/* Clean up MTD stuff. */
-	return mtd_device_unregister(&nor->mtd);
+	/* Drop the registration reference; frees @nor unless handles are open. */
+	kref_put(&nor->refcnt, spi_nor_release);
+
+	return ret;
 }
 
 static void spi_nor_shutdown(struct spi_mem *spimem)
diff --git a/drivers/mtd/spi-nor/core.h b/drivers/mtd/spi-nor/core.h
index ba2d1a862c9d..e71b788b275c 100644
--- a/drivers/mtd/spi-nor/core.h
+++ b/drivers/mtd/spi-nor/core.h
@@ -726,9 +726,11 @@ bool spi_nor_is_locked_sr(struct spi_nor *nor, loff_t ofs, u64 len, const u8 *sr
 
 #ifdef CONFIG_DEBUG_FS
 void spi_nor_debugfs_register(struct spi_nor *nor);
+void spi_nor_debugfs_unregister(struct spi_nor *nor);
 void spi_nor_debugfs_shutdown(void);
 #else
 static inline void spi_nor_debugfs_register(struct spi_nor *nor) {}
+static inline void spi_nor_debugfs_unregister(struct spi_nor *nor) {}
 static inline void spi_nor_debugfs_shutdown(void) {}
 #endif
 
diff --git a/drivers/mtd/spi-nor/debugfs.c b/drivers/mtd/spi-nor/debugfs.c
index 288e2866daed..a0639e91cb2d 100644
--- a/drivers/mtd/spi-nor/debugfs.c
+++ b/drivers/mtd/spi-nor/debugfs.c
@@ -286,10 +286,8 @@ static int spi_nor_capabilities_show(struct seq_file *s, void *data)
 }
 DEFINE_SHOW_ATTRIBUTE(spi_nor_capabilities);
 
-static void spi_nor_debugfs_unregister(void *data)
+void spi_nor_debugfs_unregister(struct spi_nor *nor)
 {
-	struct spi_nor *nor = data;
-
 	debugfs_remove(nor->debugfs_root);
 	nor->debugfs_root = NULL;
 }
@@ -299,15 +297,10 @@ static struct dentry *rootdir;
 void spi_nor_debugfs_register(struct spi_nor *nor)
 {
 	struct dentry *d;
-	int ret;
 
 	if (!rootdir)
 		rootdir = debugfs_create_dir(SPI_NOR_DEBUGFS_ROOT, NULL);
 
-	ret = devm_add_action(nor->dev, spi_nor_debugfs_unregister, nor);
-	if (ret)
-		return;
-
 	d = debugfs_create_dir(dev_name(nor->dev), rootdir);
 	nor->debugfs_root = d;
 
diff --git a/include/linux/mtd/spi-nor.h b/include/linux/mtd/spi-nor.h
index 4b92494827b1..b67d720e809d 100644
--- a/include/linux/mtd/spi-nor.h
+++ b/include/linux/mtd/spi-nor.h
@@ -352,6 +352,11 @@ struct spi_nor_flash_parameter;
  * @rww.used_banks:	bitmap of the banks in use
  * @dev:		pointer to an SPI device or an SPI NOR controller device
  * @spimem:		pointer to the SPI memory device
+ * @refcnt:		registration + open handles; the last put frees @nor.
+ *			Only used when @refcounted is set.
+ * @refcounted:		@nor is not devres-managed (spi-mem probe path)
+ * @removed:		set under @lock on unbind; later operations fail
+ * @controller_module:	controller driver module, cached at probe
  * @bouncebuf:		bounce buffer used when the buffer passed by the MTD
  *                      layer is not DMA-able
  * @bouncebuf_size:	size of the bounce buffer
@@ -393,6 +398,10 @@ struct spi_nor {
 	} rww;
 	struct device		*dev;
 	struct spi_mem		*spimem;
+	struct kref		refcnt;
+	bool			refcounted;
+	bool			removed;
+	struct module		*controller_module;
 	u8			*bouncebuf;
 	size_t			bouncebuf_size;
 	u8			*id;
-- 
2.43.7