[PATCHv2 6/6] mm/zpool: prevent zbud/zsmalloc from unloading when used

[Dan Streetman <ddstreet@ieee.org>]

Add try_module_get() to zpool_create_pool(), and module_put() to
zpool_destroy_pool().  Without module usage counting, the driver module(s)
could be unloaded while their pool(s) were active, resulting in an oops
when zpool tried to access them.

Signed-off-by: Dan Streetman <ddstreet@ieee.org>
Cc: Seth Jennings <sjennings@variantweb.net>
Cc: Minchan Kim <minchan@kernel.org>
Cc: Nitin Gupta <ngupta@vflare.org>
Cc: Weijie Yang <weijie.yang@samsung.com>
---

Changes since v1 : https://lkml.org/lkml/2014/5/24/134
  -add owner field to struct zpool_driver, pointing to driver module
  -move module usage counting from zbud/zsmalloc into zpool

 include/linux/zpool.h |  5 +++++
 mm/zbud.c             |  1 +
 mm/zpool.c            | 22 +++++++++++++++-------
 mm/zsmalloc.c         |  1 +
 4 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/include/linux/zpool.h b/include/linux/zpool.h
index a528f7c..49bd02b 100644
--- a/include/linux/zpool.h
+++ b/include/linux/zpool.h
@@ -176,6 +176,7 @@ u64 zpool_get_total_size(struct zpool *pool);
  */
 struct zpool_driver {
 	char *type;
+	struct module *owner;
 	struct list_head list;
 
 	void *(*create)(gfp_t gfp, struct zpool_ops *ops);
@@ -203,6 +204,10 @@ void zpool_register_driver(struct zpool_driver *driver);
 /**
  * zpool_unregister_driver() - unregister a zpool implementation.
  * @driver:	driver to unregister.
+ *
+ * Module usage counting is used to prevent using a driver
+ * while/after unloading.  Please only call unregister from
+ * module exit function.
  */
 void zpool_unregister_driver(struct zpool_driver *driver);
 
diff --git a/mm/zbud.c b/mm/zbud.c
index 645379e..440bab7 100644
--- a/mm/zbud.c
+++ b/mm/zbud.c
@@ -184,6 +184,7 @@ u64 zbud_zpool_total_size(void *pool)
 
 static struct zpool_driver zbud_zpool_driver = {
 	.type =		"zbud",
+	.owner =	THIS_MODULE,
 	.create =	zbud_zpool_create,
 	.destroy =	zbud_zpool_destroy,
 	.malloc =	zbud_zpool_malloc,
diff --git a/mm/zpool.c b/mm/zpool.c
index 578c379..119f340 100644
--- a/mm/zpool.c
+++ b/mm/zpool.c
@@ -72,15 +72,24 @@ static struct zpool_driver *zpool_get_driver(char *type)
 {
 	struct zpool_driver *driver;
 
-	assert_spin_locked(&drivers_lock);
+	spin_lock(&drivers_lock);
 	list_for_each_entry(driver, &drivers_head, list) {
-		if (!strcmp(driver->type, type))
-			return driver;
+		if (!strcmp(driver->type, type)) {
+			bool got = try_module_get(driver->owner);
+			spin_unlock(&drivers_lock);
+			return got ? driver : NULL;
+		}
 	}
 
+	spin_unlock(&drivers_lock);
 	return NULL;
 }
 
+static void zpool_put_driver(struct zpool_driver *driver)
+{
+	module_put(driver->owner);
+}
+
 struct zpool *zpool_create_pool(char *type, gfp_t flags,
 			struct zpool_ops *ops)
 {
@@ -89,15 +98,11 @@ struct zpool *zpool_create_pool(char *type, gfp_t flags,
 
 	pr_info("creating pool type %s\n", type);
 
-	spin_lock(&drivers_lock);
 	driver = zpool_get_driver(type);
-	spin_unlock(&drivers_lock);
 
 	if (!driver) {
 		request_module(type);
-		spin_lock(&drivers_lock);
 		driver = zpool_get_driver(type);
-		spin_unlock(&drivers_lock);
 	}
 
 	if (!driver) {
@@ -108,6 +113,7 @@ struct zpool *zpool_create_pool(char *type, gfp_t flags,
 	zpool = kmalloc(sizeof(*zpool), GFP_KERNEL);
 	if (!zpool) {
 		pr_err("couldn't create zpool - out of memory\n");
+		zpool_put_driver(driver);
 		return NULL;
 	}
 
@@ -118,6 +124,7 @@ struct zpool *zpool_create_pool(char *type, gfp_t flags,
 
 	if (!zpool->pool) {
 		pr_err("couldn't create %s pool\n", type);
+		zpool_put_driver(driver);
 		kfree(zpool);
 		return NULL;
 	}
@@ -139,6 +146,7 @@ void zpool_destroy_pool(struct zpool *zpool)
 	list_del(&zpool->list);
 	spin_unlock(&pools_lock);
 	zpool->driver->destroy(zpool->pool);
+	zpool_put_driver(zpool->driver);
 	kfree(zpool);
 }
 
diff --git a/mm/zsmalloc.c b/mm/zsmalloc.c
index feba644..ae3a28f 100644
--- a/mm/zsmalloc.c
+++ b/mm/zsmalloc.c
@@ -303,6 +303,7 @@ u64 zs_zpool_total_size(void *pool)
 
 static struct zpool_driver zs_zpool_driver = {
 	.type =		"zsmalloc",
+	.owner =	THIS_MODULE,
 	.create =	zs_zpool_create,
 	.destroy =	zs_zpool_destroy,
 	.malloc =	zs_zpool_malloc,
-- 
1.8.3.1