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
is_test_task(void)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
create_cpumask_set(struct bpf_cpumask ** out1,struct bpf_cpumask ** out2,struct bpf_cpumask ** out3,struct bpf_cpumask ** out4)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")
BPF_PROG(test_alloc_free_cpumask,struct task_struct * task,u64 clone_flags)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")
BPF_PROG(test_set_clear_cpu,struct task_struct * task,u64 clone_flags)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")
BPF_PROG(test_setall_clear_cpu,struct task_struct * task,u64 clone_flags)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")
BPF_PROG(test_first_firstzero_cpu,struct task_struct * task,u64 clone_flags)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")
BPF_PROG(test_test_and_set_clear,struct task_struct * task,u64 clone_flags)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")
BPF_PROG(test_and_or_xor,struct task_struct * task,u64 clone_flags)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")
BPF_PROG(test_intersects_subset,struct task_struct * task,u64 clone_flags)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")
BPF_PROG(test_copy_any_anyand,struct task_struct * task,u64 clone_flags)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")
BPF_PROG(test_insert_leave,struct task_struct * task,u64 clone_flags)353 int BPF_PROG(test_insert_leave, struct task_struct *task, u64 clone_flags)
354 {
355 struct bpf_cpumask *cpumask;
356 struct __cpumask_map_value *v;
357
358 cpumask = create_cpumask();
359 if (!cpumask)
360 return 0;
361
362 if (cpumask_map_insert(cpumask))
363 err = 3;
364
365 return 0;
366 }
367
368 SEC("tp_btf/task_newtask")
BPF_PROG(test_insert_remove_release,struct task_struct * task,u64 clone_flags)369 int BPF_PROG(test_insert_remove_release, struct task_struct *task, u64 clone_flags)
370 {
371 struct bpf_cpumask *cpumask;
372 struct __cpumask_map_value *v;
373
374 cpumask = create_cpumask();
375 if (!cpumask)
376 return 0;
377
378 if (cpumask_map_insert(cpumask)) {
379 err = 3;
380 return 0;
381 }
382
383 v = cpumask_map_value_lookup();
384 if (!v) {
385 err = 4;
386 return 0;
387 }
388
389 cpumask = bpf_kptr_xchg(&v->cpumask, NULL);
390 if (cpumask)
391 bpf_cpumask_release(cpumask);
392 else
393 err = 5;
394
395 return 0;
396 }
397
398 SEC("tp_btf/task_newtask")
BPF_PROG(test_insert_kptr_get_release,struct task_struct * task,u64 clone_flags)399 int BPF_PROG(test_insert_kptr_get_release, struct task_struct *task, u64 clone_flags)
400 {
401 struct bpf_cpumask *cpumask;
402 struct __cpumask_map_value *v;
403
404 cpumask = create_cpumask();
405 if (!cpumask)
406 return 0;
407
408 if (cpumask_map_insert(cpumask)) {
409 err = 3;
410 return 0;
411 }
412
413 v = cpumask_map_value_lookup();
414 if (!v) {
415 err = 4;
416 return 0;
417 }
418
419 cpumask = bpf_cpumask_kptr_get(&v->cpumask);
420 if (cpumask)
421 bpf_cpumask_release(cpumask);
422 else
423 err = 5;
424
425 return 0;
426 }
427