[PATCH bpf v1 4/8] selftests/bpf: Test rhtab kptr cancellation semantics

[Kumar Kartikeya Dwivedi <memxor@gmail.com>]

From: Nuoqi Gui <gnq25@mails.tsinghua.edu.cn>

Resizable hash-map updates and deletions must not perform full special-field
destruction in their caller context. In particular, a referenced kptr must
remain attached to the allocation until the memory allocator destructor can
release it safely.

Add separate coverage for both affected paths. The update test stores a task
kptr, replaces the ordinary value bytes with BPF_EXIST, and verifies that the
kptr survived. The delete test removes an element and exchanges its kptr
through the still-valid map-value pointer before the allocation is reclaimed.

Both cases observe a NULL kptr when rhtab uses bpf_obj_free_fields(). They
recover and release the reference after rhtab switches to cancellation
semantics.

Signed-off-by: Nuoqi Gui <gnq25@mails.tsinghua.edu.cn>
[ kkd: Split update and delete coverage and rewrote the commit log ]
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 .../testing/selftests/bpf/prog_tests/rhash.c  |   6 +
 tools/testing/selftests/bpf/progs/rhash.c     | 112 ++++++++++++++++++
 2 files changed, 118 insertions(+)

diff --git a/tools/testing/selftests/bpf/prog_tests/rhash.c b/tools/testing/selftests/bpf/prog_tests/rhash.c
index 98bb66907b7f..0641bd5b0a9e 100644
--- a/tools/testing/selftests/bpf/prog_tests/rhash.c
+++ b/tools/testing/selftests/bpf/prog_tests/rhash.c
@@ -172,6 +172,12 @@ void test_rhash(void)
 	if (test__start_subtest("test_rhash_delete_nonexistent"))
 		rhash_run("test_rhash_delete_nonexistent");
 
+	if (test__start_subtest("test_rhash_kptr_update"))
+		rhash_run("test_rhash_kptr_update");
+
+	if (test__start_subtest("test_rhash_kptr_delete"))
+		rhash_run("test_rhash_kptr_delete");
+
 	if (test__start_subtest("test_rhash_map_extra_presize"))
 		rhash_map_extra_presize();
 
diff --git a/tools/testing/selftests/bpf/progs/rhash.c b/tools/testing/selftests/bpf/progs/rhash.c
index fc2dac3a719e..aea4de8dc781 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, 1);
+	__type(key, int);
+	__type(value, struct special_elem);
+} special_fields SEC(".maps");
+
+extern struct task_struct *bpf_task_acquire(struct task_struct *p) __ksym;
+extern void bpf_task_release(struct task_struct *p) __ksym;
+
 SEC("syscall")
 int test_rhash_lookup_update(void *ctx)
 {
@@ -246,3 +262,99 @@ int test_rhash_delete_nonexistent(void *ctx)
 	err = 0;
 	return 0;
 }
+
+SEC("syscall")
+int test_rhash_kptr_update(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 || elem->val != 2)
+		return 6;
+
+	err = 7;
+	old = bpf_kptr_xchg(&elem->task, NULL);
+	if (!old)
+		return 7;
+	bpf_task_release(old);
+
+	err = 8;
+	if (bpf_map_delete_elem(&special_fields, &key))
+		return 8;
+
+	err = 0;
+	return 0;
+}
+
+SEC("syscall")
+int test_rhash_kptr_delete(void *ctx)
+{
+	struct special_elem val = {};
+	struct task_struct *task, *old;
+	struct special_elem *elem;
+	int key = 0;
+
+	err = 1;
+	if (bpf_map_update_elem(&special_fields, &key, &val, 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_delete_elem(&special_fields, &key))
+		return 5;
+
+	err = 6;
+	old = bpf_kptr_xchg(&elem->task, NULL);
+	if (!old)
+		return 6;
+	bpf_task_release(old);
+
+	err = 0;
+	return 0;
+}
-- 
2.53.0