[PATCH bpf-next v4 6/7] bpf: Support call-site kfunc specialization for far calls

[Emil Tsalapatis <emil@etsalapatis.com>]

Far-call JITs keep the kfunc BTF ID in the call immediate,
so it cannot distinguish multiple specialized targets for
the same kfunc and BTF object.

Store the finalized target descriptor index in the call offset.
Use that index for address and function-model lookup, and keep
the descriptor table in verification order so the indices remain
stable.

Signed-off-by: Emil Tsalapatis <emil@etsalapatis.com>
---
 include/linux/bpf.h          |  4 ++--
 include/linux/bpf_verifier.h |  2 ++
 kernel/bpf/fixups.c          | 18 ++++++++++++++----
 kernel/bpf/verifier.c        | 35 ++++++++++++++++++++++-------------
 4 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 7df74e47ecb0..5a9580b1769d 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -3283,7 +3283,7 @@ const struct btf_func_model *
 bpf_jit_find_kfunc_model(const struct bpf_prog *prog,
 			 const struct bpf_insn *insn);
 int bpf_get_kfunc_addr(const struct bpf_prog *prog, u32 func_id,
-		       u16 btf_fd_idx, u8 **func_addr);
+		       u16 desc_idx, u8 **func_addr);
 
 struct bpf_core_ctx {
 	struct bpf_verifier_log *log;
@@ -3626,7 +3626,7 @@ bpf_jit_find_kfunc_model(const struct bpf_prog *prog,
 
 static inline int
 bpf_get_kfunc_addr(const struct bpf_prog *prog, u32 func_id,
-		   u16 btf_fd_idx, u8 **func_addr)
+		   u16 desc_idx, u8 **func_addr)
 {
 	return -ENOTSUPP;
 }
diff --git a/include/linux/bpf_verifier.h b/include/linux/bpf_verifier.h
index e36936936418..da219eb0a9cb 100644
--- a/include/linux/bpf_verifier.h
+++ b/include/linux/bpf_verifier.h
@@ -1734,6 +1734,8 @@ struct bpf_kfunc_desc_tab {
 	 * descriptors used for verifier lookups. Call specialization may append
 	 * immutable descriptors for additional targets. Near-call JITs look up
 	 * descriptors by imm and offset after do_misc_fixups() sorts the table.
+	 * Far-call JITs use the descriptor index stored in the finalized call's
+	 * off field, so their table remains in verification order.
 	 *
 	 * Grown one entry at a time by bpf_add_kfunc_call() and during
 	 * call specialization.
diff --git a/kernel/bpf/fixups.c b/kernel/bpf/fixups.c
index 2add8001c3ec..b9f76eee5016 100644
--- a/kernel/bpf/fixups.c
+++ b/kernel/bpf/fixups.c
@@ -140,6 +140,13 @@ bpf_jit_find_kfunc_model(const struct bpf_prog *prog,
 	struct bpf_kfunc_desc_tab *tab;
 
 	tab = prog->aux->kfunc_tab;
+	if (bpf_jit_supports_far_kfunc_call()) {
+		if (insn->off < 0 || insn->off >= tab->nr_descs)
+			return NULL;
+		res = &tab->descs[insn->off];
+		return res->func_id == insn->imm ? &res->func_model : NULL;
+	}
+
 	res = bsearch(&desc, tab->descs, tab->nr_descs,
 		      sizeof(tab->descs[0]), kfunc_desc_cmp_by_imm_off);
 
@@ -2517,11 +2524,14 @@ int bpf_do_misc_fixups(struct bpf_verifier_env *env)
 		}
 	}
 
-	ret = sort_kfunc_descs_by_imm_off(env);
-	if (ret)
-		return ret;
+	/*
+	 * Do not change kfunc desc position into the table for far JIT.
+	 * because we use the indices in the instructions.
+	 */
+	if (bpf_jit_supports_far_kfunc_call())
+		return 0;
 
-	return 0;
+	return sort_kfunc_descs_by_imm_off(env);
 }
 
 static struct bpf_prog *inline_bpf_loop(struct bpf_verifier_env *env,
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 5e7c589991e9..12e33a568a3e 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -2575,12 +2575,16 @@ find_kfunc_desc(const struct bpf_prog *prog, u32 func_id, u16 offset)
 }
 
 int bpf_get_kfunc_addr(const struct bpf_prog *prog, u32 func_id,
-		       u16 btf_fd_idx, u8 **func_addr)
+		       u16 desc_idx, u8 **func_addr)
 {
+	struct bpf_kfunc_desc_tab *tab;
 	const struct bpf_kfunc_desc *desc;
 
-	desc = find_kfunc_desc(prog, func_id, btf_fd_idx);
-	if (!desc)
+	tab = prog->aux->kfunc_tab;
+	if (desc_idx >= tab->nr_descs)
+		return -EFAULT;
+	desc = &tab->descs[desc_idx];
+	if (desc->func_id != func_id)
 		return -EFAULT;
 
 	*func_addr = (u8 *)desc->addr;
@@ -21303,7 +21307,8 @@ static int specialize_kfunc(struct bpf_verifier_env *env, struct bpf_kfunc_desc
 }
 
 static int add_kfunc_desc_target(struct bpf_verifier_env *env,
-				 const struct bpf_kfunc_desc *target_desc)
+				 const struct bpf_kfunc_desc *target_desc,
+				 u16 *desc_idx)
 {
 	struct bpf_kfunc_desc desc = *target_desc;
 	struct bpf_kfunc_desc_tab *new_tab;
@@ -21316,8 +21321,10 @@ static int add_kfunc_desc_target(struct bpf_verifier_env *env,
 	for (i = 0; i < tab->nr_descs; i++) {
 		if (tab->descs[i].func_id == desc.func_id &&
 		    tab->descs[i].offset == desc.offset &&
-		    tab->descs[i].addr == desc.addr)
+		    tab->descs[i].addr == desc.addr) {
+			*desc_idx = i;
 			return 0;
+		}
 	}
 
 	if (tab->nr_descs == MAX_KFUNC_CALL_DESCS) {
@@ -21332,6 +21339,7 @@ static int add_kfunc_desc_target(struct bpf_verifier_env *env,
 	tab = new_tab;
 	prog_aux->kfunc_tab = tab;
 
+	*desc_idx = tab->nr_descs;
 	tab->descs[tab->nr_descs++] = desc;
 	return 0;
 }
@@ -21359,6 +21367,7 @@ int bpf_fixup_kfunc_call(struct bpf_verifier_env *env, struct bpf_insn *insn,
 	struct bpf_kfunc_desc desc_copy;
 	struct bpf_kfunc_desc *desc;
 	unsigned long call_imm;
+	u16 desc_idx;
 	bool near_call;
 	int err;
 
@@ -21381,10 +21390,8 @@ int bpf_fixup_kfunc_call(struct bpf_verifier_env *env, struct bpf_insn *insn,
 	}
 
 	near_call = !bpf_jit_supports_far_kfunc_call();
-	if (near_call) {
-		desc_copy = *desc;
-		desc = &desc_copy;
-	}
+	desc_copy = *desc;
+	desc = &desc_copy;
 
 	err = specialize_kfunc(env, desc, insn_idx);
 	if (err)
@@ -21398,12 +21405,14 @@ int bpf_fixup_kfunc_call(struct bpf_verifier_env *env, struct bpf_insn *insn,
 			return -EINVAL;
 		}
 		insn->imm = call_imm;
-
-		err = add_kfunc_desc_target(env, desc);
-		if (err)
-			return err;
 	}
 
+	err = add_kfunc_desc_target(env, desc, &desc_idx);
+	if (err)
+		return err;
+	if (!near_call)
+		insn->off = desc_idx;
+
 	if (is_bpf_obj_new_kfunc(desc->func_id) || is_bpf_percpu_obj_new_kfunc(desc->func_id)) {
 		struct btf_struct_meta *kptr_struct_meta = env->insn_aux_data[insn_idx].kptr_struct_meta;
 		struct bpf_insn addr[2] = { BPF_LD_IMM64(BPF_REG_2, (long)kptr_struct_meta) };
-- 
2.52.0