selftests/bpf: Test bpf_kptr_xchg stashing into local kptr

Test stashing both referenced kptr and local kptr into local kptrs. Then,
test unstashing them.

Acked-by: Martin KaFai Lau <martin.lau@kernel.org>
Acked-by: Hou Tao <houtao1@huawei.com>
Signed-off-by: Dave Marchevsky <davemarchevsky@fb.com>
Signed-off-by: Amery Hung <amery.hung@bytedance.com>
Link: https://lore.kernel.org/r/20240813212424.2871455-6-amery.hung@bytedance.com
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

diff --git a/tools/testing/selftests/bpf/progs/local_kptr_stash.c b/tools/testing/selftests/bpf/progs/local_kptr_stash.c
index 75043ff..b092a72 100644 (file)
--- a/tools/testing/selftests/bpf/progs/local_kptr_stash.c
+++ b/tools/testing/selftests/bpf/progs/local_kptr_stash.c
@@ -8,9 +8,12 @@
 #include "../bpf_experimental.h"
 #include "../bpf_testmod/bpf_testmod_kfunc.h"
 
+struct plain_local;
+
 struct node_data {
        long key;
        long data;
+       struct plain_local __kptr * stashed_in_local_kptr;
        struct bpf_rb_node node;
 };
 
@@ -85,6 +88,7 @@ static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
 
 static int create_and_stash(int idx, int val)
 {
+       struct plain_local *inner_local_kptr;
        struct map_value *mapval;
        struct node_data *res;
 
@@ -92,11 +96,25 @@ static int create_and_stash(int idx, int val)
        if (!mapval)
                return 1;
 
+       inner_local_kptr = bpf_obj_new(typeof(*inner_local_kptr));
+       if (!inner_local_kptr)
+               return 2;
+
        res = bpf_obj_new(typeof(*res));
-       if (!res)
-               return 1;
+       if (!res) {
+               bpf_obj_drop(inner_local_kptr);
+               return 3;
+       }
        res->key = val;
 
+       inner_local_kptr = bpf_kptr_xchg(&res->stashed_in_local_kptr, inner_local_kptr);
+       if (inner_local_kptr) {
+               /* Should never happen, we just obj_new'd res */
+               bpf_obj_drop(inner_local_kptr);
+               bpf_obj_drop(res);
+               return 4;
+       }
+
        res = bpf_kptr_xchg(&mapval->node, res);
        if (res)
                bpf_obj_drop(res);
@@ -169,6 +187,7 @@ long stash_local_with_root(void *ctx)
 SEC("tc")
 long unstash_rb_node(void *ctx)
 {
+       struct plain_local *inner_local_kptr = NULL;
        struct map_value *mapval;
        struct node_data *res;
        long retval;
@@ -180,6 +199,13 @@ long unstash_rb_node(void *ctx)
 
        res = bpf_kptr_xchg(&mapval->node, NULL);
        if (res) {
+               inner_local_kptr = bpf_kptr_xchg(&res->stashed_in_local_kptr, inner_local_kptr);
+               if (!inner_local_kptr) {
+                       bpf_obj_drop(res);
+                       return 1;
+               }
+               bpf_obj_drop(inner_local_kptr);
+
                retval = res->key;
                bpf_obj_drop(res);
                return retval;

diff --git a/tools/testing/selftests/bpf/progs/task_kfunc_success.c b/tools/testing/selftests/bpf/progs/task_kfunc_success.c
index 70df695..3138bb6 100644 (file)
--- a/tools/testing/selftests/bpf/progs/task_kfunc_success.c
+++ b/tools/testing/selftests/bpf/progs/task_kfunc_success.c
@@ -5,6 +5,7 @@
 #include <bpf/bpf_tracing.h>
 #include <bpf/bpf_helpers.h>
 
+#include "../bpf_experimental.h"
 #include "task_kfunc_common.h"
 
 char _license[] SEC("license") = "GPL";
@@ -143,7 +144,7 @@ SEC("tp_btf/task_newtask")
 int BPF_PROG(test_task_xchg_release, struct task_struct *task, u64 clone_flags)
 {
        struct task_struct *kptr;
-       struct __tasks_kfunc_map_value *v;
+       struct __tasks_kfunc_map_value *v, *local;
        long status;
 
        if (!is_test_kfunc_task())
@@ -167,6 +168,29 @@ int BPF_PROG(test_task_xchg_release, struct task_struct *task, u64 clone_flags)
                return 0;
        }
 
+       local = bpf_obj_new(typeof(*local));
+       if (!local) {
+               err = 4;
+               bpf_task_release(kptr);
+               return 0;
+       }
+
+       kptr = bpf_kptr_xchg(&local->task, kptr);
+       if (kptr) {
+               err = 5;
+               bpf_obj_drop(local);
+               bpf_task_release(kptr);
+               return 0;
+       }
+
+       kptr = bpf_kptr_xchg(&local->task, NULL);
+       if (!kptr) {
+               err = 6;
+               bpf_obj_drop(local);
+               return 0;
+       }
+
+       bpf_obj_drop(local);
        bpf_task_release(kptr);
 
        return 0;