[PATCH v3 09/19] mm/ksw: add probe management helpers

Jinchao Wang
2025-09-10T05:32:13Z

Provide ksw_stack_init() and ksw_stack_exit() to manage entry and
exit probes for the target function from ksw_get_config().
Use atomic PID tracking to ensure singleton watch.

Signed-off-by: Jinchao Wang <wangjinchao600@gmail.com>
---
mm/kstackwatch/kstackwatch.h | 4 ++
mm/kstackwatch/stack.c | 99 ++++++++++++++++++++++++++++++++++++
2 files changed, 103 insertions(+)

diff --git a/mm/kstackwatch/kstackwatch.h b/mm/kstackwatch/kstackwatch.h
index 2fa377843f17..79ca40e69268 100644
--- a/mm/kstackwatch/kstackwatch.h
+++ b/mm/kstackwatch/kstackwatch.h
@@ -38,6 +38,10 @@ struct ksw_config {
// singleton, only modified in kernel.c
const struct ksw_config *ksw_get_config(void);

+/* stack management */
+int ksw_stack_init(void);
+void ksw_stack_exit(void);
+
/* watch management */
int ksw_watch_init(void);
void ksw_watch_exit(void);
diff --git a/mm/kstackwatch/stack.c b/mm/kstackwatch/stack.c
index cec594032515..72409156458f 100644
--- a/mm/kstackwatch/stack.c
+++ b/mm/kstackwatch/stack.c
@@ -1 +1,100 @@
// SPDX-License-Identifier: GPL-2.0
+#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
+
+#include <linux/atomic.h>
+#include <linux/fprobe.h>
+#include <linux/kprobes.h>
+#include <linux/printk.h>
+#include <linux/spinlock.h>
+
+#include "kstackwatch.h"
+
+static struct kprobe entry_probe;
+static struct fprobe exit_probe;
+#define INVALID_PID -1
+static atomic_t ksw_stack_pid = ATOMIC_INIT(INVALID_PID);
+
+static int ksw_stack_prepare_watch(struct pt_regs *regs,
+ const struct ksw_config *config,
+ u64 *watch_addr, u64 *watch_len)
+{
+ /* implement logic will be added in following patches */
+ *watch_addr = 0;
+ *watch_len = 0;
+ return 0;
+}
+
+static void ksw_stack_entry_handler(struct kprobe *p, struct pt_regs *regs,
+ unsigned long flags)
+{
+ u64 watch_addr;
+ u64 watch_len;
+ int ret;
+
+ if (atomic_cmpxchg(&ksw_stack_pid, INVALID_PID, current->pid) !=
+ INVALID_PID)
+ return;
+
+ ret = ksw_stack_prepare_watch(regs, ksw_get_config(), &watch_addr,
+ &watch_len);
+ if (ret) {
+ atomic_set(&ksw_stack_pid, INVALID_PID);
+ pr_err("failed to prepare watch target: %d\n", ret);
+ return;
+ }
+
+ ret = ksw_watch_on(watch_addr, watch_len);
+ if (ret) {
+ atomic_set(&ksw_stack_pid, INVALID_PID);
+ pr_err("failed to watch on addr:0x%llx len:%llu %d\n",
+ watch_addr, watch_len, ret);
+ return;
+ }
+}
+
+static void ksw_stack_exit_handler(struct fprobe *fp, unsigned long ip,
+ unsigned long ret_ip,
+ struct ftrace_regs *regs, void *data)
+{
+ if (atomic_read(&ksw_stack_pid) != current->pid)
+ return;
+
+ ksw_watch_off();
+
+ atomic_set(&ksw_stack_pid, INVALID_PID);
+}
+
+int ksw_stack_init(void)
+{
+ int ret;
+ char *symbuf = NULL;
+
+ memset(&entry_probe, 0, sizeof(entry_probe));
+ entry_probe.symbol_name = ksw_get_config()->function;
+ entry_probe.offset = ksw_get_config()->ip_offset;
+ entry_probe.post_handler = ksw_stack_entry_handler;
+ ret = register_kprobe(&entry_probe);
+ if (ret) {
+ pr_err("Failed to register kprobe ret %d\n", ret);
+ return ret;
+ }
+
+ memset(&exit_probe, 0, sizeof(exit_probe));
+ exit_probe.exit_handler = ksw_stack_exit_handler;
+ symbuf = (char *)ksw_get_config()->function;
+
+ ret = register_fprobe_syms(&exit_probe, (const char **)&symbuf, 1);
+ if (ret < 0) {
+ pr_err("register_fprobe_syms fail %d\n", ret);
+ unregister_kprobe(&entry_probe);
+ return ret;
+ }
+
+ return 0;
+}
+
+void ksw_stack_exit(void)
+{
+ unregister_fprobe(&exit_probe);
+ unregister_kprobe(&entry_probe);
+}
--
2.43.0