4 files changed, 19 insertions, 11 deletions

diff --git a/net/tipc/core.h b/net/tipc/core.h
index a881f92a8537..c58a1d16563a 100644
--- a/net/tipc/core.h
+++ b/net/tipc/core.h
@@ -56,6 +56,7 @@
 #include <linux/netdevice.h>
 #include <linux/in.h>
 #include <linux/list.h>
+#include <linux/slab.h>
 #include <linux/vmalloc.h>
 
 /*
diff --git a/net/tipc/eth_media.c b/net/tipc/eth_media.c
index 524ba5696d4d..6230d16020c4 100644
--- a/net/tipc/eth_media.c
+++ b/net/tipc/eth_media.c
@@ -38,6 +38,7 @@
 #include <net/tipc/tipc_bearer.h>
 #include <net/tipc/tipc_msg.h>
 #include <linux/netdevice.h>
+#include <linux/slab.h>
 #include <net/net_namespace.h>
 
 #define MAX_ETH_BEARERS         2
diff --git a/net/tipc/ref.c b/net/tipc/ref.c
index 414fc34b8bea..8dea66500cf5 100644
--- a/net/tipc/ref.c
+++ b/net/tipc/ref.c
@@ -153,11 +153,11 @@ void tipc_ref_table_stop(void)
 
 u32 tipc_ref_acquire(void *object, spinlock_t **lock)
 {
-        struct reference *entry;
         u32 index;
         u32 index_mask;
         u32 next_plus_upper;
         u32 ref;
+        struct reference *entry = NULL;
 
         if (!object) {
                 err("Attempt to acquire reference to non-existent object\n");
@@ -175,30 +175,36 @@ u32 tipc_ref_acquire(void *object, spinlock_t **lock)
                 index = tipc_ref_table.first_free;
                 entry = &(tipc_ref_table.entries[index]);
                 index_mask = tipc_ref_table.index_mask;
-                /* take lock in case a previous user of entry still holds it */
-                spin_lock_bh(&entry->lock);
                 next_plus_upper = entry->ref;
                 tipc_ref_table.first_free = next_plus_upper & index_mask;
                 ref = (next_plus_upper & ~index_mask) + index;
-                entry->ref = ref;
-                entry->object = object;
-                *lock = &entry->lock;
         }
         else if (tipc_ref_table.init_point < tipc_ref_table.capacity) {
                 index = tipc_ref_table.init_point++;
                 entry = &(tipc_ref_table.entries[index]);
                 spin_lock_init(&entry->lock);
-                spin_lock_bh(&entry->lock);
                 ref = tipc_ref_table.start_mask + index;
-                entry->ref = ref;
-                entry->object = object;
-                *lock = &entry->lock;
         }
         else {
                 ref = 0;
         }
         write_unlock_bh(&ref_table_lock);
 
+        /*
+         * Grab the lock so no one else can modify this entry
+         * While we assign its ref value & object pointer
+         */
+        if (entry) {
+                spin_lock_bh(&entry->lock);
+                entry->ref = ref;
+                entry->object = object;
+                *lock = &entry->lock;
+                /*
+                 * keep it locked, the caller is responsible
+                 * for unlocking this when they're done with it
+                 */
+        }
+
         return ref;
 }
 
diff --git a/net/tipc/socket.c b/net/tipc/socket.c
index 4b235fc1c70f..cfb20b80b3a1 100644
--- a/net/tipc/socket.c
+++ b/net/tipc/socket.c
@@ -40,9 +40,9 @@
 #include <linux/socket.h>
 #include <linux/errno.h>
 #include <linux/mm.h>
-#include <linux/slab.h>
 #include <linux/poll.h>
 #include <linux/fcntl.h>
+#include <linux/gfp.h>
 #include <asm/string.h>
 #include <asm/atomic.h>
 #include <net/sock.h>