selftests/bpf: Add selftests for bpf_task_from_pid()

Add some selftest testcases that validate the expected behavior of the
bpf_task_from_pid() kfunc that was added in the prior patch.

Signed-off-by: David Vernet <[e-mail]>
Link: https://lore.kernel.org/r/[e-mail]
Signed-off-by: Alexei Starovoitov <[e-mail]>

tools/testing/selftests/bpf/prog_tests/task_kfunc.c

···
 	"test_task_xchg_release",
 	"test_task_get_release",
 	"test_task_current_acquire_release",
+	"test_task_from_pid_arg",
+	"test_task_from_pid_current",
+	"test_task_from_pid_invalid",
 };
 
 static struct {
···
 	{"task_kfunc_release_fp", "arg#0 pointer type STRUCT task_struct must point"},
 	{"task_kfunc_release_null", "arg#0 is ptr_or_null_ expected ptr_ or socket"},
 	{"task_kfunc_release_unacquired", "release kernel function bpf_task_release expects"},
+	{"task_kfunc_from_pid_no_null_check", "arg#0 is ptr_or_null_ expected ptr_ or socket"},
 };
 
 static void verify_fail(const char *prog_name, const char *expected_err_msg)

tools/testing/selftests/bpf/progs/task_kfunc_common.h

···
 struct task_struct *bpf_task_acquire(struct task_struct *p) __ksym;
 struct task_struct *bpf_task_kptr_get(struct task_struct **pp) __ksym;
 void bpf_task_release(struct task_struct *p) __ksym;
+struct task_struct *bpf_task_from_pid(s32 pid) __ksym;
 
 static inline struct __tasks_kfunc_map_value *tasks_kfunc_map_value_lookup(struct task_struct *p)
 {

tools/testing/selftests/bpf/progs/task_kfunc_failure.c

···
 
 	return 0;
 }
+
+SEC("tp_btf/task_newtask")
+int BPF_PROG(task_kfunc_from_pid_no_null_check, struct task_struct *task, u64 clone_flags)
+{
+	struct task_struct *acquired;
+
+	acquired = bpf_task_from_pid(task->pid);
+
+	/* Releasing bpf_task_from_pid() lookup without a NULL check. */
+	bpf_task_release(acquired);
+
+	return 0;
+}

tools/testing/selftests/bpf/progs/task_kfunc_success.c

···
 
 	return 0;
 }
+
+static void lookup_compare_pid(const struct task_struct *p)
+{
+	struct task_struct *acquired;
+
+	acquired = bpf_task_from_pid(p->pid);
+	if (!acquired) {
+		err = 1;
+		return;
+	}
+
+	if (acquired->pid != p->pid)
+		err = 2;
+	bpf_task_release(acquired);
+}
+
+SEC("tp_btf/task_newtask")
+int BPF_PROG(test_task_from_pid_arg, struct task_struct *task, u64 clone_flags)
+{
+	struct task_struct *acquired;
+
+	if (!is_test_kfunc_task())
+		return 0;
+
+	lookup_compare_pid(task);
+	return 0;
+}
+
+SEC("tp_btf/task_newtask")
+int BPF_PROG(test_task_from_pid_current, struct task_struct *task, u64 clone_flags)
+{
+	struct task_struct *current, *acquired;
+
+	if (!is_test_kfunc_task())
+		return 0;
+
+	lookup_compare_pid(bpf_get_current_task_btf());
+	return 0;
+}
+
+static int is_pid_lookup_valid(s32 pid)
+{
+	struct task_struct *acquired;
+
+	acquired = bpf_task_from_pid(pid);
+	if (acquired) {
+		bpf_task_release(acquired);
+		return 1;
+	}
+
+	return 0;
+}
+
+SEC("tp_btf/task_newtask")
+int BPF_PROG(test_task_from_pid_invalid, struct task_struct *task, u64 clone_flags)
+{
+	struct task_struct *acquired;
+
+	if (!is_test_kfunc_task())
+		return 0;
+
+	if (is_pid_lookup_valid(-1)) {
+		err = 1;
+		return 0;
+	}
+
+	if (is_pid_lookup_valid(0xcafef00d)) {
+		err = 2;
+		return 0;
+	}
+
+	return 0;
+}