[PATCH v2 08/18] mm/ksw: add stack probe support

[Jinchao Wang <wangjinchao600@gmail.com>]

Introduce stack management for KStackWatch using kprobes and fprobes to
enable dynamic watch switching:

- Entry: prepare target address/length and enable watch
- Exit: disable watch

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

diff --git a/mm/kstackwatch/kstackwatch.h b/mm/kstackwatch/kstackwatch.h
index 13ef8c79f855..bc8664af4fa6 100644
--- a/mm/kstackwatch/kstackwatch.h
+++ b/mm/kstackwatch/kstackwatch.h
@@ -38,6 +38,10 @@ struct ksw_config {
 
 extern bool panic_on_catch;
 
+/* stack management */
+int ksw_stack_init(struct ksw_config *config);
+void ksw_stack_exit(void);
+
 /* watch management */
 int ksw_watch_init(struct ksw_config *config);
 void ksw_watch_exit(void);
diff --git a/mm/kstackwatch/stack.c b/mm/kstackwatch/stack.c
index cec594032515..3b72177315cc 100644
--- a/mm/kstackwatch/stack.c
+++ b/mm/kstackwatch/stack.c
@@ -1 +1,92 @@
 // SPDX-License-Identifier: GPL-2.0
+
+#include <linux/fprobe.h>
+#include <linux/kprobes.h>
+#include <linux/sched.h>
+#include <linux/spinlock.h>
+
+#include "kstackwatch.h"
+
+struct ksw_config *probe_config;
+
+/* prepare watch_addr and watch_len for watch */
+static int ksw_stack_prepare_watch(struct pt_regs *regs,
+				   struct ksw_config *config, u64 *watch_addr,
+				   u64 *watch_len)
+{
+	/* TODO: implement logic */
+	*watch_addr = 0;
+	*watch_len = 0;
+	return 0;
+}
+
+static struct kprobe entry_probe;
+static struct fprobe exit_probe_fprobe;
+
+static void ksw_stack_entry_handler(struct kprobe *p, struct pt_regs *regs,
+				    unsigned long flags)
+{
+	int ret;
+	u64 watch_addr;
+	u64 watch_len;
+
+	ret = ksw_stack_prepare_watch(regs, probe_config, &watch_addr,
+				      &watch_len);
+	if (ret) {
+		pr_err("KSW: failed to prepare watch target: %d\n", ret);
+		return;
+	}
+
+	ret = ksw_watch_on(watch_addr, watch_len);
+	if (ret) {
+		pr_err("KSW: failed to watch on addr:0x%llx len:%llx %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)
+{
+	ksw_watch_off();
+}
+
+int ksw_stack_init(struct ksw_config *config)
+{
+	int ret;
+	char *symbuf = NULL;
+
+	/* Setup entry probe */
+	memset(&entry_probe, 0, sizeof(entry_probe));
+	entry_probe.symbol_name = config->function;
+	entry_probe.offset = config->ip_offset;
+	entry_probe.post_handler = ksw_stack_entry_handler;
+	probe_config = config;
+	ret = register_kprobe(&entry_probe);
+	if (ret < 0) {
+		pr_err("KSW: Failed to register kprobe ret %d\n", ret);
+		return ret;
+	}
+
+	/* Setup exit probe */
+	memset(&exit_probe_fprobe, 0, sizeof(exit_probe_fprobe));
+	exit_probe_fprobe.exit_handler = ksw_stack_exit_handler;
+	symbuf = probe_config->function;
+
+	ret = register_fprobe_syms(&exit_probe_fprobe, (const char **)&symbuf,
+				   1);
+	if (ret < 0) {
+		pr_err("KSW: register_fprobe_syms fail %d\n", ret);
+		unregister_kprobe(&entry_probe);
+		return ret;
+	}
+
+	return 0;
+}
+
+void ksw_stack_exit(void)
+{
+	unregister_fprobe(&exit_probe_fprobe);
+	unregister_kprobe(&entry_probe);
+}
-- 
2.43.0