[PATCH bpf-next v2] bpf: Add __prog tag to pass in prog->aux

[PATCH bpf-next v2] bpf: Add __prog tag to pass in prog->aux

[Kumar Kartikeya Dwivedi]

Instead of hardcoding the list of kfuncs that need prog->aux passed to
them with a combination of fixup_kfunc_call adjustment + __ign suffix,
combine both in __prog suffix, which ignores the argument passed in, and
fixes it up to the prog->aux. This allows kfuncs to have the prog->aux
passed into them without having to touch the verifier.

Cc: Tejun Heo <tj@kernel.org>
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
Changelog:
v1 -> v2
v1: https://lore.kernel.org/bpf/20250512210246.3741193-1-memxor@gmail.com

 * Change __aux tag to __prog. (Alexei)
---
 include/linux/bpf_verifier.h |  1 +
 kernel/bpf/helpers.c         |  4 ++--
 kernel/bpf/verifier.c        | 33 +++++++++++++++++++++++++++------
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/include/linux/bpf_verifier.h b/include/linux/bpf_verifier.h
index 9734544b6957..7dd85ed6059e 100644
--- a/include/linux/bpf_verifier.h
+++ b/include/linux/bpf_verifier.h
@@ -606,6 +606,7 @@ struct bpf_insn_aux_data {
 	bool calls_callback;
 	/* registers alive before this instruction. */
 	u16 live_regs_before;
+	u16 arg_prog;
 };

 #define MAX_USED_MAPS 64 /* max number of maps accessed by one eBPF program */
diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index fed53da75025..43cbf439b9fb 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -3012,9 +3012,9 @@ __bpf_kfunc int bpf_wq_start(struct bpf_wq *wq, unsigned int flags)
 __bpf_kfunc int bpf_wq_set_callback_impl(struct bpf_wq *wq,
 					 int (callback_fn)(void *map, int *key, void *value),
 					 unsigned int flags,
-					 void *aux__ign)
+					 void *aux__prog)
 {
-	struct bpf_prog_aux *aux = (struct bpf_prog_aux *)aux__ign;
+	struct bpf_prog_aux *aux = (struct bpf_prog_aux *)aux__prog;
 	struct bpf_async_kern *async = (struct bpf_async_kern *)wq;

 	if (flags)
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 28f5a7899bd6..f409a06099f6 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -343,6 +343,7 @@ struct bpf_kfunc_call_arg_meta {
 		int uid;
 	} map;
 	u64 mem_size;
+	u32 arg_prog;
 };

 struct btf *btf_vmlinux;
@@ -11897,6 +11898,11 @@ static bool is_kfunc_arg_irq_flag(const struct btf *btf, const struct btf_param
 	return btf_param_match_suffix(btf, arg, "__irq_flag");
 }

+static bool is_kfunc_arg_prog(const struct btf *btf, const struct btf_param *arg)
+{
+	return btf_param_match_suffix(btf, arg, "__prog");
+}
+
 static bool is_kfunc_arg_scalar_with_name(const struct btf *btf,
 					  const struct btf_param *arg,
 					  const char *name)
@@ -12938,6 +12944,17 @@ static int check_kfunc_args(struct bpf_verifier_env *env, struct bpf_kfunc_call_
 		if (is_kfunc_arg_ignore(btf, &args[i]))
 			continue;

+		if (is_kfunc_arg_prog(btf, &args[i])) {
+			/* Used to reject repeated use of __aux. */
+			if (meta->arg_prog) {
+				verbose(env, "Only 1 prog->aux argument supported per-kfunc\n");
+				return -EFAULT;
+			}
+			meta->arg_prog = regno;
+			cur_aux(env)->arg_prog = regno;
+			continue;
+		}
+
 		if (btf_type_is_scalar(t)) {
 			if (reg->type != SCALAR_VALUE) {
 				verbose(env, "R%d is not a scalar\n", regno);
@@ -21517,13 +21534,17 @@ static int fixup_kfunc_call(struct bpf_verifier_env *env, struct bpf_insn *insn,
 		   desc->func_id == special_kfunc_list[KF_bpf_rdonly_cast]) {
 		insn_buf[0] = BPF_MOV64_REG(BPF_REG_0, BPF_REG_1);
 		*cnt = 1;
-	} else if (is_bpf_wq_set_callback_impl_kfunc(desc->func_id)) {
-		struct bpf_insn ld_addrs[2] = { BPF_LD_IMM64(BPF_REG_4, (long)env->prog->aux) };
+	}

-		insn_buf[0] = ld_addrs[0];
-		insn_buf[1] = ld_addrs[1];
-		insn_buf[2] = *insn;
-		*cnt = 3;
+	if (env->insn_aux_data[insn_idx].arg_prog) {
+		u32 regno = env->insn_aux_data[insn_idx].arg_prog;
+		struct bpf_insn ld_addrs[2] = { BPF_LD_IMM64(regno, (long)env->prog->aux) };
+		int idx = *cnt;
+
+		insn_buf[idx++] = ld_addrs[0];
+		insn_buf[idx++] = ld_addrs[1];
+		insn_buf[idx++] = *insn;
+		*cnt = idx;
 	}
 	return 0;
 }
--
2.47.1

Re: [PATCH bpf-next v2] bpf: Add __prog tag to pass in prog->aux

[Eduard Zingerman]

On Mon, 2025-05-12 at 19:57 -0700, Kumar Kartikeya Dwivedi wrote:

[...]

In Documentation/bpf/kfuncs.rst there are descriptions for several
suffixes that are supported, "__prog" should be added there.

Do we want to add a separate test case for this feature?
It looks like there are no tests for other suffixes,
so relying on wq tests passing is probably fine.

Implementation lgtm, a few nits below.

>  include/linux/bpf_verifier.h |  1 +
>  kernel/bpf/helpers.c         |  4 ++--
>  kernel/bpf/verifier.c        | 33 +++++++++++++++++++++++++++------
>  3 files changed, 30 insertions(+), 8 deletions(-)
> 
> diff --git a/include/linux/bpf_verifier.h b/include/linux/bpf_verifier.h
> index 9734544b6957..7dd85ed6059e 100644
> --- a/include/linux/bpf_verifier.h
> +++ b/include/linux/bpf_verifier.h
> @@ -606,6 +606,7 @@ struct bpf_insn_aux_data {
>  	bool calls_callback;
>  	/* registers alive before this instruction. */
>  	u16 live_regs_before;
> +	u16 arg_prog;

Nit: there is a 4-bit hole after `fastcall_spills_num`,
     `arg_prog` field could be put there and 2 bytes at the tail
     would be remain for future extension.

>  };
> 
>  #define MAX_USED_MAPS 64 /* max number of maps accessed by one eBPF program */
> diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
> index fed53da75025..43cbf439b9fb 100644
> --- a/kernel/bpf/helpers.c
> +++ b/kernel/bpf/helpers.c
> @@ -3012,9 +3012,9 @@ __bpf_kfunc int bpf_wq_start(struct bpf_wq *wq, unsigned int flags)
>  __bpf_kfunc int bpf_wq_set_callback_impl(struct bpf_wq *wq,
>  					 int (callback_fn)(void *map, int *key, void *value),
>  					 unsigned int flags,
> -					 void *aux__ign)
> +					 void *aux__prog)
>  {
> -	struct bpf_prog_aux *aux = (struct bpf_prog_aux *)aux__ign;
> +	struct bpf_prog_aux *aux = (struct bpf_prog_aux *)aux__prog;
>  	struct bpf_async_kern *async = (struct bpf_async_kern *)wq;
> 
>  	if (flags)
> diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
> index 28f5a7899bd6..f409a06099f6 100644
> --- a/kernel/bpf/verifier.c
> +++ b/kernel/bpf/verifier.c
> @@ -343,6 +343,7 @@ struct bpf_kfunc_call_arg_meta {
>  		int uid;
>  	} map;
>  	u64 mem_size;
> +	u32 arg_prog;

Nit: this is a boolean flag, I'd put it after `bool arg_owning_ref`,
     as there is a 3 bytes hole there.

>  };
> 
>  struct btf *btf_vmlinux;

[...]

@@ -12906,6 +12912,17 @@ static int check_kfunc_args(struct bpf_verifier_env *env, struct bpf_kfunc_call_
 		if (is_kfunc_arg_ignore(btf, &args[i]))
 			continue;
 
+		if (is_kfunc_arg_prog(btf, &args[i])) {
+			/* Used to reject repeated use of __aux. */
                                                          ^^^^^^
                                               Nit: should be __prog.

+			if (meta->arg_prog) {
+				verbose(env, "Only 1 prog->aux argument supported per-kfunc\n");
+				return -EFAULT;
+			}
+			meta->arg_prog = regno;
+			cur_aux(env)->arg_prog = regno;
+			continue;
+		}
+
 		if (btf_type_is_scalar(t)) {
 			if (reg->type != SCALAR_VALUE) {
 				verbose(env, "R%d is not a scalar\n", regno);