[PATCH] Bump module ref during init.

From: Rusty Russell <rusty@rustcorp.com.au>

__module_get is theoretically allowed on module inside init, since we
already hold an implicit reference.  Currently this BUG()s: make the
reference count explicit, which also simplifies delete path.  Also cleans
up unload path, such that it only drops semaphore when it's actually
sleeping for rmmod --wait.

kernel/module.c

@@ -214,6 +214,8 @@ static void module_unload_init(struct module *mod)
 	INIT_LIST_HEAD(&mod->modules_which_use_me);
 	for (i = 0; i < NR_CPUS; i++)
 		atomic_set(&mod->ref[i].count, 0);
+	/* Hold reference count during initialization. */
+	atomic_set(&mod->ref[smp_processor_id()].count, 1);
 	/* Backwards compatibility macros put refcount during init. */
 	mod->waiter = current;
 }
@@ -462,6 +464,21 @@ void cleanup_module(void)
 }
 EXPORT_SYMBOL(cleanup_module);
 
+static void wait_for_zero_refcount(struct module *mod)
+{
+	/* Since we might sleep for some time, drop the semaphore first */
+	up(&module_mutex);
+	for (;;) {
+		DEBUGP("Looking at refcount...\n");
+		set_current_state(TASK_UNINTERRUPTIBLE);
+		if (module_refcount(mod) == 0)
+			break;
+		schedule();
+	}
+	current->state = TASK_RUNNING;
+	down(&module_mutex);
+}
+
 asmlinkage long
 sys_delete_module(const char __user *name_user, unsigned int flags)
 {
@@ -500,16 +517,6 @@ sys_delete_module(const char __user *name_user, unsigned int flags)
 		goto out;
 	}
 
-	/* Coming up?  Allow force on stuck modules. */
-	if (mod->state == MODULE_STATE_COMING) {
-		forced = try_force(flags);
-		if (!forced) {
-			/* This module can't be removed */
-			ret = -EBUSY;
-			goto out;
-		}
-	}
-
 	/* If it has an init func, it must have an exit func to unload */
 	if ((mod->init != init_module && mod->exit == cleanup_module)
 	    || mod->unsafe) {
@@ -529,35 +536,22 @@ sys_delete_module(const char __user *name_user, unsigned int flags)
 	/* If it's not unused, quit unless we are told to block. */
 	if ((flags & O_NONBLOCK) && module_refcount(mod) != 0) {
 		forced = try_force(flags);
-		if (!forced)
+		if (!forced) {
 			ret = -EWOULDBLOCK;
-	} else {
-		mod->waiter = current;
-		mod->state = MODULE_STATE_GOING;
-	}
 			restart_refcounts();
-
-	if (ret != 0)
 			goto out;
-
-	if (forced)
-		goto destroy;
-
-	/* Since we might sleep for some time, drop the semaphore first */
-	up(&module_mutex);
-	for (;;) {
-		DEBUGP("Looking at refcount...\n");
-		set_current_state(TASK_UNINTERRUPTIBLE);
-		if (module_refcount(mod) == 0)
-			break;
-		schedule();
 		}
-	current->state = TASK_RUNNING;
+	}
 
-	DEBUGP("Regrabbing mutex...\n");
-	down(&module_mutex);
+	/* Mark it as dying. */
+	mod->waiter = current;
+	mod->state = MODULE_STATE_GOING;
+	restart_refcounts();
+
+	/* Never wait if forced. */
+	if (!forced && module_refcount(mod) != 0)
+		wait_for_zero_refcount(mod);
 
- destroy:
 	/* Final destruction now noone is using it. */
 	mod->exit();
 	free_module(mod);
@@ -1453,6 +1447,7 @@ sys_init_module(void __user *umod,
 			printk(KERN_ERR "%s: module is now stuck!\n",
 			       mod->name);
 		else {
+			module_put(mod);
 			down(&module_mutex);
 			free_module(mod);
 			up(&module_mutex);
@@ -1463,6 +1458,8 @@ sys_init_module(void __user *umod,
 	/* Now it's a first class citizen! */
 	down(&module_mutex);
 	mod->state = MODULE_STATE_LIVE;
+	/* Drop initial reference. */
+	module_put(mod);
 	module_free(mod, mod->module_init);
 	mod->module_init = NULL;
 	mod->init_size = 0;