1 // SPDX-License-Identifier: GPL-2.0
2 /* Copyright (c) 2023 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 "cpumask_common.h"
9 
10 char _license[] SEC("license") = "GPL";
11 
12 int pid, nr_cpus;
13 
14 static bool is_test_task(void)
15 {
16 	int cur_pid = bpf_get_current_pid_tgid() >> 32;
17 
18 	return pid == cur_pid;
19 }
20 
21 static bool create_cpumask_set(struct bpf_cpumask **out1,
22 			       struct bpf_cpumask **out2,
23 			       struct bpf_cpumask **out3,
24 			       struct bpf_cpumask **out4)
25 {
26 	struct bpf_cpumask *mask1, *mask2, *mask3, *mask4;
27 
28 	mask1 = create_cpumask();
29 	if (!mask1)
30 		return false;
31 
32 	mask2 = create_cpumask();
33 	if (!mask2) {
34 		bpf_cpumask_release(mask1);
35 		err = 3;
36 		return false;
37 	}
38 
39 	mask3 = create_cpumask();
40 	if (!mask3) {
41 		bpf_cpumask_release(mask1);
42 		bpf_cpumask_release(mask2);
43 		err = 4;
44 		return false;
45 	}
46 
47 	mask4 = create_cpumask();
48 	if (!mask4) {
49 		bpf_cpumask_release(mask1);
50 		bpf_cpumask_release(mask2);
51 		bpf_cpumask_release(mask3);
52 		err = 5;
53 		return false;
54 	}
55 
56 	*out1 = mask1;
57 	*out2 = mask2;
58 	*out3 = mask3;
59 	*out4 = mask4;
60 
61 	return true;
62 }
63 
64 SEC("tp_btf/task_newtask")
65 int BPF_PROG(test_alloc_free_cpumask, struct task_struct *task, u64 clone_flags)
66 {
67 	struct bpf_cpumask *cpumask;
68 
69 	if (!is_test_task())
70 		return 0;
71 
72 	cpumask = create_cpumask();
73 	if (!cpumask)
74 		return 0;
75 
76 	bpf_cpumask_release(cpumask);
77 	return 0;
78 }
79 
80 SEC("tp_btf/task_newtask")
81 int BPF_PROG(test_set_clear_cpu, struct task_struct *task, u64 clone_flags)
82 {
83 	struct bpf_cpumask *cpumask;
84 
85 	if (!is_test_task())
86 		return 0;
87 
88 	cpumask = create_cpumask();
89 	if (!cpumask)
90 		return 0;
91 
92 	bpf_cpumask_set_cpu(0, cpumask);
93 	if (!bpf_cpumask_test_cpu(0, cast(cpumask))) {
94 		err = 3;
95 		goto release_exit;
96 	}
97 
98 	bpf_cpumask_clear_cpu(0, cpumask);
99 	if (bpf_cpumask_test_cpu(0, cast(cpumask))) {
100 		err = 4;
101 		goto release_exit;
102 	}
103 
104 release_exit:
105 	bpf_cpumask_release(cpumask);
106 	return 0;
107 }
108 
109 SEC("tp_btf/task_newtask")
110 int BPF_PROG(test_setall_clear_cpu, struct task_struct *task, u64 clone_flags)
111 {
112 	struct bpf_cpumask *cpumask;
113 
114 	if (!is_test_task())
115 		return 0;
116 
117 	cpumask = create_cpumask();
118 	if (!cpumask)
119 		return 0;
120 
121 	bpf_cpumask_setall(cpumask);
122 	if (!bpf_cpumask_full(cast(cpumask))) {
123 		err = 3;
124 		goto release_exit;
125 	}
126 
127 	bpf_cpumask_clear(cpumask);
128 	if (!bpf_cpumask_empty(cast(cpumask))) {
129 		err = 4;
130 		goto release_exit;
131 	}
132 
133 release_exit:
134 	bpf_cpumask_release(cpumask);
135 	return 0;
136 }
137 
138 SEC("tp_btf/task_newtask")
139 int BPF_PROG(test_first_firstzero_cpu, struct task_struct *task, u64 clone_flags)
140 {
141 	struct bpf_cpumask *cpumask;
142 
143 	if (!is_test_task())
144 		return 0;
145 
146 	cpumask = create_cpumask();
147 	if (!cpumask)
148 		return 0;
149 
150 	if (bpf_cpumask_first(cast(cpumask)) < nr_cpus) {
151 		err = 3;
152 		goto release_exit;
153 	}
154 
155 	if (bpf_cpumask_first_zero(cast(cpumask)) != 0) {
156 		bpf_printk("first zero: %d", bpf_cpumask_first_zero(cast(cpumask)));
157 		err = 4;
158 		goto release_exit;
159 	}
160 
161 	bpf_cpumask_set_cpu(0, cpumask);
162 	if (bpf_cpumask_first(cast(cpumask)) != 0) {
163 		err = 5;
164 		goto release_exit;
165 	}
166 
167 	if (bpf_cpumask_first_zero(cast(cpumask)) != 1) {
168 		err = 6;
169 		goto release_exit;
170 	}
171 
172 release_exit:
173 	bpf_cpumask_release(cpumask);
174 	return 0;
175 }
176 
177 SEC("tp_btf/task_newtask")
178 int BPF_PROG(test_test_and_set_clear, struct task_struct *task, u64 clone_flags)
179 {
180 	struct bpf_cpumask *cpumask;
181 
182 	if (!is_test_task())
183 		return 0;
184 
185 	cpumask = create_cpumask();
186 	if (!cpumask)
187 		return 0;
188 
189 	if (bpf_cpumask_test_and_set_cpu(0, cpumask)) {
190 		err = 3;
191 		goto release_exit;
192 	}
193 
194 	if (!bpf_cpumask_test_and_set_cpu(0, cpumask)) {
195 		err = 4;
196 		goto release_exit;
197 	}
198 
199 	if (!bpf_cpumask_test_and_clear_cpu(0, cpumask)) {
200 		err = 5;
201 		goto release_exit;
202 	}
203 
204 release_exit:
205 	bpf_cpumask_release(cpumask);
206 	return 0;
207 }
208 
209 SEC("tp_btf/task_newtask")
210 int BPF_PROG(test_and_or_xor, struct task_struct *task, u64 clone_flags)
211 {
212 	struct bpf_cpumask *mask1, *mask2, *dst1, *dst2;
213 
214 	if (!is_test_task())
215 		return 0;
216 
217 	if (!create_cpumask_set(&mask1, &mask2, &dst1, &dst2))
218 		return 0;
219 
220 	bpf_cpumask_set_cpu(0, mask1);
221 	bpf_cpumask_set_cpu(1, mask2);
222 
223 	if (bpf_cpumask_and(dst1, cast(mask1), cast(mask2))) {
224 		err = 6;
225 		goto release_exit;
226 	}
227 	if (!bpf_cpumask_empty(cast(dst1))) {
228 		err = 7;
229 		goto release_exit;
230 	}
231 
232 	bpf_cpumask_or(dst1, cast(mask1), cast(mask2));
233 	if (!bpf_cpumask_test_cpu(0, cast(dst1))) {
234 		err = 8;
235 		goto release_exit;
236 	}
237 	if (!bpf_cpumask_test_cpu(1, cast(dst1))) {
238 		err = 9;
239 		goto release_exit;
240 	}
241 
242 	bpf_cpumask_xor(dst2, cast(mask1), cast(mask2));
243 	if (!bpf_cpumask_equal(cast(dst1), cast(dst2))) {
244 		err = 10;
245 		goto release_exit;
246 	}
247 
248 release_exit:
249 	bpf_cpumask_release(mask1);
250 	bpf_cpumask_release(mask2);
251 	bpf_cpumask_release(dst1);
252 	bpf_cpumask_release(dst2);
253 	return 0;
254 }
255 
256 SEC("tp_btf/task_newtask")
257 int BPF_PROG(test_intersects_subset, struct task_struct *task, u64 clone_flags)
258 {
259 	struct bpf_cpumask *mask1, *mask2, *dst1, *dst2;
260 
261 	if (!is_test_task())
262 		return 0;
263 
264 	if (!create_cpumask_set(&mask1, &mask2, &dst1, &dst2))
265 		return 0;
266 
267 	bpf_cpumask_set_cpu(0, mask1);
268 	bpf_cpumask_set_cpu(1, mask2);
269 	if (bpf_cpumask_intersects(cast(mask1), cast(mask2))) {
270 		err = 6;
271 		goto release_exit;
272 	}
273 
274 	bpf_cpumask_or(dst1, cast(mask1), cast(mask2));
275 	if (!bpf_cpumask_subset(cast(mask1), cast(dst1))) {
276 		err = 7;
277 		goto release_exit;
278 	}
279 
280 	if (!bpf_cpumask_subset(cast(mask2), cast(dst1))) {
281 		err = 8;
282 		goto release_exit;
283 	}
284 
285 	if (bpf_cpumask_subset(cast(dst1), cast(mask1))) {
286 		err = 9;
287 		goto release_exit;
288 	}
289 
290 release_exit:
291 	bpf_cpumask_release(mask1);
292 	bpf_cpumask_release(mask2);
293 	bpf_cpumask_release(dst1);
294 	bpf_cpumask_release(dst2);
295 	return 0;
296 }
297 
298 SEC("tp_btf/task_newtask")
299 int BPF_PROG(test_copy_any_anyand, struct task_struct *task, u64 clone_flags)
300 {
301 	struct bpf_cpumask *mask1, *mask2, *dst1, *dst2;
302 	u32 cpu;
303 
304 	if (!is_test_task())
305 		return 0;
306 
307 	if (!create_cpumask_set(&mask1, &mask2, &dst1, &dst2))
308 		return 0;
309 
310 	bpf_cpumask_set_cpu(0, mask1);
311 	bpf_cpumask_set_cpu(1, mask2);
312 	bpf_cpumask_or(dst1, cast(mask1), cast(mask2));
313 
314 	cpu = bpf_cpumask_any(cast(mask1));
315 	if (cpu != 0) {
316 		err = 6;
317 		goto release_exit;
318 	}
319 
320 	cpu = bpf_cpumask_any(cast(dst2));
321 	if (cpu < nr_cpus) {
322 		err = 7;
323 		goto release_exit;
324 	}
325 
326 	bpf_cpumask_copy(dst2, cast(dst1));
327 	if (!bpf_cpumask_equal(cast(dst1), cast(dst2))) {
328 		err = 8;
329 		goto release_exit;
330 	}
331 
332 	cpu = bpf_cpumask_any(cast(dst2));
333 	if (cpu > 1) {
334 		err = 9;
335 		goto release_exit;
336 	}
337 
338 	cpu = bpf_cpumask_any_and(cast(mask1), cast(mask2));
339 	if (cpu < nr_cpus) {
340 		err = 10;
341 		goto release_exit;
342 	}
343 
344 release_exit:
345 	bpf_cpumask_release(mask1);
346 	bpf_cpumask_release(mask2);
347 	bpf_cpumask_release(dst1);
348 	bpf_cpumask_release(dst2);
349 	return 0;
350 }
351 
352 SEC("tp_btf/task_newtask")
353 int BPF_PROG(test_insert_leave, struct task_struct *task, u64 clone_flags)
354 {
355 	struct bpf_cpumask *cpumask;
356 
357 	cpumask = create_cpumask();
358 	if (!cpumask)
359 		return 0;
360 
361 	if (cpumask_map_insert(cpumask))
362 		err = 3;
363 
364 	return 0;
365 }
366 
367 SEC("tp_btf/task_newtask")
368 int BPF_PROG(test_insert_remove_release, struct task_struct *task, u64 clone_flags)
369 {
370 	struct bpf_cpumask *cpumask;
371 	struct __cpumask_map_value *v;
372 
373 	cpumask = create_cpumask();
374 	if (!cpumask)
375 		return 0;
376 
377 	if (cpumask_map_insert(cpumask)) {
378 		err = 3;
379 		return 0;
380 	}
381 
382 	v = cpumask_map_value_lookup();
383 	if (!v) {
384 		err = 4;
385 		return 0;
386 	}
387 
388 	cpumask = bpf_kptr_xchg(&v->cpumask, NULL);
389 	if (cpumask)
390 		bpf_cpumask_release(cpumask);
391 	else
392 		err = 5;
393 
394 	return 0;
395 }
396 
397 SEC("tp_btf/task_newtask")
398 int BPF_PROG(test_global_mask_rcu, struct task_struct *task, u64 clone_flags)
399 {
400 	struct bpf_cpumask *local, *prev;
401 
402 	if (!is_test_task())
403 		return 0;
404 
405 	local = create_cpumask();
406 	if (!local)
407 		return 0;
408 
409 	prev = bpf_kptr_xchg(&global_mask, local);
410 	if (prev) {
411 		bpf_cpumask_release(prev);
412 		err = 3;
413 		return 0;
414 	}
415 
416 	bpf_rcu_read_lock();
417 	local = global_mask;
418 	if (!local) {
419 		err = 4;
420 		bpf_rcu_read_unlock();
421 		return 0;
422 	}
423 
424 	bpf_cpumask_test_cpu(0, (const struct cpumask *)local);
425 	bpf_rcu_read_unlock();
426 
427 	return 0;
428 }
429