[PATCH bpf-next v2 2/2] selftests/bpf: Add tests for KF_OBTAIN kfunc

[Juntong Deng]

This patch adds test cases for KF_OBTAIN kfunc. Note that these test
cases are only used to test KF_OBTAIN and are not related to actual
usage scenarios.

kfunc_obtain_not_trusted is used to test that KF_OBTAIN kfunc only
accepts valid pointers.

kfunc_obtain_use_after_release is used to test that the returned pointer
becomes invalid after the pointer passed to KF_OBTAIN kfunc is released.

kfunc_obtain_trusted is the correct usage, valid pointers must be passed
to KF_OBTAIN kfunc and the returned pointer is only used if the passed
pointer has not been released.

Signed-off-by: Juntong Deng <juntong.deng@xxxxxxxxxxx>
---
 .../selftests/bpf/bpf_testmod/bpf_testmod.c   | 12 +++
 .../bpf/bpf_testmod/bpf_testmod_kfunc.h       |  3 +
 .../bpf/prog_tests/kfunc_obtain_test.c        | 10 +++
 .../selftests/bpf/progs/kfunc_obtain.c        | 74 +++++++++++++++++++
 4 files changed, 99 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/prog_tests/kfunc_obtain_test.c
 create mode 100644 tools/testing/selftests/bpf/progs/kfunc_obtain.c

diff --git a/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c b/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c
index bbf9442f0722..9e8a54dc88a2 100644
--- a/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c
+++ b/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod.c
@@ -183,6 +183,16 @@ __bpf_kfunc void bpf_kfunc_dynptr_test(struct bpf_dynptr *ptr,
 {
 }
 
+struct mm_struct *bpf_kfunc_obtain_test(struct task_struct *task)
+{
+	return task->mm;
+}
+
+struct task_struct *bpf_get_untrusted_task_test(struct task_struct *task)
+{
+	return task;
+}
+
 __bpf_kfunc struct bpf_testmod_ctx *
 bpf_testmod_ctx_create(int *err)
 {
@@ -541,6 +551,8 @@ BTF_ID_FLAGS(func, bpf_iter_testmod_seq_destroy, KF_ITER_DESTROY)
 BTF_ID_FLAGS(func, bpf_iter_testmod_seq_value)
 BTF_ID_FLAGS(func, bpf_kfunc_common_test)
 BTF_ID_FLAGS(func, bpf_kfunc_dynptr_test)
+BTF_ID_FLAGS(func, bpf_kfunc_obtain_test, KF_OBTAIN)
+BTF_ID_FLAGS(func, bpf_get_untrusted_task_test)
 BTF_ID_FLAGS(func, bpf_testmod_ctx_create, KF_ACQUIRE | KF_RET_NULL)
 BTF_ID_FLAGS(func, bpf_testmod_ctx_release, KF_RELEASE)
 BTF_KFUNCS_END(bpf_testmod_common_kfunc_ids)
diff --git a/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod_kfunc.h b/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod_kfunc.h
index e587a79f2239..cb38a211a9f3 100644
--- a/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod_kfunc.h
+++ b/tools/testing/selftests/bpf/bpf_testmod/bpf_testmod_kfunc.h
@@ -144,4 +144,7 @@ void bpf_kfunc_dynptr_test(struct bpf_dynptr *ptr, struct bpf_dynptr *ptr__nulla
 struct bpf_testmod_ctx *bpf_testmod_ctx_create(int *err) __ksym;
 void bpf_testmod_ctx_release(struct bpf_testmod_ctx *ctx) __ksym;
 
+struct mm_struct *bpf_kfunc_obtain_test(struct task_struct *task) __ksym;
+struct task_struct *bpf_get_untrusted_task_test(struct task_struct *task) __ksym;
+
 #endif /* _BPF_TESTMOD_KFUNC_H */
diff --git a/tools/testing/selftests/bpf/prog_tests/kfunc_obtain_test.c b/tools/testing/selftests/bpf/prog_tests/kfunc_obtain_test.c
new file mode 100644
index 000000000000..debc92fc1acc
--- /dev/null
+++ b/tools/testing/selftests/bpf/prog_tests/kfunc_obtain_test.c
@@ -0,0 +1,10 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <test_progs.h>
+#include "kfunc_obtain.skel.h"
+
+void test_kfunc_obtain(void)
+{
+	if (env.has_testmod)
+		RUN_TESTS(kfunc_obtain);
+}
diff --git a/tools/testing/selftests/bpf/progs/kfunc_obtain.c b/tools/testing/selftests/bpf/progs/kfunc_obtain.c
new file mode 100644
index 000000000000..8f0e074928ce
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/kfunc_obtain.c
@@ -0,0 +1,74 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include "vmlinux.h"
+#include <bpf/bpf_tracing.h>
+#include <bpf/bpf_helpers.h>
+#include "bpf_misc.h"
+#include "../bpf_testmod/bpf_testmod_kfunc.h"
+
+char _license[] SEC("license") = "GPL";
+
+struct task_struct *bpf_task_from_pid(s32 pid) __ksym;
+void bpf_task_release(struct task_struct *p) __ksym;
+
+/* The following test cases are only used to test KF_OBTAIN
+ * and are not related to actual usage scenarios.
+ */
+
+SEC("syscall")
+__failure __msg("must be referenced or trusted")
+int BPF_PROG(kfunc_obtain_not_trusted)
+{
+	struct task_struct *cur_task, *untrusted_task;
+
+	cur_task = bpf_get_current_task_btf();
+	untrusted_task = bpf_get_untrusted_task_test(cur_task);
+
+	bpf_kfunc_obtain_test(untrusted_task);
+
+	return 0;
+}
+
+SEC("syscall")
+__success
+int BPF_PROG(kfunc_obtain_trusted)
+{
+	struct task_struct *cur_task, *trusted_task;
+	struct mm_struct *mm;
+	int map_count = 0;
+
+	cur_task = bpf_get_current_task_btf();
+	trusted_task = bpf_task_from_pid(cur_task->pid);
+	if (trusted_task == NULL)
+		return 0;
+
+	mm = bpf_kfunc_obtain_test(trusted_task);
+
+	map_count = mm->map_count;
+
+	bpf_task_release(trusted_task);
+
+	return map_count;
+}
+
+SEC("syscall")
+__failure __msg("invalid mem access 'scalar'")
+int BPF_PROG(kfunc_obtain_use_after_release)
+{
+	struct task_struct *cur_task, *trusted_task;
+	struct mm_struct *mm;
+	int map_count = 0;
+
+	cur_task = bpf_get_current_task_btf();
+	trusted_task = bpf_task_from_pid(cur_task->pid);
+	if (trusted_task == NULL)
+		return 0;
+
+	mm = bpf_kfunc_obtain_test(trusted_task);
+
+	bpf_task_release(trusted_task);
+
+	map_count = mm->map_count;
+
+	return map_count;
+}
-- 
2.39.2