[PATCH 2/5] workqueue: add support for module-owned work

[Cen Zhang <zzzccc427@gmail.com>]

Queueing a callback on a system workqueue does not take a reference to
the module containing that callback. A caller which releases its last
module reference after queueing work can therefore leave a callback in
unloaded text. Releasing the reference from the callback itself also
leaves its return path unprotected.

Add module_work and schedule_module_work() to hold the callback's owner
from queueing until the callback returns. Run the dispatch and final
module_put() in workqueue core, which remains present when the callback's
module is unloaded. Cache the function and owner before invoking the
callback so that it can free the containing work item.

6LoWPAN needs this for deferred network-device deletion after removing
the last peer.

Assisted-by: LLM
Signed-off-by: Cen Zhang <zzzccc427@gmail.com>
---
 include/linux/workqueue.h | 15 ++++++++++++++
 kernel/workqueue.c        | 43 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/include/linux/workqueue.h b/include/linux/workqueue.h
index c8a36423cb34..9920796c8822 100644
--- a/include/linux/workqueue.h
+++ b/include/linux/workqueue.h
@@ -128,6 +128,14 @@ struct rcu_work {
 	struct workqueue_struct *wq;
 };
 
+struct module;
+
+struct module_work {
+	struct work_struct work;
+	struct module *owner;
+	work_func_t func;
+};
+
 enum wq_affn_scope {
 	WQ_AFFN_DFL,			/* use system default */
 	WQ_AFFN_CPU,			/* one pod per CPU */
@@ -220,6 +228,11 @@ static inline struct rcu_work *to_rcu_work(struct work_struct *work)
 	return container_of(work, struct rcu_work, work);
 }
 
+static inline struct module_work *to_module_work(struct work_struct *work)
+{
+	return container_of(work, struct module_work, work);
+}
+
 struct execute_work {
 	struct work_struct work;
 };
@@ -634,6 +647,8 @@ extern void __flush_workqueue(struct workqueue_struct *wq);
 extern void drain_workqueue(struct workqueue_struct *wq);
 
 extern int schedule_on_each_cpu(work_func_t func);
+bool schedule_module_work(struct module_work *mwork, work_func_t func,
+			  struct module *owner);
 
 int execute_in_process_context(work_func_t fn, struct execute_work *);
 
diff --git a/kernel/workqueue.c b/kernel/workqueue.c
index 1ae3732a2c51..1a16bc5dfb68 100644
--- a/kernel/workqueue.c
+++ b/kernel/workqueue.c
@@ -48,6 +48,7 @@
 #include <linux/hashtable.h>
 #include <linux/rculist.h>
 #include <linux/nodemask.h>
+#include <linux/module.h>
 #include <linux/moduleparam.h>
 #include <linux/uaccess.h>
 #include <linux/sched/isolation.h>
@@ -4808,6 +4809,48 @@ int execute_in_process_context(work_func_t fn, struct execute_work *ew)
 }
 EXPORT_SYMBOL_GPL(execute_in_process_context);
 
+static void module_work_func(struct work_struct *work)
+{
+	struct module_work *mwork = to_module_work(work);
+	struct module *owner = mwork->owner;
+	work_func_t func = mwork->func;
+
+	func(work);
+	module_put(owner);
+}
+
+/**
+ * schedule_module_work - schedule work owned by a module
+ * @mwork: module work to schedule
+ * @func: work function to schedule
+ * @owner: module owning @func
+ *
+ * Take a reference to @owner before scheduling @func. The reference is
+ * released by workqueue core after the callback returns. The callback may
+ * free @mwork. @mwork must not be pending.
+ *
+ * Return: %false if the module is being removed or the work could not be
+ * queued, %true otherwise.
+ */
+bool schedule_module_work(struct module_work *mwork, work_func_t func,
+			  struct module *owner)
+{
+	if (!try_module_get(owner))
+		return false;
+
+	INIT_WORK(&mwork->work, module_work_func);
+	mwork->owner = owner;
+	mwork->func = func;
+
+	if (!schedule_work(&mwork->work)) {
+		module_put(owner);
+		return false;
+	}
+
+	return true;
+}
+EXPORT_SYMBOL_GPL(schedule_module_work);
+
 /**
  * free_workqueue_attrs - free a workqueue_attrs
  * @attrs: workqueue_attrs to free
-- 
2.43.0