[PATCH bpf-next v6 01/11] kernel: Implement try_module_get_live

[Kumar Kartikeya Dwivedi <memxor@gmail.com>]

Refactor shared functionality between strong_try_module_get and
try_module_get into a common helper, and expose try_module_get_live
that returns a bool similar to try_module_get.

It will be used in the next patch for btf_try_get_module, to eliminate a
race between module __init function invocation and module_put from BPF
side.

Cc: Luis Chamberlain <mcgrof@kernel.org>
Cc: Jessica Yu <jeyu@kernel.org>
Cc: linux-kernel@vger.kernel.org
Cc: linux-modules@vger.kernel.org
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 include/linux/module.h | 26 +++++++++++++++++++-------
 kernel/module.c        | 20 ++++++++------------
 2 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/include/linux/module.h b/include/linux/module.h
index c9f1200b2312..eb83aaeaa76e 100644
--- a/include/linux/module.h
+++ b/include/linux/module.h
@@ -608,17 +608,17 @@ void symbol_put_addr(void *addr);
 /* Sometimes we know we already have a refcount, and it's easier not
    to handle the error case (which only happens with rmmod --wait). */
 extern void __module_get(struct module *module);
-
-/* This is the Right Way to get a module: if it fails, it's being removed,
- * so pretend it's not there. */
-extern bool try_module_get(struct module *module);
-
+extern int __try_module_get(struct module *module, bool strong);
 extern void module_put(struct module *module);
 
 #else /*!CONFIG_MODULE_UNLOAD*/
-static inline bool try_module_get(struct module *module)
+static inline int __try_module_get(struct module *module, bool strong)
 {
-	return !module || module_is_live(module);
+	if (module && !module_is_live(module))
+		return -ENOENT;
+	if (strong && module && module->state == MODULE_STATE_COMING)
+		return -EBUSY;
+	return 0;
 }
 static inline void module_put(struct module *module)
 {
@@ -631,6 +631,18 @@ static inline void __module_get(struct module *module)
 
 #endif /* CONFIG_MODULE_UNLOAD */
 
+/* This is the Right Way to get a module: if it fails, it's being removed,
+ * so pretend it's not there. */
+static inline bool try_module_get(struct module *module)
+{
+	return !__try_module_get(module, false);
+}
+/* Only take reference for modules which have fully initialized */
+static inline bool try_module_get_live(struct module *module)
+{
+	return !__try_module_get(module, true);
+}
+
 /* This is a #define so the string doesn't get put in every .o file */
 #define module_name(mod)			\
 ({						\
diff --git a/kernel/module.c b/kernel/module.c
index 84a9141a5e15..a9bb0a5576c8 100644
--- a/kernel/module.c
+++ b/kernel/module.c
@@ -318,12 +318,7 @@ EXPORT_SYMBOL(unregister_module_notifier);
 static inline int strong_try_module_get(struct module *mod)
 {
 	BUG_ON(mod && mod->state == MODULE_STATE_UNFORMED);
-	if (mod && mod->state == MODULE_STATE_COMING)
-		return -EBUSY;
-	if (try_module_get(mod))
-		return 0;
-	else
-		return -ENOENT;
+	return __try_module_get(mod, true);
 }
 
 static inline void add_taint_module(struct module *mod, unsigned flag,
@@ -1066,24 +1061,25 @@ void __module_get(struct module *module)
 }
 EXPORT_SYMBOL(__module_get);
 
-bool try_module_get(struct module *module)
+int __try_module_get(struct module *module, bool strong)
 {
-	bool ret = true;
+	int ret = 0;
 
 	if (module) {
 		preempt_disable();
+		if (strong && module->state == MODULE_STATE_COMING)
+			ret = -EBUSY;
 		/* Note: here, we can fail to get a reference */
-		if (likely(module_is_live(module) &&
+		else if (likely(module_is_live(module) &&
 			   atomic_inc_not_zero(&module->refcnt) != 0))
 			trace_module_get(module, _RET_IP_);
 		else
-			ret = false;
-
+			ret = -ENOENT;
 		preempt_enable();
 	}
 	return ret;
 }
-EXPORT_SYMBOL(try_module_get);
+EXPORT_SYMBOL(__try_module_get);
 
 void module_put(struct module *module)
 {
-- 
2.34.1