Re: [PATCH bpf-next v6 4/9] bpf: Mark each subprog with proper private stack modes

[Jiri Olsa <olsajiri@gmail.com>]

On Sun, Oct 20, 2024 at 12:14:05PM -0700, Yonghong Song wrote:
> Three private stack modes are used to direct jit action:
>   NO_PRIV_STACK:        do not use private stack
>   PRIV_STACK_SUB_PROG:  adjust frame pointer address (similar to normal stack)
>   PRIV_STACK_ROOT_PROG: set the frame pointer
> 
> Note that for subtree root prog (main prog or callback fn), even if the
> bpf_prog stack size is 0, PRIV_STACK_ROOT_PROG mode is still used.
> This is for bpf exception handling. More details can be found in
> subsequent jit support and selftest patches.
> 
> Signed-off-by: Yonghong Song <yonghong.song@linux.dev>
> ---
>  include/linux/bpf.h   |  9 +++++++++
>  kernel/bpf/core.c     | 19 +++++++++++++++++++
>  kernel/bpf/verifier.c | 29 +++++++++++++++++++++++++++++
>  3 files changed, 57 insertions(+)
> 
> diff --git a/include/linux/bpf.h b/include/linux/bpf.h
> index 376e43fc72b9..27430e9dcfe3 100644
> --- a/include/linux/bpf.h
> +++ b/include/linux/bpf.h
> @@ -1456,6 +1456,12 @@ struct btf_mod_pair {
>  
>  struct bpf_kfunc_desc_tab;
>  
> +enum bpf_priv_stack_mode {
> +	NO_PRIV_STACK,
> +	PRIV_STACK_SUB_PROG,
> +	PRIV_STACK_ROOT_PROG,
> +};
> +
>  struct bpf_prog_aux {
>  	atomic64_t refcnt;
>  	u32 used_map_cnt;
> @@ -1472,6 +1478,9 @@ struct bpf_prog_aux {
>  	u32 ctx_arg_info_size;
>  	u32 max_rdonly_access;
>  	u32 max_rdwr_access;
> +	enum bpf_priv_stack_mode priv_stack_mode;
> +	u16 subtree_stack_depth; /* Subtree stack depth if PRIV_STACK_ROOT_PROG, 0 otherwise */
> +	void __percpu *priv_stack_ptr;
>  	struct btf *attach_btf;
>  	const struct bpf_ctx_arg_aux *ctx_arg_info;
>  	struct mutex dst_mutex; /* protects dst_* pointers below, *after* prog becomes visible */
> diff --git a/kernel/bpf/core.c b/kernel/bpf/core.c
> index 14d9288441f2..aee0055def4f 100644
> --- a/kernel/bpf/core.c
> +++ b/kernel/bpf/core.c
> @@ -1240,6 +1240,7 @@ void __weak bpf_jit_free(struct bpf_prog *fp)
>  		struct bpf_binary_header *hdr = bpf_jit_binary_hdr(fp);
>  
>  		bpf_jit_binary_free(hdr);
> +		free_percpu(fp->aux->priv_stack_ptr);

this should be also put to the x86 version of the bpf_jit_free ?

jirka

>  		WARN_ON_ONCE(!bpf_prog_kallsyms_verify_off(fp));
>  	}
>  
> @@ -2421,6 +2422,24 @@ struct bpf_prog *bpf_prog_select_runtime(struct bpf_prog *fp, int *err)
>  		if (*err)
>  			return fp;
>  
> +		if (fp->aux->priv_stack_eligible) {
> +			if (!fp->aux->stack_depth) {
> +				fp->aux->priv_stack_mode = NO_PRIV_STACK;
> +			} else {
> +				void __percpu *priv_stack_ptr;
> +
> +				fp->aux->priv_stack_mode = PRIV_STACK_ROOT_PROG;
> +				priv_stack_ptr =
> +					__alloc_percpu_gfp(fp->aux->stack_depth, 8, GFP_KERNEL);
> +				if (!priv_stack_ptr) {
> +					*err = -ENOMEM;
> +					return fp;
> +				}
> +				fp->aux->subtree_stack_depth = fp->aux->stack_depth;
> +				fp->aux->priv_stack_ptr = priv_stack_ptr;
> +			}
> +		}
> +
>  		fp = bpf_int_jit_compile(fp);
>  		bpf_prog_jit_attempt_done(fp);
>  		if (!fp->jited && jit_needed) {
> diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
> index 25283ee6f86f..f770015d6ad1 100644
> --- a/kernel/bpf/verifier.c
> +++ b/kernel/bpf/verifier.c
> @@ -20018,6 +20018,8 @@ static int jit_subprogs(struct bpf_verifier_env *env)
>  {
>  	struct bpf_prog *prog = env->prog, **func, *tmp;
>  	int i, j, subprog_start, subprog_end = 0, len, subprog;
> +	int subtree_top_idx, subtree_stack_depth;
> +	void __percpu *priv_stack_ptr;
>  	struct bpf_map *map_ptr;
>  	struct bpf_insn *insn;
>  	void *old_bpf_func;
> @@ -20096,6 +20098,33 @@ static int jit_subprogs(struct bpf_verifier_env *env)
>  		func[i]->is_func = 1;
>  		func[i]->sleepable = prog->sleepable;
>  		func[i]->aux->func_idx = i;
> +
> +		subtree_top_idx = env->subprog_info[i].subtree_top_idx;
> +		if (env->subprog_info[subtree_top_idx].priv_stack_eligible) {
> +			if (subtree_top_idx == i)
> +				func[i]->aux->subtree_stack_depth =
> +					env->subprog_info[i].subtree_stack_depth;
> +
> +			subtree_stack_depth = func[i]->aux->subtree_stack_depth;
> +			if (subtree_top_idx != i) {
> +				if (env->subprog_info[subtree_top_idx].subtree_stack_depth)
> +					func[i]->aux->priv_stack_mode = PRIV_STACK_SUB_PROG;
> +				else
> +					func[i]->aux->priv_stack_mode = NO_PRIV_STACK;
> +			} else if (!subtree_stack_depth) {
> +				func[i]->aux->priv_stack_mode = PRIV_STACK_ROOT_PROG;
> +			} else {
> +				func[i]->aux->priv_stack_mode = PRIV_STACK_ROOT_PROG;
> +				priv_stack_ptr =
> +					__alloc_percpu_gfp(subtree_stack_depth, 8, GFP_KERNEL);
> +				if (!priv_stack_ptr) {
> +					err = -ENOMEM;
> +					goto out_free;
> +				}
> +				func[i]->aux->priv_stack_ptr = priv_stack_ptr;
> +			}
> +		}
> +
>  		/* Below members will be freed only at prog->aux */
>  		func[i]->aux->btf = prog->aux->btf;
>  		func[i]->aux->func_info = prog->aux->func_info;
> -- 
> 2.43.5
> 
>