1// SPDX-License-Identifier: GPL-2.0
2/* Copyright (C) 2023 Chuyi Zhou <zhouchuyi@bytedance.com> */
3
4#include "vmlinux.h"
5#include <bpf/bpf_helpers.h>
6#include <bpf/bpf_tracing.h>
7#include "bpf_misc.h"
8#include "bpf_experimental.h"
9
10char _license[] SEC("license") = "GPL";
11
12pid_t target_pid;
13u64 root_cg_id, leaf_cg_id;
14u64 first_cg_id, last_cg_id;
15
16int pre_order_cnt, post_order_cnt, tree_high;
17
18struct cgroup *bpf_cgroup_from_id(u64 cgid) __ksym;
19void bpf_cgroup_release(struct cgroup *p) __ksym;
20void bpf_rcu_read_lock(void) __ksym;
21void bpf_rcu_read_unlock(void) __ksym;
22
23SEC("fentry.s/" SYS_PREFIX "sys_getpgid")
24int iter_css_for_each(const void *ctx)
25{
26	struct task_struct *cur_task = bpf_get_current_task_btf();
27	struct cgroup_subsys_state *root_css, *leaf_css, *pos;
28	struct cgroup *root_cgrp, *leaf_cgrp, *cur_cgrp;
29
30	if (cur_task->pid != target_pid)
31		return 0;
32
33	root_cgrp = bpf_cgroup_from_id(root_cg_id);
34
35	if (!root_cgrp)
36		return 0;
37
38	leaf_cgrp = bpf_cgroup_from_id(leaf_cg_id);
39
40	if (!leaf_cgrp) {
41		bpf_cgroup_release(root_cgrp);
42		return 0;
43	}
44	root_css = &root_cgrp->self;
45	leaf_css = &leaf_cgrp->self;
46	pre_order_cnt = post_order_cnt = tree_high = 0;
47	first_cg_id = last_cg_id = 0;
48
49	bpf_rcu_read_lock();
50	bpf_for_each(css, pos, root_css, BPF_CGROUP_ITER_DESCENDANTS_POST) {
51		cur_cgrp = pos->cgroup;
52		post_order_cnt++;
53		last_cg_id = cur_cgrp->kn->id;
54	}
55
56	bpf_for_each(css, pos, root_css, BPF_CGROUP_ITER_DESCENDANTS_PRE) {
57		cur_cgrp = pos->cgroup;
58		pre_order_cnt++;
59		if (!first_cg_id)
60			first_cg_id = cur_cgrp->kn->id;
61	}
62
63	bpf_for_each(css, pos, leaf_css, BPF_CGROUP_ITER_ANCESTORS_UP)
64		tree_high++;
65
66	bpf_for_each(css, pos, root_css, BPF_CGROUP_ITER_ANCESTORS_UP)
67		tree_high--;
68	bpf_rcu_read_unlock();
69	bpf_cgroup_release(root_cgrp);
70	bpf_cgroup_release(leaf_cgrp);
71	return 0;
72}
73