Re: [PATCH bpf-next v3 2/9] bpf: Support __arena and __arena_nullable on struct_ops arguments

[Eduard Zingerman <eddyz87@gmail.com>]

[-- Attachment #1: Type: text/plain, Size: 3822 bytes --]

On Tue, 2026-08-04 at 16:55 -0700, Eduard Zingerman wrote:
> On Mon, 2026-08-03 at 14:51 +0200, Kumar Kartikeya Dwivedi wrote:
> 
> Lgtm except for several nits below.
> 
> ...
> 
> > diff --git a/kernel/bpf/bpf_struct_ops.c b/kernel/bpf/bpf_struct_ops.c
> 
> ...
> 
> > @@ -226,26 +229,30 @@ static int prepare_arg_info(struct btf *btf,
> >  	info = info_buf;
> >  	for (arg_no = 0; arg_no < nargs; arg_no++) {
> >  		/* Skip arguments that is not suffixed with
> > -		 * "__nullable or __ref".
> > +		 * "__nullable", "__ref", "__arena" or "__arena_nullable".
> >  		 */
> >  		is_nullable = btf_param_match_suffix(btf, &stub_args[arg_no],
> >  						     MAYBE_NULL_SUFFIX);
> >  		is_refcounted = btf_param_match_suffix(btf, &stub_args[arg_no],
> >  						       REFCOUNTED_SUFFIX);
> > +		is_arena_nullable = btf_param_match_suffix(btf, &stub_args[arg_no],
> > +							   ARENA_MAYBE_NULL_SUFFIX);
> > +		is_arena = btf_param_match_suffix(btf, &stub_args[arg_no], ARENA_SUFFIX);
> >  
> >  		if (is_nullable)
> >  			suffix = MAYBE_NULL_SUFFIX;
> >  		else if (is_refcounted)
> >  			suffix = REFCOUNTED_SUFFIX;
> > +		else if (is_arena_nullable)
> > +			suffix = ARENA_MAYBE_NULL_SUFFIX;
> > +		else if (is_arena)
> > +			suffix = ARENA_SUFFIX;
> 
> As discussed already, that's unfortunate that p__arena and p__nullable
> are valid, but p__arena__nullable is not, substituted by
> p__arena_nullable. Same goes for global subprogs.
> 
> >  		else
> >  			continue;
> >  
> > -		/* Should be a pointer to struct */
> > -		pointed_type = btf_type_resolve_ptr(btf,
> > -						    args[arg_no].type,
> > -						    &arg_btf_id);
> > -		if (!pointed_type ||
> > -		    !btf_type_is_struct(pointed_type)) {
> > +		/* Should be a pointer to struct, or any pointer for __arena/__arena_nullable */
> > +		pointed_type = btf_type_resolve_ptr(btf, args[arg_no].type, &arg_btf_id);
> > +		if (!pointed_type || (!is_arena && !is_arena_nullable && !btf_type_is_struct(pointed_type))) {
> >  			pr_warn("stub function %s has %s tagging to an unsupported type\n",
> >  				stub_fname, suffix);
> >  			goto err_out;
> 
> ...
> 
> > diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
> > index ed7999ad6c66..d2d7a2904345 100644
> > --- a/kernel/bpf/trampoline.c
> > +++ b/kernel/bpf/trampoline.c
> 
> ...
> 
> > @@ -695,6 +743,22 @@ static int bpf_trampoline_update(struct bpf_trampoline *tr, bool lock_direct_mut
> >  		goto out;
> >  	}
> >  
> > +	/*
> > +	 * Arena ctx args are converted only by the struct_ops indirect
> > +	 * trampoline, which dispatches to a single known prog. Generic
> > +	 * trampolines can mix progs with different arenas, so no conversion
> > +	 * is possible here. Not reachable today: only struct_ops progs get
> > +	 * arena ctx args and they never ride generic trampolines.
> > +	 */
> > +	for (kind = 0; kind < BPF_TRAMP_MAX; kind++) {
> > +		for (i = 0; i < tnodes[kind].nr_nodes; i++) {
> > +			if (bpf_prog_has_arena_ctx_arg(tnodes[kind].nodes[i]->link->prog)) {
> > +				err = -ENOTSUPP;
> > +				goto out;
> > +			}
> > +		}
> > +	}
> 
> Nit: should this be done in __bpf_trampoline_link_prog() or
>      bpf_trampoline_add_prog()? To avoid doing the same check on unlink.
>      Since this is an invariant violation, do we want to add WARN_ONCE?
> 
> > +
> >  	/* clear all bits except SHARE_IPMODIFY and TAIL_CALL_CTX */
> >  	tr->flags &= (BPF_TRAMP_F_SHARE_IPMODIFY | BPF_TRAMP_F_TAIL_CALL_CTX);
> >  
> 
> ...

Thinking a bit more about this. Since kfuncs already use flags in
btf_func_model, I think it would be nice to use them for struct_ops as
well. E.g. as in the patch attached (applied on top of this series).
Wdyt?

[-- Attachment #2: struct-ops-flags.diff --]
[-- Type: text/x-patch, Size: 9423 bytes --]

diff --git a/arch/x86/net/bpf_jit_comp.c b/arch/x86/net/bpf_jit_comp.c
index 4b349ae35ebf..dcc68eb5632a 100644
--- a/arch/x86/net/bpf_jit_comp.c
+++ b/arch/x86/net/bpf_jit_comp.c
@@ -3071,12 +3071,12 @@ static void emit_arena_arg_conv(u8 **pprog, u32 src_reg, bool nullable, u32 base
 
 static void save_args(const struct btf_func_model *m, u8 **prog,
 		      int stack_size, bool for_call_origin, u32 flags,
-		      const struct bpf_tramp_arena_args *aargs)
+		      u64 arena_base)
 {
 	int arg_regs, first_off = 0, nr_regs = 0, nr_stack_slots = 0;
 	bool use_jmp = bpf_trampoline_use_jmp(flags);
 	int stack_args_off = (use_jmp || (flags & BPF_TRAMP_F_INDIRECT)) ? 16 : 24;
-	int i, j, slot = 0;
+	int i, j;
 
 	/* Store function arguments to stack.
 	 * For a function that accepts two pointers the sequence will be:
@@ -3084,6 +3084,9 @@ static void save_args(const struct btf_func_model *m, u8 **prog,
 	 * mov QWORD PTR [rbp-0x8],rsi
 	 */
 	for (i = 0; i < min_t(int, m->nr_args, MAX_BPF_FUNC_ARGS); i++) {
+		bool arena_arg = arena_base && (m->arg_flags[i] & BTF_FMODEL_ARENA_ARG);
+		bool nullable = m->arg_flags[i] & BTF_FMODEL_NULLABLE_ARG;
+
 		arg_regs = (m->arg_size[i] + 7) / 8;
 
 		/* According to the research of Yonghong, struct members
@@ -3117,10 +3120,9 @@ static void save_args(const struct btf_func_model *m, u8 **prog,
 			for (j = 0; j < arg_regs; j++) {
 				emit_ldx(prog, BPF_DW, BPF_REG_0, BPF_REG_FP,
 					 nr_stack_slots * 8 + stack_args_off);
-				if (aargs && (aargs->slots & BIT(slot)))
-					emit_arena_arg_conv(prog, BPF_REG_0,
-							    aargs->nullable_slots & BIT(slot),
-							    (u32)aargs->kern_vm_start);
+				if (arena_arg)
+					emit_arena_arg_conv(prog, BPF_REG_0, nullable,
+							    (u32)arena_base);
 				emit_stx(prog, BPF_DW, BPF_REG_FP, BPF_REG_0,
 					 -stack_size);
 
@@ -3128,7 +3130,6 @@ static void save_args(const struct btf_func_model *m, u8 **prog,
 					first_off = stack_size;
 				stack_size -= 8;
 				nr_stack_slots++;
-				slot++;
 			}
 		} else {
 			/* Only copy the arguments on-stack to current
@@ -3137,7 +3138,6 @@ static void save_args(const struct btf_func_model *m, u8 **prog,
 			 */
 			if (for_call_origin) {
 				nr_regs += arg_regs;
-				slot += arg_regs;
 				continue;
 			}
 
@@ -3145,16 +3145,13 @@ static void save_args(const struct btf_func_model *m, u8 **prog,
 			for (j = 0; j < arg_regs; j++) {
 				u32 src = nr_regs == 5 ? X86_REG_R9 : BPF_REG_1 + nr_regs;
 
-				if (aargs && (aargs->slots & BIT(slot))) {
-					emit_arena_arg_conv(prog, src,
-							    aargs->nullable_slots & BIT(slot),
-							    (u32)aargs->kern_vm_start);
+				if (arena_arg) {
+					emit_arena_arg_conv(prog, src, nullable, (u32)arena_base);
 					src = BPF_REG_0;
 				}
 				emit_stx(prog, BPF_DW, BPF_REG_FP, src, -stack_size);
 				stack_size -= 8;
 				nr_regs++;
-				slot++;
 			}
 		}
 	}
@@ -3445,13 +3442,12 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 	struct bpf_tramp_nodes *fentry = &tnodes[BPF_TRAMP_FENTRY];
 	struct bpf_tramp_nodes *fexit = &tnodes[BPF_TRAMP_FEXIT];
 	struct bpf_tramp_nodes *fmod_ret = &tnodes[BPF_TRAMP_MODIFY_RETURN];
-	struct bpf_tramp_arena_args aargs;
 	void *orig_call = func_addr;
 	int cookie_off, cookie_cnt;
 	u8 **branches = NULL;
+	u64 arena_base;
 	u64 func_meta;
 	u8 *prog;
-	bool has_aargs;
 	bool save_ret;
 
 	/*
@@ -3462,7 +3458,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 	WARN_ON_ONCE((flags & BPF_TRAMP_F_INDIRECT) &&
 		     (flags & ~(BPF_TRAMP_F_INDIRECT | BPF_TRAMP_F_RET_FENTRY_RET)));
 
-	has_aargs = bpf_tramp_collect_arena_args(tnodes, flags, &aargs);
+	arena_base = bpf_tramp_arena_base(m, tnodes, flags);
 
 	/* extra registers for struct arguments */
 	for (i = 0; i < m->nr_args; i++) {
@@ -3601,8 +3597,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 		emit_store_stack_imm64(&prog, BPF_REG_0, -ip_off, (long)func_addr);
 	}
 
-	save_args(m, &prog, regs_off, false, flags,
-		  has_aargs ? &aargs : NULL);
+	save_args(m, &prog, regs_off, false, flags, arena_base);
 
 	if (flags & BPF_TRAMP_F_CALL_ORIG) {
 		/* arg1: mov rdi, im */
@@ -3644,7 +3639,7 @@ static int __arch_prepare_bpf_trampoline(struct bpf_tramp_image *im, void *rw_im
 
 	if (flags & BPF_TRAMP_F_CALL_ORIG) {
 		restore_regs(m, &prog, regs_off);
-		save_args(m, &prog, arg_stack_off, true, flags, NULL);
+		save_args(m, &prog, arg_stack_off, true, flags, 0);
 
 		if (flags & BPF_TRAMP_F_TAIL_CALL_CTX) {
 			/* Before calling the original function, load the
diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index c839f039729a..a959efc3468f 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h
@@ -1293,18 +1293,13 @@ struct bpf_tramp_nodes {
 };
 
 /*
- * Which 8-byte ctx slots of a struct_ops trampoline hold arena kernel
- * pointers that save_args() converts to the arena pointer form,
- * ctx[slot] = (u32)(kaddr - kern_vm_start).
+ * The arena base against which a struct_ops trampoline converts the
+ * arguments marked with BTF_FMODEL_ARENA_ARG while saving them into the BPF
+ * ctx, ctx[arg] = (u32)(kaddr - kern_vm_start). Zero when the trampoline
+ * converts nothing.
  */
-struct bpf_tramp_arena_args {
-	u32 slots;
-	u32 nullable_slots;	/* subset of @slots where NULL is preserved */
-	u64 kern_vm_start;
-};
-
-bool bpf_tramp_collect_arena_args(struct bpf_tramp_nodes *tnodes, u32 flags,
-				  struct bpf_tramp_arena_args *aargs);
+u64 bpf_tramp_arena_base(const struct btf_func_model *m,
+			 struct bpf_tramp_nodes *tnodes, u32 flags);
 
 struct bpf_tramp_run_ctx;
 
@@ -1707,11 +1702,6 @@ struct bpf_ctx_arg_aux {
 	u32 btf_id;
 	u32 ref_id;
 	bool refcounted;
-	/*
-	 * We don't encode NULL-ness in the type for the program, but still need
-	 * to distinguish it for the purposes of telling JITs what sequence to emit.
-	 */
-	bool arena_nullable;
 };
 
 struct btf_mod_pair {
diff --git a/kernel/bpf/bpf_struct_ops.c b/kernel/bpf/bpf_struct_ops.c
index 827a6216a620..c6a7c0213bb0 100644
--- a/kernel/bpf/bpf_struct_ops.c
+++ b/kernel/bpf/bpf_struct_ops.c
@@ -174,6 +174,7 @@ static int prepare_arg_info(struct btf *btf,
 			    const char *st_ops_name,
 			    const char *member_name,
 			    const struct btf_type *func_proto, void *stub_func_addr,
+			    struct btf_func_model *model,
 			    struct bpf_struct_ops_arg_info *arg_info)
 {
 	const struct btf_type *stub_func_proto, *pointed_type;
@@ -288,7 +289,9 @@ static int prepare_arg_info(struct btf *btf,
 			 * precision around it, since it has no safety implication.
 			 */
 			info->reg_type = PTR_TO_ARENA;
-			info->arena_nullable = is_arena_nullable;
+			model->arg_flags[arg_no] |= BTF_FMODEL_ARENA_ARG;
+			if (is_arena_nullable)
+				model->arg_flags[arg_no] |= BTF_FMODEL_NULLABLE_ARG;
 		}
 
 		info++;
@@ -476,6 +479,7 @@ int bpf_struct_ops_desc_init(struct bpf_struct_ops_desc *st_ops_desc,
 		stub_func_addr = *(void **)(st_ops->cfi_stubs + moff);
 		err = prepare_arg_info(btf, st_ops->name, mname,
 				       func_proto, stub_func_addr,
+				       &st_ops->func_models[i],
 				       arg_info + i);
 		if (err)
 			goto errout;
diff --git a/kernel/bpf/trampoline.c b/kernel/bpf/trampoline.c
index d2d7a2904345..412a7d32e94a 100644
--- a/kernel/bpf/trampoline.c
+++ b/kernel/bpf/trampoline.c
@@ -540,40 +540,33 @@ static bool bpf_prog_has_arena_ctx_arg(const struct bpf_prog *prog)
 }
 
 /*
- * Collect which ctx slots of a struct_ops trampoline hold arena kernel
- * pointers that save_args() must convert to the arena pointer form. Only
- * the struct_ops indirect trampoline converts: it dispatches to a single
- * prog whose arena is known at generation time. Return false when there
- * is nothing to convert.
+ * The arena base against which save_args() converts the arguments marked
+ * with BTF_FMODEL_ARENA_ARG. Only the struct_ops indirect trampoline
+ * converts: it dispatches to a single prog whose arena is known at
+ * generation time. Return 0 when there is nothing to convert.
  */
-bool bpf_tramp_collect_arena_args(struct bpf_tramp_nodes *tnodes, u32 flags,
-				  struct bpf_tramp_arena_args *aargs)
+u64 bpf_tramp_arena_base(const struct btf_func_model *m,
+			 struct bpf_tramp_nodes *tnodes, u32 flags)
 {
 	const struct bpf_prog *prog;
 	int i;
 
-	memset(aargs, 0, sizeof(*aargs));
-
 	if (!(flags & BPF_TRAMP_F_INDIRECT) ||
 	    tnodes[BPF_TRAMP_FENTRY].nr_nodes != 1)
-		return false;
+		return 0;
 
-	prog = tnodes[BPF_TRAMP_FENTRY].nodes[0]->link->prog;
-	for (i = 0; i < prog->aux->ctx_arg_info_size; i++) {
-		const struct bpf_ctx_arg_aux *info = &prog->aux->ctx_arg_info[i];
+	for (i = 0; i < m->nr_args; i++)
+		if (m->arg_flags[i] & BTF_FMODEL_ARENA_ARG)
+			break;
+	if (i == m->nr_args)
+		return 0;
 
-		if (base_type(info->reg_type) != PTR_TO_ARENA)
-			continue;
-		aargs->slots |= BIT(info->offset / 8);
-		if (info->arena_nullable)
-			aargs->nullable_slots |= BIT(info->offset / 8);
-	}
-	if (!aargs->slots)
-		return false;
+	/* Verification rejects an arena argument without an arena. */
+	prog = tnodes[BPF_TRAMP_FENTRY].nodes[0]->link->prog;
 	if (WARN_ON_ONCE(!prog->aux->arena))
-		return false;
-	aargs->kern_vm_start = bpf_arena_get_kern_vm_start(prog->aux->arena);
-	return true;
+		return 0;
+
+	return bpf_arena_get_kern_vm_start(prog->aux->arena);
 }
 
 static void bpf_tramp_image_free(struct bpf_tramp_image *im)