Re: [PATCH bpf-next v1 02/13] selftests/bpf: Test referenced kptr arguments of struct_ops programs

[Yonghong Song <yonghong.song@linux.dev>]

On 12/13/24 3:29 PM, Amery Hung wrote:
> Test referenced kptr acquired through struct_ops argument tagged with
> "__ref". The success case checks whether 1) a reference to the correct
> type is acquired, and 2) the referenced kptr argument can be accessed in
> multiple paths as long as it hasn't been released. In the fail cases,
> we first confirm that a referenced kptr acquried through a struct_ops
> argument is not allowed to be leaked. Then, we make sure this new
> referenced kptr acquiring mechanism does not accidentally allow referenced
> kptrs to flow into global subprograms through their arguments.
>
> Signed-off-by: Amery Hung <amery.hung@bytedance.com>
> ---
>   .../selftests/bpf/bpf_testmod/bpf_testmod.c   |  7 ++
>   .../selftests/bpf/bpf_testmod/bpf_testmod.h   |  2 +
>   .../prog_tests/test_struct_ops_refcounted.c   | 58 ++++++++++++++++
>   .../bpf/progs/struct_ops_refcounted.c         | 67 +++++++++++++++++++
>   ...ruct_ops_refcounted_fail__global_subprog.c | 32 +++++++++
>   .../struct_ops_refcounted_fail__ref_leak.c    | 17 +++++
>   6 files changed, 183 insertions(+)
>   create mode 100644 tools/testing/selftests/bpf/prog_tests/test_struct_ops_refcounted.c
>   create mode 100644 tools/testing/selftests/bpf/progs/struct_ops_refcounted.c
>   create mode 100644 tools/testing/selftests/bpf/progs/struct_ops_refcounted_fail__global_subprog.c
>   create mode 100644 tools/testing/selftests/bpf/progs/struct_ops_refcounted_fail__ref_leak.c
>
> diff --git a/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c b/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c
> index 987d41af71d2..244234546ae2 100644
> --- a/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c
> +++ b/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c
> @@ -1135,10 +1135,17 @@ static int bpf_testmod_ops__test_maybe_null(int dummy,
>   	return 0;
>   }
>   
> +static int bpf_testmod_ops__test_refcounted(int dummy,
> +					    struct task_struct *task__ref)
> +{
> +	return 0;
> +}
> +
>   static struct bpf_testmod_ops __bpf_testmod_ops = {
>   	.test_1 = bpf_testmod_test_1,
>   	.test_2 = bpf_testmod_test_2,
>   	.test_maybe_null = bpf_testmod_ops__test_maybe_null,
> +	.test_refcounted = bpf_testmod_ops__test_refcounted,
>   };
>   
>   struct bpf_struct_ops bpf_bpf_testmod_ops = {
> diff --git a/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.h b/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.h
> index fb7dff47597a..0e31586c1353 100644
> --- a/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.h
> +++ b/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.h
> @@ -36,6 +36,8 @@ struct bpf_testmod_ops {
>   	/* Used to test nullable arguments. */
>   	int (*test_maybe_null)(int dummy, struct task_struct *task);
>   	int (*unsupported_ops)(void);
> +	/* Used to test ref_acquired arguments. */
> +	int (*test_refcounted)(int dummy, struct task_struct *task);
>   
>   	/* The following fields are used to test shadow copies. */
>   	char onebyte;
> diff --git a/tools/testing/selftests/bpf/prog_tests/test_struct_ops_refcounted.c b/tools/testing/selftests/bpf/prog_tests/test_struct_ops_refcounted.c
> new file mode 100644
> index 000000000000..976df951b700
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/prog_tests/test_struct_ops_refcounted.c
> @@ -0,0 +1,58 @@
> +#include <test_progs.h>
> +
> +#include "struct_ops_refcounted.skel.h"
> +#include "struct_ops_refcounted_fail__ref_leak.skel.h"
> +#include "struct_ops_refcounted_fail__global_subprog.skel.h"
> +
> +/* Test that the verifier accepts a program that first acquires a referenced
> + * kptr through context and then releases the reference
> + */
> +static void refcounted(void)
> +{
> +	struct struct_ops_refcounted *skel;
> +
> +	skel = struct_ops_refcounted__open_and_load();
> +	if (!ASSERT_OK_PTR(skel, "struct_ops_module_open_and_load"))
> +		return;
> +
> +	struct_ops_refcounted__destroy(skel);
> +}
> +
> +/* Test that the verifier rejects a program that acquires a referenced
> + * kptr through context without releasing the reference
> + */
> +static void refcounted_fail__ref_leak(void)
> +{
> +	struct struct_ops_refcounted_fail__ref_leak *skel;
> +
> +	skel = struct_ops_refcounted_fail__ref_leak__open_and_load();
> +	if (ASSERT_ERR_PTR(skel, "struct_ops_module_fail__open_and_load"))
> +		return;
> +
> +	struct_ops_refcounted_fail__ref_leak__destroy(skel);
> +}
> +
> +/* Test that the verifier rejects a program that contains a global
> + * subprogram with referenced kptr arguments
> + */
> +static void refcounted_fail__global_subprog(void)
> +{
> +	struct struct_ops_refcounted_fail__global_subprog *skel;
> +
> +	skel = struct_ops_refcounted_fail__global_subprog__open_and_load();
> +	if (ASSERT_ERR_PTR(skel, "struct_ops_module_fail__open_and_load"))
> +		return;
> +
> +	struct_ops_refcounted_fail__global_subprog__destroy(skel);
> +}
> +
> +void test_struct_ops_refcounted(void)
> +{
> +	if (test__start_subtest("refcounted"))
> +		refcounted();
> +	if (test__start_subtest("refcounted_fail__ref_leak"))
> +		refcounted_fail__ref_leak();
> +	if (test__start_subtest("refcounted_fail__global_subprog"))
> +		refcounted_fail__global_subprog();
> +}
> +
> diff --git a/tools/testing/selftests/bpf/progs/struct_ops_refcounted.c b/tools/testing/selftests/bpf/progs/struct_ops_refcounted.c
> new file mode 100644
> index 000000000000..2c1326668b92
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/progs/struct_ops_refcounted.c
> @@ -0,0 +1,67 @@
> +#include <vmlinux.h>
> +#include <bpf/bpf_tracing.h>
> +#include "../bpf_testmod/bpf_testmod.h"
> +#include "bpf_misc.h"
> +
> +char _license[] SEC("license") = "GPL";
> +
> +extern void bpf_task_release(struct task_struct *p) __ksym;
> +
> +/* This is a test BPF program that uses struct_ops to access a referenced
> + * kptr argument. This is a test for the verifier to ensure that it
> + * 1) recongnizes the task as a referenced object (i.e., ref_obj_id > 0), and
> + * 2) the same reference can be acquired from multiple paths as long as it
> + *    has not been released.
> + *
> + * test_refcounted() is equivalent to the C code below. It is written in assembly
> + * to avoid reads from task (i.e., getting referenced kptrs to task) being merged
> + * into single path by the compiler.
> + *
> + * int test_refcounted(int dummy, struct task_struct *task)
> + * {
> + *         if (dummy % 2)
> + *                 bpf_task_release(task);
> + *         else
> + *                 bpf_task_release(task);
> + *         return 0;
> + * }
> + */
> +SEC("struct_ops/test_refcounted")
> +int test_refcounted(unsigned long long *ctx)
> +{
> +	asm volatile ("					\
> +	/* r6 = dummy */				\
> +	r6 = *(u64 *)(r1 + 0x0);			\
> +	/* if (r6 & 0x1 != 0) */			\
> +	r6 &= 0x1;					\
> +	if r6 == 0 goto l0_%=;				\
> +	/* r1 = task */					\
> +	r1 = *(u64 *)(r1 + 0x8);			\
> +	call %[bpf_task_release];			\
> +	goto l1_%=;					\
> +l0_%=:	/* r1 = task */					\
> +	r1 = *(u64 *)(r1 + 0x8);			\
> +	call %[bpf_task_release];			\
> +l1_%=:	/* return 0 */					\
> +"	:
> +	: __imm(bpf_task_release)
> +	: __clobber_all);
> +	return 0;
> +}

You can use clang nomerge attribute to prevent bpf_task_release(task) merging. For example,

$ cat t.c
struct task_struct {
         int a;
         int b;
         int d[20];
};


__attribute__((nomerge)) extern void bpf_task_release(struct task_struct *task);

int test_refcounted(int dummy, struct task_struct *task)
{
         if (dummy % 2)
                 bpf_task_release(task);
         else
                 bpf_task_release(task);
         return 0;
}

$ clang --version
clang version 19.1.5 (https://github.com/llvm/llvm-project.git ab4b5a2db582958af1ee308a790cfdb42bd24720)
Target: x86_64-unknown-linux-gnu
Thread model: posix
InstalledDir: /home/yhs/work/llvm-project/llvm/build.19/Release/bin
$ clang --target=bpf -O2 -mcpu=v3 -S t.c
$ cat t.s
         .text
         .file   "t.c"
         .globl  test_refcounted                 # -- Begin function test_refcounted
         .p2align        3
         .type   test_refcounted,@function
test_refcounted:                        # @test_refcounted
# %bb.0:
         w1 &= 1
         if w1 == 0 goto LBB0_2
# %bb.1:
         r1 = r2
         call bpf_task_release
         goto LBB0_3
LBB0_2:
         r1 = r2
         call bpf_task_release
LBB0_3:
         w0 = 0
         exit
.Lfunc_end0:
         .size   test_refcounted, .Lfunc_end0-test_refcounted
                                         # -- End function
         .addrsig

> +
> +/* BTF FUNC records are not generated for kfuncs referenced
> + * from inline assembly. These records are necessary for
> + * libbpf to link the program. The function below is a hack
> + * to ensure that BTF FUNC records are generated.
> + */
> +void __btf_root(void)
> +{
> +	bpf_task_release(NULL);
> +}
> +
> +SEC(".struct_ops.link")
> +struct bpf_testmod_ops testmod_refcounted = {
> +	.test_refcounted = (void *)test_refcounted,
> +};
> +
> +
> diff --git a/tools/testing/selftests/bpf/progs/struct_ops_refcounted_fail__global_subprog.c b/tools/testing/selftests/bpf/progs/struct_ops_refcounted_fail__global_subprog.c
> new file mode 100644
> index 000000000000..c7e84e63b053
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/progs/struct_ops_refcounted_fail__global_subprog.c
> @@ -0,0 +1,32 @@
> +#include <vmlinux.h>
> +#include <bpf/bpf_tracing.h>
> +#include "../bpf_testmod/bpf_testmod.h"
> +
> +char _license[] SEC("license") = "GPL";
> +
> +extern void bpf_task_release(struct task_struct *p) __ksym;
> +
> +__noinline int subprog_release(__u64 *ctx __arg_ctx)
> +{
> +	struct task_struct *task = (struct task_struct *)ctx[1];
> +	int dummy = (int)ctx[0];
> +
> +	bpf_task_release(task);
> +
> +	return dummy + 1;
> +}
> +
> +SEC("struct_ops/test_refcounted")
> +int test_refcounted(unsigned long long *ctx)
> +{
> +	struct task_struct *task = (struct task_struct *)ctx[1];
> +
> +	bpf_task_release(task);
> +
> +	return subprog_release(ctx);
> +}
> +
> +SEC(".struct_ops.link")
> +struct bpf_testmod_ops testmod_ref_acquire = {
> +	.test_refcounted = (void *)test_refcounted,
> +};
> diff --git a/tools/testing/selftests/bpf/progs/struct_ops_refcounted_fail__ref_leak.c b/tools/testing/selftests/bpf/progs/struct_ops_refcounted_fail__ref_leak.c
> new file mode 100644
> index 000000000000..6e82859eb187
> --- /dev/null
> +++ b/tools/testing/selftests/bpf/progs/struct_ops_refcounted_fail__ref_leak.c
> @@ -0,0 +1,17 @@
> +#include <vmlinux.h>
> +#include <bpf/bpf_tracing.h>
> +#include "../bpf_testmod/bpf_testmod.h"
> +
> +char _license[] SEC("license") = "GPL";
> +
> +SEC("struct_ops/test_refcounted")
> +int BPF_PROG(test_refcounted, int dummy,
> +	     struct task_struct *task)
> +{
> +	return 0;
> +}
> +
> +SEC(".struct_ops.link")
> +struct bpf_testmod_ops testmod_ref_acquire = {
> +	.test_refcounted = (void *)test_refcounted,
> +};