[PATCH bpf-next 2/2] selftests/bpf: Cover RHash special-field recycle

[Nuoqi Gui]

Add regression coverage for RHash value recycling with referenced kptrs.

Store a referenced task kptr and update the value with BPF_EXIST. Verify
that bpf_kptr_xchg() still returns the kptr. Then store another task kptr
and delete the element to exercise the delayed reclamation path.

Signed-off-by: Nuoqi Gui <gnq25@xxxxxxxxxxxxxxxxxxxxx>
---
 tools/testing/selftests/bpf/prog_tests/rhash.c |  3 +
 tools/testing/selftests/bpf/progs/rhash.c      | 81 ++++++++++++++++++++++++++
 2 files changed, 84 insertions(+)

diff --git a/tools/testing/selftests/bpf/prog_tests/rhash.c b/tools/testing/selftests/bpf/prog_tests/rhash.c
index 98bb66907b7f..5bf2de828361 100644
--- a/tools/testing/selftests/bpf/prog_tests/rhash.c
+++ b/tools/testing/selftests/bpf/prog_tests/rhash.c
@@ -180,4 +180,7 @@ void test_rhash(void)
 
 	if (test__start_subtest("test_rhash_iter"))
 		rhash_iter_test();
+
+	if (test__start_subtest("test_rhash_special_fields_recycle"))
+		rhash_run("test_rhash_special_fields_recycle");
 }
diff --git a/tools/testing/selftests/bpf/progs/rhash.c b/tools/testing/selftests/bpf/progs/rhash.c
index fc2dac3a719e..d28d115cb74e 100644
--- a/tools/testing/selftests/bpf/progs/rhash.c
+++ b/tools/testing/selftests/bpf/progs/rhash.c
@@ -19,6 +19,11 @@ struct elem {
 	int val;
 };
 
+struct special_elem {
+	struct task_struct __kptr * task;
+	int val;
+};
+
 struct {
 	__uint(type, BPF_MAP_TYPE_RHASH);
 	__uint(map_flags, BPF_F_NO_PREALLOC);
@@ -27,6 +32,17 @@ struct {
 	__type(value, struct elem);
 } rhmap SEC(".maps");
 
+struct {
+	__uint(type, BPF_MAP_TYPE_RHASH);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+	__uint(max_entries, 4);
+	__type(key, int);
+	__type(value, struct special_elem);
+} special_fields SEC(".maps");
+
+struct task_struct *bpf_task_acquire(struct task_struct *p) __ksym;
+void bpf_task_release(struct task_struct *p) __ksym;
+
 SEC("syscall")
 int test_rhash_lookup_update(void *ctx)
 {
@@ -246,3 +262,68 @@ int test_rhash_delete_nonexistent(void *ctx)
 	err = 0;
 	return 0;
 }
+
+SEC("syscall")
+int test_rhash_special_fields_recycle(void *ctx)
+{
+	struct special_elem val1 = { .val = 1 };
+	struct special_elem val2 = { .val = 2 };
+	struct task_struct *task, *old;
+	struct special_elem *elem;
+	int key = 0;
+
+	err = 1;
+	if (bpf_map_update_elem(&special_fields, &key, &val1, BPF_NOEXIST))
+		return 1;
+
+	err = 2;
+	elem = bpf_map_lookup_elem(&special_fields, &key);
+	if (!elem)
+		return 2;
+
+	err = 3;
+	task = bpf_task_acquire(bpf_get_current_task_btf());
+	if (!task)
+		return 3;
+
+	err = 4;
+	old = bpf_kptr_xchg(&elem->task, task);
+	if (old) {
+		bpf_task_release(old);
+		return 4;
+	}
+
+	err = 5;
+	if (bpf_map_update_elem(&special_fields, &key, &val2, BPF_EXIST))
+		return 5;
+
+	err = 6;
+	elem = bpf_map_lookup_elem(&special_fields, &key);
+	if (!elem)
+		return 6;
+
+	err = 7;
+	task = bpf_kptr_xchg(&elem->task, NULL);
+	if (!task)
+		return 7;
+	bpf_task_release(task);
+
+	err = 8;
+	task = bpf_task_acquire(bpf_get_current_task_btf());
+	if (!task)
+		return 8;
+
+	err = 9;
+	old = bpf_kptr_xchg(&elem->task, task);
+	if (old) {
+		bpf_task_release(old);
+		return 9;
+	}
+
+	err = 10;
+	if (bpf_map_delete_elem(&special_fields, &key))
+		return 10;
+
+	err = 0;
+	return 0;
+}

-- 
2.34.1