[RFC PATCH 05/13] mm/kstackwatch: Add atomic HWBP arm/disarm operations

[Jinchao Wang <wangjinchao600@gmail.com>]

Implement the critical atomic operations for dynamically arming and
disarming hardware breakpoints across all CPUs without allocation
overhead.

This patch adds the core functionality that enables kstackwatch to
operate in atomic contexts (such as kprobe handlers):

Key features:
1. ksw_watch_on() - Atomically arm breakpoints on all CPUs with
   specified address and length
2. ksw_watch_off() - Disarm all breakpoints by resetting to dummy marker
3. HWBP updates using arch_reinstall_hw_breakpoint()
4. SMP-safe coordination using work queues and async function calls

The implementation uses a hybrid approach for SMP coordination:
- Current CPU: Direct function call for immediate effect
- Other CPUs: Asynchronous smp_call_function_single_async() for
  non-blocking operation in queue worker

This enables the kprobe handlers (added in subsequent patches) to
instantly arm breakpoints on function entry and disarm on exit,
providing real-time stack corruption detection without the performance
penalties or atomic context limitations of traditional approaches.

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

diff --git a/mm/kstackwatch/kstackwatch.h b/mm/kstackwatch/kstackwatch.h
index 256574cd9cb2..910f49014715 100644
--- a/mm/kstackwatch/kstackwatch.h
+++ b/mm/kstackwatch/kstackwatch.h
@@ -41,5 +41,7 @@ extern bool panic_on_catch;
 /* watch management */
 int ksw_watch_init(struct ksw_config *config);
 void ksw_watch_exit(void);
+int ksw_watch_on(u64 watch_addr, u64 watch_len);
+void ksw_watch_off(void);
 
 #endif /* _KSTACKWATCH_H */
diff --git a/mm/kstackwatch/watch.c b/mm/kstackwatch/watch.c
index 5cc2dfef140b..7ab247531961 100644
--- a/mm/kstackwatch/watch.c
+++ b/mm/kstackwatch/watch.c
@@ -1,4 +1,6 @@
 // SPDX-License-Identifier: GPL-2.0
+
+#include "linux/printk.h"
 #include <linux/kprobes.h>
 #include <linux/hw_breakpoint.h>
 #include <linux/perf_event.h>
@@ -11,11 +13,24 @@
 
 #include "kstackwatch.h"
 
+#define MAX_STACK_ENTRIES 64
+
 struct perf_event *__percpu *watch_events;
 struct ksw_config *watch_config;
+static DEFINE_SPINLOCK(watch_lock);
 
 static unsigned long long marker;
 
+struct watch_worker {
+	struct work_struct work;
+	int original_cpu;
+} myworker;
+
+static void ksw_watch_on_local_cpu(void *useless);
+
+static DEFINE_PER_CPU(call_single_data_t,
+		      hwbp_csd) = CSD_INIT(ksw_watch_on_local_cpu, NULL);
+
 /* Enhanced breakpoint handler with watch identification */
 static void ksw_watch_handler(struct perf_event *bp,
 			      struct perf_sample_data *data,
@@ -31,6 +46,59 @@ static void ksw_watch_handler(struct perf_event *bp,
 		panic("KSW: Stack corruption detected");
 }
 
+/* Setup single hardware breakpoint on current CPU */
+static void ksw_watch_on_local_cpu(void *useless)
+{
+	struct perf_event *bp;
+	int cpu = smp_processor_id();
+	int ret;
+
+	bp = *per_cpu_ptr(watch_events, cpu);
+	if (!bp)
+		return;
+
+	/* Update breakpoint address */
+	ret = hw_breakpoint_arch_parse(bp, &bp->attr, counter_arch_bp(bp));
+	if (ret) {
+		pr_err("KSW: Failed to parse HWBP for CPU %d ret %d\n", cpu,
+		       ret);
+		return;
+	}
+	ret = arch_reinstall_hw_breakpoint(bp);
+	if (ret) {
+		pr_err("KSW: Failed to install HWBP on CPU %d ret %d\n", cpu,
+		       ret);
+		return;
+	}
+
+	if (bp->attr.bp_addr == (unsigned long)&marker) {
+		pr_info("KSW: HWBP disarmed on CPU %d\n", cpu);
+	} else {
+		pr_info("KSW: HWBP armed on CPU %d at 0x%px (len %llu)\n", cpu,
+			(void *)bp->attr.bp_addr, bp->attr.bp_len);
+	}
+}
+
+static void ksw_watch_on_work_fn(struct work_struct *work)
+{
+	struct watch_worker *worker =
+		container_of(work, struct watch_worker, work);
+	int original_cpu = READ_ONCE(worker->original_cpu);
+	int local_cpu = smp_processor_id();
+	call_single_data_t *csd;
+	int cpu;
+
+	for_each_online_cpu(cpu) {
+		if (cpu == original_cpu)
+			continue;
+		if (cpu == local_cpu)
+			continue;
+		csd = &per_cpu(hwbp_csd, cpu);
+		smp_call_function_single_async(cpu, csd);
+	}
+	ksw_watch_on_local_cpu(NULL);
+}
+
 /* Initialize hardware breakpoint  */
 int ksw_watch_init(struct ksw_config *config)
 {
@@ -50,6 +118,8 @@ int ksw_watch_init(struct ksw_config *config)
 		return ret;
 	}
 
+	/* Initialize work structure */
+	INIT_WORK(&myworker.work, ksw_watch_on_work_fn);
 	watch_config = config;
 	pr_info("KSW: HWBP  initialized\n");
 	return 0;
@@ -63,3 +133,51 @@ void ksw_watch_exit(void)
 
 	pr_info("KSW: HWBP  cleaned up\n");
 }
+
+/* Legacy API: Arm single hardware breakpoint (backward compatibility) */
+int ksw_watch_on(u64 watch_addr, u64 watch_len)
+{
+	struct perf_event *bp;
+	unsigned long flags;
+	int cpu;
+
+	if (!watch_addr) {
+		pr_err("KSW: Invalid address for arming HWBP\n");
+		return -EINVAL;
+	}
+
+	spin_lock_irqsave(&watch_lock, flags);
+
+	/* Check if already armed - only need to check one CPU since all share same addr */
+	bp = *this_cpu_ptr(watch_events);
+	if (bp->attr.bp_addr != 0 &&
+	    bp->attr.bp_addr != (unsigned long)&marker && // installted
+	    watch_addr != (unsigned long)&marker) { //restore
+		spin_unlock_irqrestore(&watch_lock, flags);
+		return -EBUSY;
+	}
+
+	/* Update address in all minimal breakpoint structures */
+	for_each_possible_cpu(cpu) {
+		bp = *per_cpu_ptr(watch_events, cpu);
+		WRITE_ONCE(bp->attr.bp_addr, watch_addr);
+		WRITE_ONCE(bp->attr.bp_len, watch_len);
+	}
+
+	WRITE_ONCE(myworker.original_cpu, smp_processor_id());
+
+	spin_unlock_irqrestore(&watch_lock, flags);
+
+	/* Then install on all CPUs */
+	/* Run on current CPU directly */
+	queue_work(system_highpri_wq, &myworker.work);
+	ksw_watch_on_local_cpu(NULL);
+	return 0;
+}
+
+void ksw_watch_off(void)
+{
+	pr_info("KSW: Disarming all HWBPs\n");
+	ksw_watch_on((unsigned long)&marker, sizeof(marker));
+	pr_info("KSW: All HWBPs disarmed\n");
+}
-- 
2.43.0