[PATCH bpf-next 06/12] bpf: Add a JIT helper for the outgoing stack of kfunc calls

[Yonghong Song <yonghong.song@linux.dev>]

Add bpf_jit_kfunc_stack_slots() helper which will be used
in x86 and arm64 bpf jit.

Signed-off-by: Yonghong Song <yonghong.song@linux.dev>
---
 include/linux/filter.h |  2 ++
 kernel/bpf/core.c      | 31 +++++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/include/linux/filter.h b/include/linux/filter.h
index e288df4124d3..eddcb9987056 100644
--- a/include/linux/filter.h
+++ b/include/linux/filter.h
@@ -1239,6 +1239,8 @@ bool bpf_jit_supports_percpu_insn(void);
 bool bpf_jit_supports_kfunc_call(void);
 bool bpf_jit_supports_kfunc_ret_reg_pair(void);
 bool bpf_jit_supports_kfunc_arg_slot(u32 slots_used, u32 nslots, u32 align);
+u16 bpf_jit_kfunc_stack_slots(const struct bpf_prog *prog, u32 nr_regs,
+			      int (*layout)(const struct btf_func_model *fm, u8 *pos, int max));
 bool bpf_jit_supports_stack_args(void);
 bool bpf_jit_supports_arena_args(void);
 bool bpf_jit_supports_far_kfunc_call(void);
diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
index 5a20d3e1ace9..8fea01377298 100644
--- a/kernel/bpf/core.c
+++ b/kernel/bpf/core.c
@@ -3303,6 +3303,37 @@ bool __weak bpf_jit_supports_kfunc_arg_slot(u32 slots_used, u32 nslots, u32 alig
 	return false;
 }
 
+/*
+ * The most outgoing stack eightbytes any kfunc call in @prog needs.
+ * @layout writes a position per eightbyte and returns how many, or an error.
+ * @nr_regs is how many of those positions are argument registers.
+ */
+u16 bpf_jit_kfunc_stack_slots(const struct bpf_prog *prog, u32 nr_regs,
+			      int (*layout)(const struct btf_func_model *fm, u8 *pos, int max))
+{
+	const struct bpf_insn *insn = prog->insnsi;
+	u8 pos[MAX_BPF_FUNC_ARGS];
+	int i, k, n, slots = 0;
+
+	for (i = 0; i < prog->len; i++, insn++) {
+		const struct btf_func_model *fm;
+
+		if (insn->code != (BPF_JMP | BPF_CALL) ||
+		    insn->src_reg != BPF_PSEUDO_KFUNC_CALL)
+			continue;
+		fm = bpf_jit_find_kfunc_model(prog, insn);
+		if (!fm)
+			continue;
+		n = layout(fm, pos, ARRAY_SIZE(pos));
+		if (n < 0)
+			continue;
+		for (k = 0; k < n; k++)
+			if (pos[k] >= nr_regs)
+				slots = max(slots, pos[k] - (int)nr_regs + 1);
+	}
+	return slots;
+}
+
 bool __weak bpf_jit_supports_stack_args(void)
 {
 	return false;
-- 
2.53.0-Meta