Re: [PATCH bpf-next v4 2/3] bpf: struct_ops supports more than one page for trampolines.

[Martin KaFai Lau <martin.lau@linux.dev>]

On 2/24/24 2:34 PM, Kui-Feng Lee wrote:
> +void bpf_struct_ops_tramp_buf_free(void *image)
> +{
> +	arch_free_bpf_trampoline(image, PAGE_SIZE);
> +	if (image)

Moved the "arch_free_bpf_trampoline(image, PAGE_SIZE);" after the image NULL 
test. I think it was an overlook at the bpf_dummy_struct_ops.c. The exisiting 
__bpf_struct_ops_map_free(image, PAGE_SIZE) had been testing NULL before free also.

> +		bpf_jit_uncharge_modmem(PAGE_SIZE);
> +}
> +
>   #define MAYBE_NULL_SUFFIX "__nullable"
>   #define MAX_STUB_NAME 128
>   
> @@ -461,6 +486,15 @@ static void bpf_struct_ops_map_put_progs(struct bpf_struct_ops_map *st_map)
>   	}
>   }
>   
> +static void bpf_struct_ops_map_free_image(struct bpf_struct_ops_map *st_map)
> +{
> +	int i;
> +
> +	for (i = 0; i < st_map->image_pages_cnt; i++)
> +		bpf_struct_ops_tramp_buf_free(st_map->image_pages[i]);
> +	st_map->image_pages_cnt = 0;
> +}
> +
>   static int check_zero_holes(const struct btf *btf, const struct btf_type *t, void *data)
>   {
>   	const struct btf_member *member;
> @@ -503,12 +537,21 @@ const struct bpf_link_ops bpf_struct_ops_link_lops = {
>   	.dealloc = bpf_struct_ops_link_dealloc,
>   };
>   
> +/* *image should be NULL and allow_alloc should be true if a caller wants
> + * this function to allocate a image buffer for it. Otherwise, this
> + * function allocate a new image buffer only if allow_alloc is true and the
> + * size of the trampoline is larger than the space left in the current
> + * image buffer.
> + */
>   int bpf_struct_ops_prepare_trampoline(struct bpf_tramp_links *tlinks,
>   				      struct bpf_tramp_link *link,
>   				      const struct btf_func_model *model,
> -				      void *stub_func, void *image, void *image_end)
> +				      void *stub_func,
> +				      void **_image, u32 *image_off,
> +				      bool allow_alloc)
>   {
>   	u32 flags = BPF_TRAMP_F_INDIRECT;
> +	void *image = *_image;
>   	int size;
>   
>   	tlinks[BPF_TRAMP_FENTRY].links[0] = link;
> @@ -518,14 +561,30 @@ int bpf_struct_ops_prepare_trampoline(struct bpf_tramp_links *tlinks,
>   		flags |= BPF_TRAMP_F_RET_FENTRY_RET;
>   
>   	size = arch_bpf_trampoline_size(model, flags, tlinks, NULL);
> -	if (size < 0)
> +	if (size <= 0)
>   		return size;
> -	if (size > (unsigned long)image_end - (unsigned long)image)
> -		return -E2BIG;
> -	return arch_prepare_bpf_trampoline(NULL, image, image_end,
> +
> +	/* Allocate image buffer if necessary */
> +	if (!image || size > PAGE_SIZE - *image_off) {
> +		if (!allow_alloc)
> +			return -E2BIG;
> +
> +		image = bpf_struct_ops_tramp_buf_alloc();
> +		if (IS_ERR(image))
> +			return PTR_ERR(image);
> +		*_image = image;
> +		*image_off = 0;
> +	}
> +
> +	size = arch_prepare_bpf_trampoline(NULL, image + *image_off,
> +					   image + PAGE_SIZE,
>   					   model, flags, tlinks, stub_func);
> -}
> +	if (size > 0)
> +		*image_off += size;
> +	/* The caller should free the allocated memory even if size < 0 */

I massage the logic a little such that the caller does not need to free the new 
unused page when this function errored out. I kept both the "*_image" and 
"*image_off" param unchanged on the error case.

Applied. Thanks.

[ ... ]

> @@ -781,10 +850,7 @@ static void __bpf_struct_ops_map_free(struct bpf_map *map)
>   	if (st_map->links)
>   		bpf_struct_ops_map_put_progs(st_map);
>   	bpf_map_area_free(st_map->links);
> -	if (st_map->image) {
> -		arch_free_bpf_trampoline(st_map->image, PAGE_SIZE);
> -		bpf_jit_uncharge_modmem(PAGE_SIZE);
> -	}
> +	bpf_struct_ops_map_free_image(st_map);
>   	bpf_map_area_free(st_map->uvalue);
>   	bpf_map_area_free(st_map);
>   }

[ ... ]

> diff --git a/net/bpf/bpf_dummy_struct_ops.c b/net/bpf/bpf_dummy_struct_ops.c
> index 02de71719aed..da73905eff4a 100644
> --- a/net/bpf/bpf_dummy_struct_ops.c
> +++ b/net/bpf/bpf_dummy_struct_ops.c
> @@ -91,6 +91,7 @@ int bpf_struct_ops_test_run(struct bpf_prog *prog, const union bpf_attr *kattr,
>   	struct bpf_tramp_link *link = NULL;
>   	void *image = NULL;
>   	unsigned int op_idx;
> +	u32 image_off = 0;
>   	int prog_ret;
>   	s32 type_id;
>   	int err;
> @@ -114,12 +115,6 @@ int bpf_struct_ops_test_run(struct bpf_prog *prog, const union bpf_attr *kattr,
>   		goto out;
>   	}
>   
> -	image = arch_alloc_bpf_trampoline(PAGE_SIZE);
> -	if (!image) {
> -		err = -ENOMEM;
> -		goto out;
> -	}
> -
>   	link = kzalloc(sizeof(*link), GFP_USER);
>   	if (!link) {
>   		err = -ENOMEM;
> @@ -133,7 +128,8 @@ int bpf_struct_ops_test_run(struct bpf_prog *prog, const union bpf_attr *kattr,
>   	err = bpf_struct_ops_prepare_trampoline(tlinks, link,
>   						&st_ops->func_models[op_idx],
>   						&dummy_ops_test_ret_function,
> -						image, image + PAGE_SIZE);
> +						&image, &image_off,
> +						true);
>   	if (err < 0)
>   		goto out;
>   
> @@ -147,7 +143,7 @@ int bpf_struct_ops_test_run(struct bpf_prog *prog, const union bpf_attr *kattr,
>   		err = -EFAULT;
>   out:
>   	kfree(args);
> -	arch_free_bpf_trampoline(image, PAGE_SIZE);
> +	bpf_struct_ops_tramp_buf_free(image);
>   	if (link)
>   		bpf_link_put(&link->link);
>   	kfree(tlinks);