[PATCH] bpf: Cancel special fields on rhashtab value recycle instead of freeing

[Muhammad Falak R Wani <falakreyaz@gmail.com>]

Commit a3a81d2476512 ("bpf: Cancel special fields on map value recycle")
changed array and hashtab update/delete paths to avoid full special-field
destruction while recycling a map value. Such paths can run in NMI
context for tracing programs, where referenced kptr destructors are not
generally safe. Instead, bpf_obj_cancel_fields() cancels the NMI-safe
timer, workqueue, and task_work fields and leaves referenced kptr cleanup
to a later safe destruction path.

Resizable hashtable special-field support was added four days earlier by
commit 6905f8601298e ("bpf: Allow special fields in resizable hashtab"),
but its equivalent paths were not converted. BPF_MAP_TYPE_RHASH permits
referenced kptr fields and can be used by perf-event programs, which may
run in NMI context. rhtab_map_update_existing() and rhtab_delete_elem()
therefore still call bpf_obj_free_fields() directly and may invoke a
referenced kptr destructor from NMI context.

bpf_disable_instrumentation() around the rhashtable operations only
prevents a nested instrumentation program from reentering a bucket lock.
It does not change the execution context of the current program and
therefore does not make a subsequent destructor call NMI-safe.

Use bpf_obj_cancel_fields() in both paths, matching array and hashtab.
On delete, the allocator destructor performs full cleanup after the RCU
grace periods. On an in-place update, the kptr remains attached to the
map value until a BPF program explicitly removes it or the element is
later freed, matching the array-map semantics promised when rhashtable
special-field support was introduced.

Extend the map kptr lifetime test with an rhashtable variant. It stashes
a referenced kptr, updates the ordinary fields of the existing value,
and verifies that the update does not release the reference.

Fixes: 6905f8601298e ("bpf: Allow special fields in resizable hashtab")
Signed-off-by: Muhammad Falak R Wani <falakreyaz@gmail.com>
---
 kernel/bpf/hashtab.c                          | 26 ++++++++-----------
 .../selftests/bpf/prog_tests/map_kptr.c       | 11 ++++++++
 tools/testing/selftests/bpf/progs/map_kptr.c  | 12 +++++++++
 3 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/kernel/bpf/hashtab.c b/kernel/bpf/hashtab.c
index d40cb5dd446ca..a395928a3cf20 100644
--- a/kernel/bpf/hashtab.c
+++ b/kernel/bpf/hashtab.c
@@ -2864,16 +2864,6 @@ static int rhtab_map_alloc_check(union bpf_attr *attr)
 	return htab_map_alloc_check(attr);
 }
 
-static void rhtab_check_and_free_fields(struct bpf_rhtab *rhtab,
-					struct rhtab_elem *elem)
-{
-	if (IS_ERR_OR_NULL(rhtab->map.record))
-		return;
-
-	bpf_obj_free_fields(rhtab->map.record,
-			    rhtab_elem_value(elem, rhtab->map.key_size));
-}
-
 static void rhtab_mem_dtor(void *obj, void *ctx)
 {
 	struct htab_btf_record *hrec = ctx;
@@ -2963,8 +2953,12 @@ static int rhtab_delete_elem(struct bpf_rhtab *rhtab, struct rhtab_elem *elem, v
 		rhtab_read_elem_value(&rhtab->map, copy, elem, flags);
 		check_and_init_map_value(&rhtab->map, copy);
 	}
-	/* Release internal structs: kptr, bpf_timer, task_work, wq */
-	rhtab_check_and_free_fields(rhtab, elem);
+	/*
+	 * Cancel timer, workqueue, and task_work fields before deferring the
+	 * element free. Referenced kptr destruction is not NMI-safe, so leave
+	 * it for rhtab_mem_dtor() after the RCU grace periods.
+	 */
+	bpf_obj_cancel_fields(&rhtab->map, rhtab_elem_value(elem, rhtab->map.key_size));
 	bpf_mem_cache_free_rcu(&rhtab->ma, elem);
 	return 0;
 }
@@ -3022,10 +3016,12 @@ static long rhtab_map_update_existing(struct bpf_map *map, struct rhtab_elem *el
 	 * BPF_F_LOCK, matching arraymap semantics.
 	 *
 	 * copy_map_value() skips special-field offsets, so old timers/
-	 * kptrs/etc. still sit in the slot. Cancel them after the copy
-	 * to match arraymap's update semantics.
+	 * kptrs/etc. still sit in the slot. This path may run in NMI context,
+	 * so only cancel timer/workqueue/task_work here. Keep kptr fields
+	 * attached to the value, matching arraymap semantics; referenced
+	 * kptrs are destroyed when the element is eventually freed.
 	 */
-	rhtab_check_and_free_fields(rhtab, elem);
+	bpf_obj_cancel_fields(&rhtab->map, old_val);
 	return 0;
 }
 
diff --git a/tools/testing/selftests/bpf/prog_tests/map_kptr.c b/tools/testing/selftests/bpf/prog_tests/map_kptr.c
index 17e707dddda8d..9fddf03387bb8 100644
--- a/tools/testing/selftests/bpf/prog_tests/map_kptr.c
+++ b/tools/testing/selftests/bpf/prog_tests/map_kptr.c
@@ -98,6 +98,12 @@ static void test_map_kptr_success(bool test_run)
 	ASSERT_OK(ret, "test_map_kptr_ref3 refcount");
 	ASSERT_OK(opts.retval, "test_map_kptr_ref3 retval");
 
+	ret = bpf_map__delete_elem(skel->maps.rhash_map, &key, sizeof(key), 0);
+	ASSERT_OK(ret, "rhash_map delete");
+	ret = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.test_map_kptr_ref3), &opts);
+	ASSERT_OK(ret, "test_map_kptr_ref3 refcount");
+	ASSERT_OK(opts.retval, "test_map_kptr_ref3 retval");
+
 	ret = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.test_ls_map_kptr_ref_del), &lopts);
 	ASSERT_OK(ret, "test_ls_map_kptr_ref_del delete");
 	skel->data->ref--;
@@ -147,6 +153,7 @@ enum map_update_kptr_case {
 	MAP_UPDATE_KPTR_ARRAY,
 	MAP_UPDATE_KPTR_HASH,
 	MAP_UPDATE_KPTR_HASH_MALLOC,
+	MAP_UPDATE_KPTR_RHASH,
 };
 
 static struct bpf_program *map_update_kptr_prog(struct map_kptr *skel,
@@ -159,6 +166,8 @@ static struct bpf_program *map_update_kptr_prog(struct map_kptr *skel,
 		return skel->progs.test_hash_map_update_kptr;
 	case MAP_UPDATE_KPTR_HASH_MALLOC:
 		return skel->progs.test_hash_malloc_map_update_kptr;
+	case MAP_UPDATE_KPTR_RHASH:
+		return skel->progs.test_rhash_map_update_kptr;
 	}
 
 	return NULL;
@@ -204,6 +213,8 @@ void serial_test_map_kptr(void)
 		test_map_update_kptr(MAP_UPDATE_KPTR_HASH);
 	if (test__start_subtest("update_hash_malloc_map_kptr"))
 		test_map_update_kptr(MAP_UPDATE_KPTR_HASH_MALLOC);
+	if (test__start_subtest("update_rhash_map_kptr"))
+		test_map_update_kptr(MAP_UPDATE_KPTR_RHASH);
 
 	skel = rcu_tasks_trace_gp__open_and_load();
 	if (!ASSERT_OK_PTR(skel, "rcu_tasks_trace_gp__open_and_load"))
diff --git a/tools/testing/selftests/bpf/progs/map_kptr.c b/tools/testing/selftests/bpf/progs/map_kptr.c
index 0d87c97dac991..44210dd3c0ec9 100644
--- a/tools/testing/selftests/bpf/progs/map_kptr.c
+++ b/tools/testing/selftests/bpf/progs/map_kptr.c
@@ -57,6 +57,14 @@ struct hash_malloc_map {
 	__uint(map_flags, BPF_F_NO_PREALLOC);
 } hash_malloc_map SEC(".maps");
 
+struct {
+	__uint(type, BPF_MAP_TYPE_RHASH);
+	__type(key, int);
+	__type(value, struct map_value);
+	__uint(max_entries, 1);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+} rhash_map SEC(".maps");
+
 struct pcpu_hash_malloc_map {
 	__uint(type, BPF_MAP_TYPE_PERCPU_HASH);
 	__type(key, int);
@@ -421,6 +429,7 @@ int test_map_kptr_ref1(struct __sk_buff *ctx)
 	bpf_map_update_elem(&hash_map, &key, &val, 0);
 	bpf_map_update_elem(&hash_malloc_map, &key, &val, 0);
 	bpf_map_update_elem(&lru_hash_map, &key, &val, 0);
+	bpf_map_update_elem(&rhash_map, &key, &val, 0);
 
 	bpf_map_update_elem(&pcpu_hash_map, &key, &val, 0);
 	bpf_map_update_elem(&pcpu_hash_malloc_map, &key, &val, 0);
@@ -430,6 +439,7 @@ int test_map_kptr_ref1(struct __sk_buff *ctx)
 	TEST(hash_map);
 	TEST(hash_malloc_map);
 	TEST(lru_hash_map);
+	TEST(rhash_map);
 
 	TEST_PCPU(pcpu_array_map);
 	TEST_PCPU(pcpu_hash_map);
@@ -468,6 +478,7 @@ int test_map_kptr_ref2(struct __sk_buff *ctx)
 	TEST(hash_map);
 	TEST(hash_malloc_map);
 	TEST(lru_hash_map);
+	TEST(rhash_map);
 
 	TEST_PCPU(pcpu_array_map);
 	TEST_PCPU(pcpu_hash_map);
@@ -599,6 +610,7 @@ int name(void *ctx)						\
 
 DEFINE_HASH_UPDATE_KPTR_TEST(test_hash_map_update_kptr, hash_map)
 DEFINE_HASH_UPDATE_KPTR_TEST(test_hash_malloc_map_update_kptr, hash_malloc_map)
+DEFINE_HASH_UPDATE_KPTR_TEST(test_rhash_map_update_kptr, rhash_map)
 
 SEC("syscall")
 int test_ls_map_kptr_ref1(void *ctx)
-- 
2.55.0