[PATCH bpf-next v2 12/23] bpf: Resolve BTF ID of ARG_PTR_TO_BTF_ID in kfunc bpf_func_proto

[Amery Hung <ameryhung@gmail.com>]

check_kfunc_args() walks the kfunc's BTF on every verification of a call
to work out which BTF ID an ARG_PTR_TO_BTF_ID argument expects. Helpers
name theirs in bpf_func_proto::arg_btf_id[], as a pointer to a BTF ID
that resolve_btfids fills in at build time.

The ID of a kfunc argument's referent is already stored in the
immutable BTF records that describe its pointer and modifier chain. Make
arg_btf_id[] point to the BTF field containing the resolved ID. A __map
argument instead uses the existing vmlinux BTF ID pointer. Produce this
metadata alongside the argument classification in get_kfunc_arg_type().

Declare the argument BTF ID pointers const, since the verifier only
reads through them. The BTF object owns their storage and remains
alive while the generated prototype is used, so the pointers remain
valid when the kfunc descriptor array is reallocated or sorted.

Both helper and kfunc callers can now read the expected BTF ID through
the same bpf_func_proto field, which lets check_func_arg() take over
the ARG_PTR_TO_BTF_ID case.

arg_btf_id[] shares a union with arg_size[], so an argument cannot
store both. The scalar-struct memory fallback keeps resolving its byte
size at verification time, as it does today.

No functional change.

Signed-off-by: Amery Hung <ameryhung@gmail.com>
---
 include/linux/bpf.h   | 12 ++++++------
 kernel/bpf/verifier.c | 42 ++++++++++++++++++++++++++++++------------
 2 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index aa4d3bb5e8cc..d0066d744ceb 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1021,13 +1021,13 @@ struct bpf_func_proto {
 	};
 	union {
 		struct {
-			u32 *arg1_btf_id;
-			u32 *arg2_btf_id;
-			u32 *arg3_btf_id;
-			u32 *arg4_btf_id;
-			u32 *arg5_btf_id;
+			const u32 *arg1_btf_id;
+			const u32 *arg2_btf_id;
+			const u32 *arg3_btf_id;
+			const u32 *arg4_btf_id;
+			const u32 *arg5_btf_id;
 		};
-		u32 *arg_btf_id[MAX_BPF_FUNC_ARGS];
+		const u32 *arg_btf_id[MAX_BPF_FUNC_ARGS];
 		struct {
 			size_t arg1_size;
 			size_t arg2_size;
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 39f7f633b4d6..0547fbeeebbe 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -8839,7 +8839,7 @@ static int check_func_arg(struct bpf_verifier_env *env, u32 arg,
 	enum bpf_arg_type arg_type = fn->arg_type[arg];
 	int regno = reg_from_argno(argno);
 	enum bpf_reg_type type = reg->type;
-	u32 *arg_btf_id = NULL;
+	const u32 *arg_btf_id = NULL;
 	u32 key_size;
 	int err = 0;
 
@@ -12179,13 +12179,18 @@ bool bpf_is_kfunc_pkt_changing(struct bpf_call_arg_meta *meta)
 
 static int
 get_kfunc_arg_type(struct bpf_verifier_env *env, struct bpf_call_arg_meta *meta,
-		   const struct btf_param *args, int arg, int nargs)
+		   const struct btf_param *args, int arg, int nargs,
+		   struct bpf_func_proto *proto)
 {
 	const struct btf_type *t, *ref_t = NULL;
+	const u32 *ref_id_ptr = NULL;
 	argno_t argno = argno_from_arg(arg + 1);
 	const char *ref_tname = NULL;
+	u32 ref_id;
 	int arg_type;
 
+	proto->arg_btf_id[arg] = NULL;
+
 	if (is_kfunc_arg_prog_aux(meta->btf, &args[arg]))
 		return ARG_PTR_TO_PROG_AUX;
 
@@ -12213,7 +12218,11 @@ get_kfunc_arg_type(struct bpf_verifier_env *env, struct bpf_call_arg_meta *meta,
 			reg_arg_name(env, argno), btf_type_str(t));
 		return -EINVAL;
 	}
-	ref_t = btf_type_skip_modifiers(meta->btf, t->type, NULL);
+	/* Keep a pointer to the BTF field containing the resolved referent ID. */
+	ref_id_ptr = &t->type;
+	ref_t = btf_type_skip_modifiers(meta->btf, *ref_id_ptr, &ref_id);
+	while (*ref_id_ptr != ref_id)
+		ref_id_ptr = &btf_type_by_id(meta->btf, *ref_id_ptr)->type;
 	ref_tname = btf_name_by_offset(meta->btf, ref_t->name_off);
 
 	/* In this function, we verify the kfunc's BTF as per the argument type,
@@ -12336,13 +12345,20 @@ get_kfunc_arg_type(struct bpf_verifier_env *env, struct bpf_call_arg_meta *meta,
 	if (is_kfunc_release(meta) && arg == 0)
 		arg_type |= OBJ_RELEASE;
 
-	/*
-	 * A KF_RCU kfunc accepts an RCU-protected pointer where it would
-	 * otherwise demand a referenced or trusted one. Other argument kinds
-	 * have their own provenance requirements and must not inherit MEM_RCU.
-	 */
-	if (base_type(arg_type) == ARG_PTR_TO_BTF_ID && is_kfunc_rcu(meta))
-		arg_type |= MEM_RCU;
+	if (base_type(arg_type) == ARG_PTR_TO_BTF_ID) {
+		if (is_kfunc_arg_map(meta->btf, &args[arg]))
+			proto->arg_btf_id[arg] = reg2btf_ids[CONST_PTR_TO_MAP];
+		else
+			proto->arg_btf_id[arg] = ref_id_ptr;
+
+		/*
+		 * A KF_RCU kfunc accepts an RCU-protected pointer where it would
+		 * otherwise demand a referenced or trusted one. Other argument kinds
+		 * have their own provenance requirements and must not inherit MEM_RCU.
+		 */
+		if (is_kfunc_rcu(meta))
+			arg_type |= MEM_RCU;
+	}
 
 	return arg_type;
 }
@@ -12368,7 +12384,7 @@ static int gen_kfunc_arg_proto(struct bpf_verifier_env *env, struct bpf_call_arg
 	}
 
 	for (i = 0; i < nargs; i++) {
-		arg_type = get_kfunc_arg_type(env, meta, args, i, nargs);
+		arg_type = get_kfunc_arg_type(env, meta, args, i, nargs, proto);
 		if (arg_type < 0)
 			return arg_type;
 
@@ -13024,8 +13040,10 @@ static int check_kfunc_args(struct bpf_verifier_env *env, struct bpf_call_arg_me
 			continue;
 		}
 
+		if (base_type(arg_type) == ARG_PTR_TO_BTF_ID)
+			ref_id = *meta->fn->arg_btf_id[i];
+
 		if (is_kfunc_arg_map(btf, &args[i])) {
-			ref_id = *reg2btf_ids[CONST_PTR_TO_MAP];
 			ref_t = btf_type_by_id(btf_vmlinux, ref_id);
 			ref_tname = btf_name_by_offset(btf, ref_t->name_off);
 		}
-- 
2.52.0