2 files changed, 27 insertions, 26 deletions

diff --git a/include/net/af_vsock.h b/include/net/af_vsock.h
index 7d64d3609ec9..428277869400 100644
--- a/include/net/af_vsock.h
+++ b/include/net/af_vsock.h
@@ -155,7 +155,11 @@ struct vsock_transport {
 
 /**** CORE ****/
 
-int vsock_core_init(const struct vsock_transport *t);
+int __vsock_core_init(const struct vsock_transport *t, struct module *owner);
+static inline int vsock_core_init(const struct vsock_transport *t)
+{
+        return __vsock_core_init(t, THIS_MODULE);
+}
 void vsock_core_exit(void);
 
 /**** UTILS ****/
diff --git a/net/vmw_vsock/af_vsock.c b/net/vmw_vsock/af_vsock.c
index 5adfd94c5b85..85d232bed87d 100644
--- a/net/vmw_vsock/af_vsock.c
+++ b/net/vmw_vsock/af_vsock.c
@@ -1925,9 +1925,23 @@ static struct miscdevice vsock_device = {
         .fops           = &vsock_device_ops,
 };
 
-static int __vsock_core_init(void)
+int __vsock_core_init(const struct vsock_transport *t, struct module *owner)
 {
-        int err;
+        int err = mutex_lock_interruptible(&vsock_register_mutex);
+
+        if (err)
+                return err;
+
+        if (transport) {
+                err = -EBUSY;
+                goto err_busy;
+        }
+
+        /* Transport must be the owner of the protocol so that it can't
+         * unload while there are open sockets.
+         */
+        vsock_proto.owner = owner;
+        transport = t;
 
         vsock_init_tables();
 
@@ -1951,36 +1965,19 @@ static int __vsock_core_init(void)
                 goto err_unregister_proto;
         }
 
+        mutex_unlock(&vsock_register_mutex);
         return 0;
 
 err_unregister_proto:
         proto_unregister(&vsock_proto);
 err_misc_deregister:
         misc_deregister(&vsock_device);
-        return err;
-}
-
-int vsock_core_init(const struct vsock_transport *t)
-{
-        int retval = mutex_lock_interruptible(&vsock_register_mutex);
-        if (retval)
-                return retval;
-
-        if (transport) {
-                retval = -EBUSY;
-                goto out;
-        }
-
-        transport = t;
-        retval = __vsock_core_init();
-        if (retval)
-                transport = NULL;
-
-out:
+        transport = NULL;
+err_busy:
         mutex_unlock(&vsock_register_mutex);
-        return retval;
+        return err;
 }
-EXPORT_SYMBOL_GPL(vsock_core_init);
+EXPORT_SYMBOL_GPL(__vsock_core_init);
 
 void vsock_core_exit(void)
 {
@@ -2000,5 +1997,5 @@ EXPORT_SYMBOL_GPL(vsock_core_exit);
 
 MODULE_AUTHOR("VMware, Inc.");
 MODULE_DESCRIPTION("VMware Virtual Socket Family");
-MODULE_VERSION("1.0.0.0-k");
+MODULE_VERSION("1.0.1.0-k");
 MODULE_LICENSE("GPL v2");