LoongArch/ftrace: Add HAVE_DYNAMIC_FTRACE_WITH_REGS support

This patch implements CONFIG_DYNAMIC_FTRACE_WITH_REGS on LoongArch,
which allows a traced function's arguments (and some other registers)
to be captured into a struct pt_regs, allowing these to be inspected
and modified.

Co-developed-by: Jinyang He <hejinyang@loongson.cn>
Signed-off-by: Jinyang He <hejinyang@loongson.cn>
Signed-off-by: Qing Zhang <zhangqing@loongson.cn>
Signed-off-by: Huacai Chen <chenhuacai@loongson.cn>

4 files changed, 51 insertions, 2 deletions

diff --git a/arch/loongarch/Kconfig b/arch/loongarch/Kconfig
index 6e9aaa747ef7..f2d1b2aef2d4 100644
--- a/arch/loongarch/Kconfig
+++ b/arch/loongarch/Kconfig
@@ -89,6 +89,7 @@ config LOONGARCH
 	select HAVE_DEBUG_STACKOVERFLOW
 	select HAVE_DMA_CONTIGUOUS
 	select HAVE_DYNAMIC_FTRACE
+	select HAVE_DYNAMIC_FTRACE_WITH_REGS
 	select HAVE_EBPF_JIT
 	select HAVE_EXIT_THREAD
 	select HAVE_FAST_GUP
diff --git a/arch/loongarch/include/asm/ftrace.h b/arch/loongarch/include/asm/ftrace.h
index 09ff0f84663d..dd4a0c8efd24 100644
--- a/arch/loongarch/include/asm/ftrace.h
+++ b/arch/loongarch/include/asm/ftrace.h
@@ -23,6 +23,8 @@ extern void prepare_ftrace_return(unsigned long self_addr, unsigned long callsit
 struct dyn_ftrace;
 struct dyn_arch_ftrace { };
 
+#define ARCH_SUPPORTS_FTRACE_OPS 1
+
 #define ftrace_init_nop ftrace_init_nop
 int ftrace_init_nop(struct module *mod, struct dyn_ftrace *rec);
 
diff --git a/arch/loongarch/kernel/ftrace_dyn.c b/arch/loongarch/kernel/ftrace_dyn.c
index 5a801c328e2a..d6f30918f94f 100644
--- a/arch/loongarch/kernel/ftrace_dyn.c
+++ b/arch/loongarch/kernel/ftrace_dyn.c
@@ -28,6 +28,21 @@ static int ftrace_modify_code(unsigned long pc, u32 old, u32 new, bool validate)
 	return 0;
 }
 
+#ifdef CONFIG_DYNAMIC_FTRACE_WITH_REGS
+int ftrace_modify_call(struct dyn_ftrace *rec, unsigned long old_addr, unsigned long addr)
+{
+	u32 old, new;
+	unsigned long pc;
+
+	pc = rec->ip + LOONGARCH_INSN_SIZE;
+
+	new = larch_insn_gen_bl(pc, addr);
+	old = larch_insn_gen_bl(pc, old_addr);
+
+	return ftrace_modify_code(pc, old, new, true);
+}
+#endif /* CONFIG_DYNAMIC_FTRACE_WITH_REGS */
+
 int ftrace_update_ftrace_func(ftrace_func_t func)
 {
 	u32 new;
diff --git a/arch/loongarch/kernel/mcount_dyn.S b/arch/loongarch/kernel/mcount_dyn.S
index cce3daa2eb1e..bbabf06244c2 100644
--- a/arch/loongarch/kernel/mcount_dyn.S
+++ b/arch/loongarch/kernel/mcount_dyn.S
@@ -27,7 +27,7 @@
  * follows the LoongArch's psABI as well.
  */
 
-	.macro  ftrace_regs_entry
+	.macro  ftrace_regs_entry allregs=0
 	PTR_ADDI	sp, sp, -PT_SIZE
 	PTR_S		t0, sp, PT_R1  /* Save parent ra at PT_R1(RA) */
 	PTR_S		a0, sp, PT_R4
@@ -39,6 +39,30 @@
 	PTR_S		a6, sp, PT_R10
 	PTR_S		a7, sp, PT_R11
 	PTR_S		fp, sp, PT_R22
+	.if \allregs
+	PTR_S		tp, sp, PT_R2
+	PTR_S		t0, sp, PT_R12
+	PTR_S		t1, sp, PT_R13
+	PTR_S		t2, sp, PT_R14
+	PTR_S		t3, sp, PT_R15
+	PTR_S		t4, sp, PT_R16
+	PTR_S		t5, sp, PT_R17
+	PTR_S		t6, sp, PT_R18
+	PTR_S		t7, sp, PT_R19
+	PTR_S		t8, sp, PT_R20
+	PTR_S		u0, sp, PT_R21
+	PTR_S		s0, sp, PT_R23
+	PTR_S		s1, sp, PT_R24
+	PTR_S		s2, sp, PT_R25
+	PTR_S		s3, sp, PT_R26
+	PTR_S		s4, sp, PT_R27
+	PTR_S		s5, sp, PT_R28
+	PTR_S		s6, sp, PT_R29
+	PTR_S		s7, sp, PT_R30
+	PTR_S		s8, sp, PT_R31
+	/* Clear it for later use as a flag sometimes. */
+	PTR_S		zero, sp, PT_R0
+	.endif
 	PTR_S		ra, sp, PT_ERA /* Save trace function ra at PT_ERA */
 	PTR_ADDI	t8, sp, PT_SIZE
 	PTR_S		t8, sp, PT_R3
@@ -85,10 +109,17 @@ ftrace_common_return:
 SYM_CODE_END(ftrace_common)
 
 SYM_CODE_START(ftrace_caller)
-	ftrace_regs_entry
+	ftrace_regs_entry allregs=0
 	b		ftrace_common
 SYM_CODE_END(ftrace_caller)
 
+#ifdef CONFIG_DYNAMIC_FTRACE_WITH_REGS
+SYM_CODE_START(ftrace_regs_caller)
+	ftrace_regs_entry allregs=1
+	b		ftrace_common
+SYM_CODE_END(ftrace_regs_caller)
+#endif
+
 #ifdef CONFIG_FUNCTION_GRAPH_TRACER
 SYM_CODE_START(ftrace_graph_caller)
 	PTR_L		a0, sp, PT_ERA