2 files changed, 37 insertions, 7 deletions

diff --git a/drivers/mtd/mtdcore.c b/drivers/mtd/mtdcore.c
index 06ec9f836ae5..f11f55f02413 100644
--- a/drivers/mtd/mtdcore.c
+++ b/drivers/mtd/mtdcore.c
@@ -214,12 +214,23 @@ struct mtd_info *get_mtd_device(struct mtd_info *mtd, int num)
                         ret = NULL;
         }
 
-        if (ret && !try_module_get(ret->owner))
+        if (!ret)
+                goto out_unlock;
+
+        if (!try_module_get(ret->owner)) {
+                ret = NULL;
+                goto out_unlock;
+        }
+
+        if (ret->get_device && ret->get_device(ret)) {
+                module_put(ret->owner);
                 ret = NULL;
+                goto out_unlock;
+        }
 
-        if (ret)
-                ret->usecount++;
+        ret->usecount++;
 
+out_unlock:
         mutex_unlock(&mtd_table_mutex);
         return ret;
 }
@@ -235,8 +246,8 @@ struct mtd_info *get_mtd_device(struct mtd_info *mtd, int num)
 
 struct mtd_info *get_mtd_device_nm(const char *name)
 {
-        int i;
-        struct mtd_info *mtd = ERR_PTR(-ENODEV);
+        int i, err = -ENODEV;
+        struct mtd_info *mtd = NULL;
 
         mutex_lock(&mtd_table_mutex);
 
@@ -247,17 +258,27 @@ struct mtd_info *get_mtd_device_nm(const char *name)
                 }
         }
 
-        if (i == MAX_MTD_DEVICES)
+        if (!mtd)
                 goto out_unlock;
 
         if (!try_module_get(mtd->owner))
                 goto out_unlock;
 
+        if (mtd->get_device) {
+                err = mtd->get_device(mtd);
+                if (err)
+                        goto out_put;
+        }
+
         mtd->usecount++;
+        mutex_unlock(&mtd_table_mutex);
+        return mtd;
 
+out_put:
+        module_put(mtd->owner);
 out_unlock:
         mutex_unlock(&mtd_table_mutex);
-        return mtd;
+        return ERR_PTR(err);
 }
 
 void put_mtd_device(struct mtd_info *mtd)
@@ -266,6 +287,8 @@ void put_mtd_device(struct mtd_info *mtd)
 
         mutex_lock(&mtd_table_mutex);
         c = --mtd->usecount;
+        if (mtd->put_device)
+                mtd->put_device(mtd);
         mutex_unlock(&mtd_table_mutex);
         BUG_ON(c < 0);
 
diff --git a/include/linux/mtd/mtd.h b/include/linux/mtd/mtd.h
index 89e937dfef55..d644e57703ad 100644
--- a/include/linux/mtd/mtd.h
+++ b/include/linux/mtd/mtd.h
@@ -207,6 +207,13 @@ struct mtd_info {
 
         struct module *owner;
         int usecount;
+
+        /* If the driver is something smart, like UBI, it may need to maintain
+         * its own reference counting. The below functions are only for driver.
+         * The driver may register its callbacks. These callbacks are not
+         * supposed to be called by MTD users */
+        int (*get_device) (struct mtd_info *mtd);
+        void (*put_device) (struct mtd_info *mtd);
 };