[RFC PATCH bpf-next 6/6] selftests/bpf: Add kfunc call test in gen_prologue and gen_epilogue

[Martin KaFai Lau <martin.lau@linux.dev>]

From: Martin KaFai Lau <martin.lau@kernel.org>

This patch changes the .gen_pro/epilogue of the bpf_testmod_st_ops
to call kfunc. It will call the inc10 and inc100 kfunc.
The value of the PROLOGUE_A and EPILOGUE_A macro are adjusted
to reflect this change also.

The inc100 kfunc is newly added in this patch which does
args->a += 100. Note that it is not in the register_btf_kfunc_id_set(),
so no need to declare in the bpf_testmod_kfunc.h.
It is enclosed with __bpf_kfunc_{start,edn}_defs to avoid the
compiler warning.

Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>
---
 .../selftests/bpf/bpf_testmod/bpf_testmod.c   | 42 ++++++++++++++++++-
 .../bpf/prog_tests/struct_ops_syscall.c       |  5 ++-
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c b/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c
index 4c75346376d9..6f745d29e124 100644
--- a/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c
+++ b/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c
@@ -966,6 +966,16 @@ __bpf_kfunc int bpf_kfunc_st_ops_inc10(struct st_ops_args *args)
 	return args->a;
 }
 
+__bpf_kfunc_start_defs();
+
+__bpf_kfunc int bpf_kfunc_st_ops_inc100(struct st_ops_args *args)
+{
+	args->a += 100;
+	return args->a;
+}
+
+__bpf_kfunc_end_defs();
+
 BTF_KFUNCS_START(bpf_testmod_check_kfunc_ids)
 BTF_ID_FLAGS(func, bpf_testmod_test_mod_kfunc)
 BTF_ID_FLAGS(func, bpf_kfunc_call_test1)
@@ -1140,6 +1150,10 @@ static int bpf_test_mod_st_ops__test_pro_epilogue(struct st_ops_args *args)
 	return 0;
 }
 
+BTF_ID_LIST(st_ops_epilogue_kfunc_list)
+BTF_ID(func, bpf_kfunc_st_ops_inc10)
+BTF_ID(func, bpf_kfunc_st_ops_inc100)
+
 static int st_ops_gen_prologue(struct bpf_insn *insn_buf, bool direct_write,
 			       const struct bpf_prog *prog, struct module **module)
 {
@@ -1153,13 +1167,28 @@ static int st_ops_gen_prologue(struct bpf_insn *insn_buf, bool direct_write,
 	 * r7 = r6->a;
 	 * r7 += 1000;
 	 * r6->a = r7;
+	 * r7 = r1;
+	 * r1 = r6;
+	 * bpf_kfunc_st_ops_in10(r1)
+	 * r1 = r6;
+	 * bpf_kfunc_st_ops_in100(r1)
+	 * r1 = r7;
 	 */
 	*insn++ = BPF_LDX_MEM(BPF_DW, BPF_REG_6, BPF_REG_1, 0);
 	*insn++ = BPF_LDX_MEM(BPF_W, BPF_REG_7, BPF_REG_6, offsetof(struct st_ops_args, a));
 	*insn++ = BPF_ALU32_IMM(BPF_ADD, BPF_REG_7, 1000);
 	*insn++ = BPF_STX_MEM(BPF_W, BPF_REG_6, BPF_REG_7, offsetof(struct st_ops_args, a));
+	*insn++ = BPF_MOV64_REG(BPF_REG_7, BPF_REG_1);
+	*insn++ = BPF_MOV64_REG(BPF_REG_1, BPF_REG_6);
+	*insn++ = BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, BPF_PSEUDO_KFUNC_CALL, 0,
+			       st_ops_epilogue_kfunc_list[0]);
+	*insn++ = BPF_MOV64_REG(BPF_REG_1, BPF_REG_6);
+	*insn++ = BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, BPF_PSEUDO_KFUNC_CALL, 0,
+			       st_ops_epilogue_kfunc_list[1]);
+	*insn++ = BPF_MOV64_REG(BPF_REG_1, BPF_REG_7);
 	*insn++ = prog->insnsi[0];
 
+	*module = THIS_MODULE;
 	return insn - insn_buf;
 }
 
@@ -1177,7 +1206,10 @@ static int st_ops_gen_epilogue(struct bpf_insn *insn_buf, const struct bpf_prog
 	 * r6 = r1->a;
 	 * r6 += 10000;
 	 * r1->a = r6;
-	 * r0 = r6;
+	 * r6 = r1;
+	 * bpf_kfunc_st_ops_in10(r1)
+	 * r1 = r6;
+	 * bpf_kfunc_st_ops_in100(r1)
 	 * r0 *= 2;
 	 * BPF_EXIT;
 	 */
@@ -1186,10 +1218,16 @@ static int st_ops_gen_epilogue(struct bpf_insn *insn_buf, const struct bpf_prog
 	*insn++ = BPF_LDX_MEM(BPF_W, BPF_REG_6, BPF_REG_1, offsetof(struct st_ops_args, a));
 	*insn++ = BPF_ALU32_IMM(BPF_ADD, BPF_REG_6, 10000);
 	*insn++ = BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_6, offsetof(struct st_ops_args, a));
-	*insn++ = BPF_MOV32_REG(BPF_REG_0, BPF_REG_6);
+	*insn++ = BPF_MOV64_REG(BPF_REG_6, BPF_REG_1);
+	*insn++ = BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, BPF_PSEUDO_KFUNC_CALL, 0,
+			       st_ops_epilogue_kfunc_list[0]);
+	*insn++ = BPF_MOV64_REG(BPF_REG_1, BPF_REG_6);
+	*insn++ = BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, BPF_PSEUDO_KFUNC_CALL, 0,
+			       st_ops_epilogue_kfunc_list[1]);
 	*insn++ = BPF_ALU32_IMM(BPF_MUL, BPF_REG_0, 2);
 	*insn++ = BPF_EXIT_INSN();
 
+	*module = THIS_MODULE;
 	return insn - insn_buf;
 }
 
diff --git a/tools/testing/selftests/bpf/prog_tests/struct_ops_syscall.c b/tools/testing/selftests/bpf/prog_tests/struct_ops_syscall.c
index a293a35b0dcc..2a73066adbf5 100644
--- a/tools/testing/selftests/bpf/prog_tests/struct_ops_syscall.c
+++ b/tools/testing/selftests/bpf/prog_tests/struct_ops_syscall.c
@@ -3,10 +3,11 @@
 #include <test_progs.h>
 #include "struct_ops_syscall.skel.h"
 
-#define EPILOGUE_A  10000
-#define PROLOGUE_A   1000
+#define KFUNC_A100    100
 #define KFUNC_A10      10
 #define SUBPROG_A       1
+#define EPILOGUE_A (10000 + KFUNC_A100 + KFUNC_A10)
+#define PROLOGUE_A  (1000 + KFUNC_A100 + KFUNC_A10)
 
 #define SUBPROG_TEST_MAIN	SUBPROG_A
 #define KFUNC_TEST_MAIN		(KFUNC_A10 + SUBPROG_A)
-- 
2.43.5