[RFC PATCH 1/7] mm/damon/core: refcount ops owner module to prevent rmmod UAF

[Ravi Jonnalagadda]

damon_select_ops() copies the registered damon_operations struct into
ctx->ops by value.  After damon_unregister_ops() is called from a
backend module's exit path, the registry slot is cleared but any
surviving ctx still holds function pointers that resolve into the
unloaded module's text.  Restarting kdamond on such a ctx, or invoking
any ops callback, jumps into freed code.

Add a struct module *owner field to damon_operations.  In
damon_select_ops(), take a reference to ops->owner via try_module_get()
after locating the registry entry; on failure return -EBUSY without
binding the ctx.  If the ctx already had an ops bound (re-select
case), drop the previous owner's reference before installing the new
one to keep the refcount balanced.  In damon_destroy_ctx(), release
the reference via module_put(ctx->ops.owner).

In damon_commit_ctx(), the live ops field is overwritten by a value
copy from src.  Balance the refcount when the owner changes: take a
ref on the new owner (return -EBUSY on failure) and put the ref on the
old owner before the assignment.

Built-in ops sets (vaddr, paddr) leave owner = NULL; try_module_get(NULL)
returns true and module_put(NULL) is a no-op.  Loadable backends set
owner = THIS_MODULE in their registration.

Also add damon_unregister_ops() so loadable backends have a clean exit
path.

Signed-off-by: Ravi Jonnalagadda <ravis.opensrc@xxxxxxxxx>
---
 include/linux/damon.h       |  4 ++++
 mm/damon/core.c             | 46 ++++++++++++++++++++++++++++++++++---
 mm/damon/tests/core-kunit.h |  2 +-
 3 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/include/linux/damon.h b/include/linux/damon.h
index df7910a39b407..8e6e1cd89e551 100644
--- a/include/linux/damon.h
+++ b/include/linux/damon.h
@@ -682,6 +682,8 @@ enum damon_ops_id {
  * struct damon_operations - Monitoring operations for given use cases.
  *
  * @id:				Identifier of this operations set.
+ * @owner:			Module that provides this operations set, or NULL
+ *				for built-in ops.
  * @init:			Initialize operations-related data structures.
  * @update:			Update operations-related data structures.
  * @prepare_access_checks:	Prepare next access check of target regions.
@@ -728,6 +730,7 @@ enum damon_ops_id {
  */
 struct damon_operations {
 	enum damon_ops_id id;
+	struct module *owner;
 	void (*init)(struct damon_ctx *context);
 	void (*update)(struct damon_ctx *context);
 	void (*prepare_access_checks)(struct damon_ctx *context);
@@ -1206,6 +1209,7 @@ int damon_commit_ctx(struct damon_ctx *old_ctx, struct damon_ctx *new_ctx);
 int damon_nr_running_ctxs(void);
 bool damon_is_registered_ops(enum damon_ops_id id);
 int damon_register_ops(struct damon_operations *ops);
+int damon_unregister_ops(enum damon_ops_id id);
 int damon_select_ops(struct damon_ctx *ctx, enum damon_ops_id id);
 
 static inline bool damon_target_has_pid(const struct damon_ctx *ctx)
diff --git a/mm/damon/core.c b/mm/damon/core.c
index e4b9adc0a64dd..b605d36b29b1a 100644
--- a/mm/damon/core.c
+++ b/mm/damon/core.c
@@ -12,6 +12,7 @@
 #include <linux/kthread.h>
 #include <linux/memcontrol.h>
 #include <linux/mm.h>
+#include <linux/module.h>
 #include <linux/psi.h>
 #include <linux/sched.h>
 #include <linux/slab.h>
@@ -93,6 +94,31 @@ int damon_register_ops(struct damon_operations *ops)
 	mutex_unlock(&damon_ops_lock);
 	return err;
 }
+EXPORT_SYMBOL_GPL(damon_register_ops);
+
+/**
+ * damon_unregister_ops() - Unregister a monitoring operations set.
+ * @id:	ID of the operations set to unregister.
+ *
+ * Return: 0 on success, negative error code otherwise.
+ */
+int damon_unregister_ops(enum damon_ops_id id)
+{
+	if (id >= NR_DAMON_OPS)
+		return -EINVAL;
+
+	/*
+	 * Callers (typically the owning module exit path) hold a
+	 * module ref via try_module_get() in damon_select_ops(); the
+	 * unregister cannot race with active ctxs because module_exit
+	 * runs only at owner refcount 0.
+	 */
+	mutex_lock(&damon_ops_lock);
+	memset(&damon_registered_ops[id], 0, sizeof(damon_registered_ops[id]));
+	mutex_unlock(&damon_ops_lock);
+	return 0;
+}
+EXPORT_SYMBOL_GPL(damon_unregister_ops);
 
 /**
  * damon_select_ops() - Select a monitoring operations to use with the context.
@@ -112,10 +138,18 @@ int damon_select_ops(struct damon_ctx *ctx, enum damon_ops_id id)
 		return -EINVAL;
 
 	mutex_lock(&damon_ops_lock);
-	if (!__damon_is_registered_ops(id))
+	if (!__damon_is_registered_ops(id)) {
 		err = -EINVAL;
-	else
-		ctx->ops = damon_registered_ops[id];
+		goto out;
+	}
+	if (!try_module_get(damon_registered_ops[id].owner)) {
+		err = -EBUSY;
+		goto out;
+	}
+	/* Drop previous owner ref if this ctx had ops selected before. */
+	module_put(ctx->ops.owner);
+	ctx->ops = damon_registered_ops[id];
+out:
 	mutex_unlock(&damon_ops_lock);
 	return err;
 }
@@ -835,6 +869,7 @@ void damon_destroy_ctx(struct damon_ctx *ctx)
 	damon_for_each_sample_filter_safe(f, next_f, &ctx->sample_control)
 		damon_destroy_sample_filter(f, &ctx->sample_control);
 
+	module_put(ctx->ops.owner);
 	kfree(ctx);
 }
 
@@ -1749,6 +1784,11 @@ int damon_commit_ctx(struct damon_ctx *dst, struct damon_ctx *src)
 			return err;
 	}
 	dst->pause = src->pause;
+	if (src->ops.owner != dst->ops.owner) {
+		if (!try_module_get(src->ops.owner))
+			return -EBUSY;
+		module_put(dst->ops.owner);
+	}
 	dst->ops = src->ops;
 	err = damon_commit_probes(dst, src);
 	if (err)
diff --git a/mm/damon/tests/core-kunit.h b/mm/damon/tests/core-kunit.h
index 0369c717b93db..300659b115602 100644
--- a/mm/damon/tests/core-kunit.h
+++ b/mm/damon/tests/core-kunit.h
@@ -342,7 +342,7 @@ static void damon_test_split_regions_of(struct kunit *test)
 static void damon_test_ops_registration(struct kunit *test)
 {
 	struct damon_ctx *c = damon_new_ctx();
-	struct damon_operations ops = {.id = DAMON_OPS_VADDR}, bak;
+	struct damon_operations ops = {.id = DAMON_OPS_VADDR}, bak = {};
 	bool need_cleanup = false;
 
 	if (!c)
-- 
2.43.0