[PATCH v3 bpf-next 5/5] selftests/bpf: Test bpf_kptr_xchg stashing into local kptr

[Amery Hung <ameryhung@xxxxxxxxx>]

From: Dave Marchevsky <davemarchevsky@xxxxxx>

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

Acked-by: Martin KaFai Lau <martin.lau@xxxxxxxxxx>
Signed-off-by: Dave Marchevsky <davemarchevsky@xxxxxx>
Signed-off-by: Amery Hung <amery.hung@xxxxxxxxxxxxx>
---
 .../selftests/bpf/progs/local_kptr_stash.c    | 58 ++++++++++++++++++-
 1 file changed, 56 insertions(+), 2 deletions(-)

diff --git a/tools/testing/selftests/bpf/progs/local_kptr_stash.c b/tools/testing/selftests/bpf/progs/local_kptr_stash.c
index 75043ffc5dad..ce7bf7790917 100644
--- a/tools/testing/selftests/bpf/progs/local_kptr_stash.c
+++ b/tools/testing/selftests/bpf/progs/local_kptr_stash.c
@@ -8,9 +8,13 @@
 #include "../bpf_experimental.h"
 #include "../bpf_testmod/bpf_testmod_kfunc.h"
 
+struct plain_local;
+
 struct node_data {
 	long key;
 	long data;
+	struct prog_test_ref_kfunc __kptr *stashed_in_kptr;
+	struct plain_local __kptr *stashed_in_local_kptr;
 	struct bpf_rb_node node;
 };
 
@@ -85,18 +89,52 @@ static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
 
 static int create_and_stash(int idx, int val)
 {
+	struct prog_test_ref_kfunc *inner_kptr;
+	struct plain_local *inner_local_kptr;
 	struct map_value *mapval;
 	struct node_data *res;
+	unsigned long dummy;
 
 	mapval = bpf_map_lookup_elem(&some_nodes, &idx);
 	if (!mapval)
 		return 1;
 
+	dummy = 0;
+	inner_kptr = bpf_kfunc_call_test_acquire(&dummy);
+	if (!inner_kptr)
+		return 2;
+
+	inner_local_kptr = bpf_obj_new(typeof(*inner_local_kptr));
+	if (!inner_local_kptr) {
+		bpf_kfunc_call_test_release(inner_kptr);
+		return 3;
+	}
+
 	res = bpf_obj_new(typeof(*res));
-	if (!res)
-		return 1;
+	if (!res) {
+		bpf_kfunc_call_test_release(inner_kptr);
+		bpf_obj_drop(inner_local_kptr);
+		return 4;
+	}
 	res->key = val;
 
+	inner_kptr = bpf_kptr_xchg(&res->stashed_in_kptr, inner_kptr);
+	if (inner_kptr) {
+		/* Should never happen, we just obj_new'd res */
+		bpf_kfunc_call_test_release(inner_kptr);
+		bpf_obj_drop(inner_local_kptr);
+		bpf_obj_drop(res);
+		return 5;
+	}
+
+	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 6;
+	}
+
 	res = bpf_kptr_xchg(&mapval->node, res);
 	if (res)
 		bpf_obj_drop(res);
@@ -169,6 +207,8 @@ long stash_local_with_root(void *ctx)
 SEC("tc")
 long unstash_rb_node(void *ctx)
 {
+	struct prog_test_ref_kfunc *inner_kptr = NULL;
+	struct plain_local *inner_local_kptr = NULL;
 	struct map_value *mapval;
 	struct node_data *res;
 	long retval;
@@ -180,6 +220,20 @@ long unstash_rb_node(void *ctx)
 
 	res = bpf_kptr_xchg(&mapval->node, NULL);
 	if (res) {
+		inner_kptr = bpf_kptr_xchg(&res->stashed_in_kptr, inner_kptr);
+		if (!inner_kptr) {
+			bpf_obj_drop(res);
+			return 1;
+		}
+		bpf_kfunc_call_test_release(inner_kptr);
+
+		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;
-- 
2.20.1