memstick: avert possible race condition between idr_pre_get and idr_get_new

Implement the usual pattern around idr_pre_get() and idr_get_new() to
handlethe situation where another thread concurrently steals this thread's
idr_pre_get() preallocation.

Signed-off-by: Alex Dubov <oakad@yahoo.com>
Signed-off-by: Andrew Morton <akpm@linux-foundation.org>
Signed-off-by: Linus Torvalds <torvalds@linux-foundation.org>

2 files changed, 15 insertions, 9 deletions

diff --git a/drivers/memstick/core/memstick.c b/drivers/memstick/core/memstick.c
index 4303b7ef73e2..e9a3eab7b0cf 100644
--- a/drivers/memstick/core/memstick.c
+++ b/drivers/memstick/core/memstick.c
@@ -511,14 +511,18 @@ int memstick_add_host(struct memstick_host *host)
 {
         int rc;
 
-        if (!idr_pre_get(&memstick_host_idr, GFP_KERNEL))
-                return -ENOMEM;
+        while (1) {
+                if (!idr_pre_get(&memstick_host_idr, GFP_KERNEL))
+                        return -ENOMEM;
 
-        spin_lock(&memstick_host_lock);
-        rc = idr_get_new(&memstick_host_idr, host, &host->id);
-        spin_unlock(&memstick_host_lock);
-        if (rc)
-                return rc;
+                spin_lock(&memstick_host_lock);
+                rc = idr_get_new(&memstick_host_idr, host, &host->id);
+                spin_unlock(&memstick_host_lock);
+                if (!rc)
+                        break;
+                else if (rc != -EAGAIN)
+                        return rc;
+        }
 
         dev_set_name(&host->dev, "memstick%u", host->id);
 
diff --git a/drivers/memstick/core/mspro_block.c b/drivers/memstick/core/mspro_block.c
index 02362eccc588..a167938d4517 100644
--- a/drivers/memstick/core/mspro_block.c
+++ b/drivers/memstick/core/mspro_block.c
@@ -1206,10 +1206,12 @@ static int mspro_block_init_disk(struct memstick_dev *card)
 
         msb->page_size = be16_to_cpu(sys_info->unit_size);
 
-        if (!idr_pre_get(&mspro_block_disk_idr, GFP_KERNEL))
+        mutex_lock(&mspro_block_disk_lock);
+        if (!idr_pre_get(&mspro_block_disk_idr, GFP_KERNEL)) {
+                mutex_unlock(&mspro_block_disk_lock);
                 return -ENOMEM;
+        }
 
-        mutex_lock(&mspro_block_disk_lock);
         rc = idr_get_new(&mspro_block_disk_idr, card, &disk_id);
         mutex_unlock(&mspro_block_disk_lock);