[PATCH bpf-next 3/4] bpf: Add new kfunc bpf_cpumask_set_from_pid

[Yafang Shao <laoar.shao@gmail.com>]

Introducing a new kfunc: bpf_cpumask_set_from_pid. This function serves the
purpose of retrieving the cpumask associated with a specific PID. Its
utility is particularly evident within container environments. For
instance, it allows for extracting the cpuset of a container using the
init task within it.

Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 kernel/bpf/cpumask.c | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/kernel/bpf/cpumask.c b/kernel/bpf/cpumask.c
index 4ae07a4..5755bb6 100644
--- a/kernel/bpf/cpumask.c
+++ b/kernel/bpf/cpumask.c
@@ -467,6 +467,22 @@ __bpf_kfunc void bpf_iter_cpumask_destroy(struct bpf_iter_cpumask *it)
 	bpf_mem_free(&bpf_global_ma, kit->cpu);
 }
 
+__bpf_kfunc bool bpf_cpumask_set_from_pid(struct cpumask *cpumask, u32 pid)
+{
+	struct task_struct *task;
+
+	if (!cpumask)
+		return false;
+
+	task = get_pid_task(find_vpid(pid), PIDTYPE_PID);
+	if (!task)
+		return false;
+
+	cpumask_copy(cpumask, task->cpus_ptr);
+	put_task_struct(task);
+	return true;
+}
+
 __bpf_kfunc_end_defs();
 
 BTF_SET8_START(cpumask_kfunc_btf_ids)
@@ -498,6 +514,7 @@ __bpf_kfunc void bpf_iter_cpumask_destroy(struct bpf_iter_cpumask *it)
 BTF_ID_FLAGS(func, bpf_iter_cpumask_new, KF_ITER_NEW | KF_RCU)
 BTF_ID_FLAGS(func, bpf_iter_cpumask_next, KF_ITER_NEXT | KF_RET_NULL | KF_RCU)
 BTF_ID_FLAGS(func, bpf_iter_cpumask_destroy, KF_ITER_DESTROY)
+BTF_ID_FLAGS(func, bpf_cpumask_set_from_pid, KF_RCU)
 BTF_SET8_END(cpumask_kfunc_btf_ids)
 
 static const struct btf_kfunc_id_set cpumask_kfunc_set = {
-- 
1.8.3.1