[PATCH bpf-next 1/2] bpf: Add bpf_task_under_cgroup helper

[Feng zhou <zhoufeng.zf@bytedance.com>]

From: Feng Zhou <zhoufeng.zf@bytedance.com>

This adds a bpf helper that's similar to the
bpf_current_task_under_cgroup. The difference is that it is a
designated task.

When hook sched related functions, sometimes it is necessary to
specify a task instead of the current task.

Signed-off-by: Feng Zhou <zhoufeng.zf@bytedance.com>
---
 include/uapi/linux/bpf.h       | 13 +++++++++++++
 kernel/bpf/verifier.c          |  4 +++-
 kernel/trace/bpf_trace.c       | 31 +++++++++++++++++++++++++++++++
 tools/include/uapi/linux/bpf.h | 13 +++++++++++++
 4 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index 4b20a7269bee..3d31ddb39e10 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -5550,6 +5550,18 @@ union bpf_attr {
  *		0 on success.
  *
  *		**-ENOENT** if the bpf_local_storage cannot be found.
+ *
+ * long bpf_task_under_cgroup(struct bpf_map *map, struct task_struct *task, u32 index)
+ *	Description
+ *		Check whether the probe is being run is the context of a given
+ *		subset of the cgroup2 hierarchy. The cgroup2 to test is held by
+ *		*map* of type **BPF_MAP_TYPE_CGROUP_ARRAY**, at *index*.
+ *	Return
+ *		The return value depends on the result of the test, and can be:
+ *
+ *		* 1, if assigned task belongs to the cgroup2.
+ *		* 0, if assigned task does not belong to the cgroup2.
+ *		* A negative error code, if an error occurred.
  */
 #define ___BPF_FUNC_MAPPER(FN, ctx...)			\
 	FN(unspec, 0, ##ctx)				\
@@ -5764,6 +5776,7 @@ union bpf_attr {
 	FN(user_ringbuf_drain, 209, ##ctx)		\
 	FN(cgrp_storage_get, 210, ##ctx)		\
 	FN(cgrp_storage_delete, 211, ##ctx)		\
+	FN(task_under_cgroup, 212, ##ctx)		\
 	/* */
 
 /* backwards-compatibility macros for users of __BPF_FUNC_MAPPER that don't
diff --git a/kernel/bpf/verifier.c b/kernel/bpf/verifier.c
index 1e05355facdc..1e2c3c3e8d5f 100644
--- a/kernel/bpf/verifier.c
+++ b/kernel/bpf/verifier.c
@@ -7771,7 +7771,8 @@ static int check_map_func_compatibility(struct bpf_verifier_env *env,
 		break;
 	case BPF_MAP_TYPE_CGROUP_ARRAY:
 		if (func_id != BPF_FUNC_skb_under_cgroup &&
-		    func_id != BPF_FUNC_current_task_under_cgroup)
+		    func_id != BPF_FUNC_current_task_under_cgroup &&
+		    func_id != BPF_FUNC_task_under_cgroup)
 			goto error;
 		break;
 	case BPF_MAP_TYPE_CGROUP_STORAGE:
@@ -7902,6 +7903,7 @@ static int check_map_func_compatibility(struct bpf_verifier_env *env,
 			goto error;
 		break;
 	case BPF_FUNC_current_task_under_cgroup:
+	case BPF_FUNC_task_under_cgroup:
 	case BPF_FUNC_skb_under_cgroup:
 		if (map->map_type != BPF_MAP_TYPE_CGROUP_ARRAY)
 			goto error;
diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
index bcf91bc7bf71..b02a04768824 100644
--- a/kernel/trace/bpf_trace.c
+++ b/kernel/trace/bpf_trace.c
@@ -814,6 +814,35 @@ static const struct bpf_func_proto bpf_current_task_under_cgroup_proto = {
 	.arg2_type      = ARG_ANYTHING,
 };
 
+BPF_CALL_3(bpf_task_under_cgroup, struct bpf_map *, map, struct task_struct *,
+	   task, u32, idx)
+{
+	struct bpf_array *array = container_of(map, struct bpf_array, map);
+	struct cgroup *cgrp;
+
+	if (unlikely(!task))
+		return -ENOENT;
+
+	if (unlikely(idx >= array->map.max_entries))
+		return -E2BIG;
+
+	cgrp = READ_ONCE(array->ptrs[idx]);
+	if (unlikely(!cgrp))
+		return -EAGAIN;
+
+	return task_under_cgroup_hierarchy(task, cgrp);
+}
+
+static const struct bpf_func_proto bpf_task_under_cgroup_proto = {
+	.func           = bpf_task_under_cgroup,
+	.gpl_only       = false,
+	.ret_type       = RET_INTEGER,
+	.arg1_type      = ARG_CONST_MAP_PTR,
+	.arg2_type	= ARG_PTR_TO_BTF_ID,
+	.arg2_btf_id	= &btf_tracing_ids[BTF_TRACING_TYPE_TASK],
+	.arg3_type      = ARG_ANYTHING,
+};
+
 struct send_signal_irq_work {
 	struct irq_work irq_work;
 	struct task_struct *task;
@@ -1510,6 +1539,8 @@ bpf_tracing_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 		return &bpf_find_vma_proto;
 	case BPF_FUNC_trace_vprintk:
 		return bpf_get_trace_vprintk_proto();
+	case BPF_FUNC_task_under_cgroup:
+		return &bpf_task_under_cgroup_proto;
 	default:
 		return bpf_base_func_proto(func_id);
 	}
diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
index 4b20a7269bee..3d31ddb39e10 100644
--- a/tools/include/uapi/linux/bpf.h
+++ b/tools/include/uapi/linux/bpf.h
@@ -5550,6 +5550,18 @@ union bpf_attr {
  *		0 on success.
  *
  *		**-ENOENT** if the bpf_local_storage cannot be found.
+ *
+ * long bpf_task_under_cgroup(struct bpf_map *map, struct task_struct *task, u32 index)
+ *	Description
+ *		Check whether the probe is being run is the context of a given
+ *		subset of the cgroup2 hierarchy. The cgroup2 to test is held by
+ *		*map* of type **BPF_MAP_TYPE_CGROUP_ARRAY**, at *index*.
+ *	Return
+ *		The return value depends on the result of the test, and can be:
+ *
+ *		* 1, if assigned task belongs to the cgroup2.
+ *		* 0, if assigned task does not belong to the cgroup2.
+ *		* A negative error code, if an error occurred.
  */
 #define ___BPF_FUNC_MAPPER(FN, ctx...)			\
 	FN(unspec, 0, ##ctx)				\
@@ -5764,6 +5776,7 @@ union bpf_attr {
 	FN(user_ringbuf_drain, 209, ##ctx)		\
 	FN(cgrp_storage_get, 210, ##ctx)		\
 	FN(cgrp_storage_delete, 211, ##ctx)		\
+	FN(task_under_cgroup, 212, ##ctx)		\
 	/* */
 
 /* backwards-compatibility macros for users of __BPF_FUNC_MAPPER that don't
-- 
2.20.1