1// SPDX-License-Identifier: GPL-2.0
2/* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
3
4#include <vmlinux.h>
5#include <bpf/bpf_tracing.h>
6#include <bpf/bpf_helpers.h>
7
8#include "cgrp_kfunc_common.h"
9
10char _license[] SEC("license") = "GPL";
11
12int err, pid, invocations;
13
14/* Prototype for all of the program trace events below:
15 *
16 * TRACE_EVENT(cgroup_mkdir,
17 *         TP_PROTO(struct cgroup *cgrp, const char *path),
18 *         TP_ARGS(cgrp, path)
19 */
20
21static bool is_test_kfunc_task(void)
22{
23	int cur_pid = bpf_get_current_pid_tgid() >> 32;
24	bool same = pid == cur_pid;
25
26	if (same)
27		__sync_fetch_and_add(&invocations, 1);
28
29	return same;
30}
31
32SEC("tp_btf/cgroup_mkdir")
33int BPF_PROG(test_cgrp_acquire_release_argument, struct cgroup *cgrp, const char *path)
34{
35	struct cgroup *acquired;
36
37	if (!is_test_kfunc_task())
38		return 0;
39
40	acquired = bpf_cgroup_acquire(cgrp);
41	if (!acquired)
42		err = 1;
43	else
44		bpf_cgroup_release(acquired);
45
46	return 0;
47}
48
49SEC("tp_btf/cgroup_mkdir")
50int BPF_PROG(test_cgrp_acquire_leave_in_map, struct cgroup *cgrp, const char *path)
51{
52	long status;
53
54	if (!is_test_kfunc_task())
55		return 0;
56
57	status = cgrps_kfunc_map_insert(cgrp);
58	if (status)
59		err = 1;
60
61	return 0;
62}
63
64SEC("tp_btf/cgroup_mkdir")
65int BPF_PROG(test_cgrp_xchg_release, struct cgroup *cgrp, const char *path)
66{
67	struct cgroup *kptr, *cg;
68	struct __cgrps_kfunc_map_value *v;
69	long status;
70
71	if (!is_test_kfunc_task())
72		return 0;
73
74	status = cgrps_kfunc_map_insert(cgrp);
75	if (status) {
76		err = 1;
77		return 0;
78	}
79
80	v = cgrps_kfunc_map_value_lookup(cgrp);
81	if (!v) {
82		err = 2;
83		return 0;
84	}
85
86	kptr = v->cgrp;
87	if (!kptr) {
88		err = 4;
89		return 0;
90	}
91
92	cg = bpf_cgroup_ancestor(kptr, 1);
93	if (cg)	/* verifier only check */
94		bpf_cgroup_release(cg);
95
96	kptr = bpf_kptr_xchg(&v->cgrp, NULL);
97	if (!kptr) {
98		err = 3;
99		return 0;
100	}
101
102	bpf_cgroup_release(kptr);
103
104	return 0;
105}
106
107SEC("tp_btf/cgroup_mkdir")
108int BPF_PROG(test_cgrp_get_release, struct cgroup *cgrp, const char *path)
109{
110	struct cgroup *kptr;
111	struct __cgrps_kfunc_map_value *v;
112	long status;
113
114	if (!is_test_kfunc_task())
115		return 0;
116
117	status = cgrps_kfunc_map_insert(cgrp);
118	if (status) {
119		err = 1;
120		return 0;
121	}
122
123	v = cgrps_kfunc_map_value_lookup(cgrp);
124	if (!v) {
125		err = 2;
126		return 0;
127	}
128
129	bpf_rcu_read_lock();
130	kptr = v->cgrp;
131	if (!kptr)
132		err = 3;
133	bpf_rcu_read_unlock();
134
135	return 0;
136}
137
138SEC("tp_btf/cgroup_mkdir")
139int BPF_PROG(test_cgrp_get_ancestors, struct cgroup *cgrp, const char *path)
140{
141	struct cgroup *self, *ancestor1, *invalid;
142
143	if (!is_test_kfunc_task())
144		return 0;
145
146	self = bpf_cgroup_ancestor(cgrp, cgrp->level);
147	if (!self) {
148		err = 1;
149		return 0;
150	}
151
152	if (self->self.id != cgrp->self.id) {
153		bpf_cgroup_release(self);
154		err = 2;
155		return 0;
156	}
157	bpf_cgroup_release(self);
158
159	ancestor1 = bpf_cgroup_ancestor(cgrp, cgrp->level - 1);
160	if (!ancestor1) {
161		err = 3;
162		return 0;
163	}
164	bpf_cgroup_release(ancestor1);
165
166	invalid = bpf_cgroup_ancestor(cgrp, 10000);
167	if (invalid) {
168		bpf_cgroup_release(invalid);
169		err = 4;
170		return 0;
171	}
172
173	invalid = bpf_cgroup_ancestor(cgrp, -1);
174	if (invalid) {
175		bpf_cgroup_release(invalid);
176		err = 5;
177		return 0;
178	}
179
180	return 0;
181}
182
183SEC("tp_btf/cgroup_mkdir")
184int BPF_PROG(test_cgrp_from_id, struct cgroup *cgrp, const char *path)
185{
186	struct cgroup *parent, *res;
187	u64 parent_cgid;
188
189	if (!is_test_kfunc_task())
190		return 0;
191
192	/* @cgrp's ID is not visible yet, let's test with the parent */
193	parent = bpf_cgroup_ancestor(cgrp, cgrp->level - 1);
194	if (!parent) {
195		err = 1;
196		return 0;
197	}
198
199	parent_cgid = parent->kn->id;
200	bpf_cgroup_release(parent);
201
202	res = bpf_cgroup_from_id(parent_cgid);
203	if (!res) {
204		err = 2;
205		return 0;
206	}
207
208	bpf_cgroup_release(res);
209
210	if (res != parent) {
211		err = 3;
212		return 0;
213	}
214
215	res = bpf_cgroup_from_id((u64)-1);
216	if (res) {
217		bpf_cgroup_release(res);
218		err = 4;
219		return 0;
220	}
221
222	return 0;
223}
224