1 | // SPDX-License-Identifier: GPL-2.0-only |
2 | /* Copyright (c) 2023 Meta, Inc */ |
3 | #include <linux/bpf.h> |
4 | #include <linux/bpf_mem_alloc.h> |
5 | #include <linux/btf.h> |
6 | #include <linux/btf_ids.h> |
7 | #include <linux/cpumask.h> |
8 | |
9 | /** |
10 | * struct bpf_cpumask - refcounted BPF cpumask wrapper structure |
11 | * @cpumask: The actual cpumask embedded in the struct. |
12 | * @usage: Object reference counter. When the refcount goes to 0, the |
13 | * memory is released back to the BPF allocator, which provides |
14 | * RCU safety. |
15 | * |
16 | * Note that we explicitly embed a cpumask_t rather than a cpumask_var_t. This |
17 | * is done to avoid confusing the verifier due to the typedef of cpumask_var_t |
18 | * changing depending on whether CONFIG_CPUMASK_OFFSTACK is defined or not. See |
19 | * the details in <linux/cpumask.h>. The consequence is that this structure is |
20 | * likely a bit larger than it needs to be when CONFIG_CPUMASK_OFFSTACK is |
21 | * defined due to embedding the whole NR_CPUS-size bitmap, but the extra memory |
22 | * overhead is minimal. For the more typical case of CONFIG_CPUMASK_OFFSTACK |
23 | * not being defined, the structure is the same size regardless. |
24 | */ |
25 | struct bpf_cpumask { |
26 | cpumask_t cpumask; |
27 | refcount_t usage; |
28 | }; |
29 | |
30 | static struct bpf_mem_alloc bpf_cpumask_ma; |
31 | |
32 | static bool cpu_valid(u32 cpu) |
33 | { |
34 | return cpu < nr_cpu_ids; |
35 | } |
36 | |
37 | __bpf_kfunc_start_defs(); |
38 | |
39 | /** |
40 | * bpf_cpumask_create() - Create a mutable BPF cpumask. |
41 | * |
42 | * Allocates a cpumask that can be queried, mutated, acquired, and released by |
43 | * a BPF program. The cpumask returned by this function must either be embedded |
44 | * in a map as a kptr, or freed with bpf_cpumask_release(). |
45 | * |
46 | * bpf_cpumask_create() allocates memory using the BPF memory allocator, and |
47 | * will not block. It may return NULL if no memory is available. |
48 | */ |
49 | __bpf_kfunc struct bpf_cpumask *bpf_cpumask_create(void) |
50 | { |
51 | struct bpf_cpumask *cpumask; |
52 | |
53 | /* cpumask must be the first element so struct bpf_cpumask be cast to struct cpumask. */ |
54 | BUILD_BUG_ON(offsetof(struct bpf_cpumask, cpumask) != 0); |
55 | |
56 | cpumask = bpf_mem_cache_alloc(ma: &bpf_cpumask_ma); |
57 | if (!cpumask) |
58 | return NULL; |
59 | |
60 | memset(cpumask, 0, sizeof(*cpumask)); |
61 | refcount_set(r: &cpumask->usage, n: 1); |
62 | |
63 | return cpumask; |
64 | } |
65 | |
66 | /** |
67 | * bpf_cpumask_acquire() - Acquire a reference to a BPF cpumask. |
68 | * @cpumask: The BPF cpumask being acquired. The cpumask must be a trusted |
69 | * pointer. |
70 | * |
71 | * Acquires a reference to a BPF cpumask. The cpumask returned by this function |
72 | * must either be embedded in a map as a kptr, or freed with |
73 | * bpf_cpumask_release(). |
74 | */ |
75 | __bpf_kfunc struct bpf_cpumask *bpf_cpumask_acquire(struct bpf_cpumask *cpumask) |
76 | { |
77 | refcount_inc(r: &cpumask->usage); |
78 | return cpumask; |
79 | } |
80 | |
81 | /** |
82 | * bpf_cpumask_release() - Release a previously acquired BPF cpumask. |
83 | * @cpumask: The cpumask being released. |
84 | * |
85 | * Releases a previously acquired reference to a BPF cpumask. When the final |
86 | * reference of the BPF cpumask has been released, it is subsequently freed in |
87 | * an RCU callback in the BPF memory allocator. |
88 | */ |
89 | __bpf_kfunc void bpf_cpumask_release(struct bpf_cpumask *cpumask) |
90 | { |
91 | if (!refcount_dec_and_test(r: &cpumask->usage)) |
92 | return; |
93 | |
94 | migrate_disable(); |
95 | bpf_mem_cache_free_rcu(ma: &bpf_cpumask_ma, ptr: cpumask); |
96 | migrate_enable(); |
97 | } |
98 | |
99 | __bpf_kfunc void bpf_cpumask_release_dtor(void *cpumask) |
100 | { |
101 | bpf_cpumask_release(cpumask); |
102 | } |
103 | CFI_NOSEAL(bpf_cpumask_release_dtor); |
104 | |
105 | /** |
106 | * bpf_cpumask_first() - Get the index of the first nonzero bit in the cpumask. |
107 | * @cpumask: The cpumask being queried. |
108 | * |
109 | * Find the index of the first nonzero bit of the cpumask. A struct bpf_cpumask |
110 | * pointer may be safely passed to this function. |
111 | */ |
112 | __bpf_kfunc u32 bpf_cpumask_first(const struct cpumask *cpumask) |
113 | { |
114 | return cpumask_first(srcp: cpumask); |
115 | } |
116 | |
117 | /** |
118 | * bpf_cpumask_first_zero() - Get the index of the first unset bit in the |
119 | * cpumask. |
120 | * @cpumask: The cpumask being queried. |
121 | * |
122 | * Find the index of the first unset bit of the cpumask. A struct bpf_cpumask |
123 | * pointer may be safely passed to this function. |
124 | */ |
125 | __bpf_kfunc u32 bpf_cpumask_first_zero(const struct cpumask *cpumask) |
126 | { |
127 | return cpumask_first_zero(srcp: cpumask); |
128 | } |
129 | |
130 | /** |
131 | * bpf_cpumask_first_and() - Return the index of the first nonzero bit from the |
132 | * AND of two cpumasks. |
133 | * @src1: The first cpumask. |
134 | * @src2: The second cpumask. |
135 | * |
136 | * Find the index of the first nonzero bit of the AND of two cpumasks. |
137 | * struct bpf_cpumask pointers may be safely passed to @src1 and @src2. |
138 | */ |
139 | __bpf_kfunc u32 bpf_cpumask_first_and(const struct cpumask *src1, |
140 | const struct cpumask *src2) |
141 | { |
142 | return cpumask_first_and(srcp1: src1, srcp2: src2); |
143 | } |
144 | |
145 | /** |
146 | * bpf_cpumask_set_cpu() - Set a bit for a CPU in a BPF cpumask. |
147 | * @cpu: The CPU to be set in the cpumask. |
148 | * @cpumask: The BPF cpumask in which a bit is being set. |
149 | */ |
150 | __bpf_kfunc void bpf_cpumask_set_cpu(u32 cpu, struct bpf_cpumask *cpumask) |
151 | { |
152 | if (!cpu_valid(cpu)) |
153 | return; |
154 | |
155 | cpumask_set_cpu(cpu, dstp: (struct cpumask *)cpumask); |
156 | } |
157 | |
158 | /** |
159 | * bpf_cpumask_clear_cpu() - Clear a bit for a CPU in a BPF cpumask. |
160 | * @cpu: The CPU to be cleared from the cpumask. |
161 | * @cpumask: The BPF cpumask in which a bit is being cleared. |
162 | */ |
163 | __bpf_kfunc void bpf_cpumask_clear_cpu(u32 cpu, struct bpf_cpumask *cpumask) |
164 | { |
165 | if (!cpu_valid(cpu)) |
166 | return; |
167 | |
168 | cpumask_clear_cpu(cpu, dstp: (struct cpumask *)cpumask); |
169 | } |
170 | |
171 | /** |
172 | * bpf_cpumask_test_cpu() - Test whether a CPU is set in a cpumask. |
173 | * @cpu: The CPU being queried for. |
174 | * @cpumask: The cpumask being queried for containing a CPU. |
175 | * |
176 | * Return: |
177 | * * true - @cpu is set in the cpumask |
178 | * * false - @cpu was not set in the cpumask, or @cpu is an invalid cpu. |
179 | */ |
180 | __bpf_kfunc bool bpf_cpumask_test_cpu(u32 cpu, const struct cpumask *cpumask) |
181 | { |
182 | if (!cpu_valid(cpu)) |
183 | return false; |
184 | |
185 | return cpumask_test_cpu(cpu, cpumask: (struct cpumask *)cpumask); |
186 | } |
187 | |
188 | /** |
189 | * bpf_cpumask_test_and_set_cpu() - Atomically test and set a CPU in a BPF cpumask. |
190 | * @cpu: The CPU being set and queried for. |
191 | * @cpumask: The BPF cpumask being set and queried for containing a CPU. |
192 | * |
193 | * Return: |
194 | * * true - @cpu is set in the cpumask |
195 | * * false - @cpu was not set in the cpumask, or @cpu is invalid. |
196 | */ |
197 | __bpf_kfunc bool bpf_cpumask_test_and_set_cpu(u32 cpu, struct bpf_cpumask *cpumask) |
198 | { |
199 | if (!cpu_valid(cpu)) |
200 | return false; |
201 | |
202 | return cpumask_test_and_set_cpu(cpu, cpumask: (struct cpumask *)cpumask); |
203 | } |
204 | |
205 | /** |
206 | * bpf_cpumask_test_and_clear_cpu() - Atomically test and clear a CPU in a BPF |
207 | * cpumask. |
208 | * @cpu: The CPU being cleared and queried for. |
209 | * @cpumask: The BPF cpumask being cleared and queried for containing a CPU. |
210 | * |
211 | * Return: |
212 | * * true - @cpu is set in the cpumask |
213 | * * false - @cpu was not set in the cpumask, or @cpu is invalid. |
214 | */ |
215 | __bpf_kfunc bool bpf_cpumask_test_and_clear_cpu(u32 cpu, struct bpf_cpumask *cpumask) |
216 | { |
217 | if (!cpu_valid(cpu)) |
218 | return false; |
219 | |
220 | return cpumask_test_and_clear_cpu(cpu, cpumask: (struct cpumask *)cpumask); |
221 | } |
222 | |
223 | /** |
224 | * bpf_cpumask_setall() - Set all of the bits in a BPF cpumask. |
225 | * @cpumask: The BPF cpumask having all of its bits set. |
226 | */ |
227 | __bpf_kfunc void bpf_cpumask_setall(struct bpf_cpumask *cpumask) |
228 | { |
229 | cpumask_setall(dstp: (struct cpumask *)cpumask); |
230 | } |
231 | |
232 | /** |
233 | * bpf_cpumask_clear() - Clear all of the bits in a BPF cpumask. |
234 | * @cpumask: The BPF cpumask being cleared. |
235 | */ |
236 | __bpf_kfunc void bpf_cpumask_clear(struct bpf_cpumask *cpumask) |
237 | { |
238 | cpumask_clear(dstp: (struct cpumask *)cpumask); |
239 | } |
240 | |
241 | /** |
242 | * bpf_cpumask_and() - AND two cpumasks and store the result. |
243 | * @dst: The BPF cpumask where the result is being stored. |
244 | * @src1: The first input. |
245 | * @src2: The second input. |
246 | * |
247 | * Return: |
248 | * * true - @dst has at least one bit set following the operation |
249 | * * false - @dst is empty following the operation |
250 | * |
251 | * struct bpf_cpumask pointers may be safely passed to @src1 and @src2. |
252 | */ |
253 | __bpf_kfunc bool bpf_cpumask_and(struct bpf_cpumask *dst, |
254 | const struct cpumask *src1, |
255 | const struct cpumask *src2) |
256 | { |
257 | return cpumask_and(dstp: (struct cpumask *)dst, src1p: src1, src2p: src2); |
258 | } |
259 | |
260 | /** |
261 | * bpf_cpumask_or() - OR two cpumasks and store the result. |
262 | * @dst: The BPF cpumask where the result is being stored. |
263 | * @src1: The first input. |
264 | * @src2: The second input. |
265 | * |
266 | * struct bpf_cpumask pointers may be safely passed to @src1 and @src2. |
267 | */ |
268 | __bpf_kfunc void bpf_cpumask_or(struct bpf_cpumask *dst, |
269 | const struct cpumask *src1, |
270 | const struct cpumask *src2) |
271 | { |
272 | cpumask_or(dstp: (struct cpumask *)dst, src1p: src1, src2p: src2); |
273 | } |
274 | |
275 | /** |
276 | * bpf_cpumask_xor() - XOR two cpumasks and store the result. |
277 | * @dst: The BPF cpumask where the result is being stored. |
278 | * @src1: The first input. |
279 | * @src2: The second input. |
280 | * |
281 | * struct bpf_cpumask pointers may be safely passed to @src1 and @src2. |
282 | */ |
283 | __bpf_kfunc void bpf_cpumask_xor(struct bpf_cpumask *dst, |
284 | const struct cpumask *src1, |
285 | const struct cpumask *src2) |
286 | { |
287 | cpumask_xor(dstp: (struct cpumask *)dst, src1p: src1, src2p: src2); |
288 | } |
289 | |
290 | /** |
291 | * bpf_cpumask_equal() - Check two cpumasks for equality. |
292 | * @src1: The first input. |
293 | * @src2: The second input. |
294 | * |
295 | * Return: |
296 | * * true - @src1 and @src2 have the same bits set. |
297 | * * false - @src1 and @src2 differ in at least one bit. |
298 | * |
299 | * struct bpf_cpumask pointers may be safely passed to @src1 and @src2. |
300 | */ |
301 | __bpf_kfunc bool bpf_cpumask_equal(const struct cpumask *src1, const struct cpumask *src2) |
302 | { |
303 | return cpumask_equal(src1p: src1, src2p: src2); |
304 | } |
305 | |
306 | /** |
307 | * bpf_cpumask_intersects() - Check two cpumasks for overlap. |
308 | * @src1: The first input. |
309 | * @src2: The second input. |
310 | * |
311 | * Return: |
312 | * * true - @src1 and @src2 have at least one of the same bits set. |
313 | * * false - @src1 and @src2 don't have any of the same bits set. |
314 | * |
315 | * struct bpf_cpumask pointers may be safely passed to @src1 and @src2. |
316 | */ |
317 | __bpf_kfunc bool bpf_cpumask_intersects(const struct cpumask *src1, const struct cpumask *src2) |
318 | { |
319 | return cpumask_intersects(src1p: src1, src2p: src2); |
320 | } |
321 | |
322 | /** |
323 | * bpf_cpumask_subset() - Check if a cpumask is a subset of another. |
324 | * @src1: The first cpumask being checked as a subset. |
325 | * @src2: The second cpumask being checked as a superset. |
326 | * |
327 | * Return: |
328 | * * true - All of the bits of @src1 are set in @src2. |
329 | * * false - At least one bit in @src1 is not set in @src2. |
330 | * |
331 | * struct bpf_cpumask pointers may be safely passed to @src1 and @src2. |
332 | */ |
333 | __bpf_kfunc bool bpf_cpumask_subset(const struct cpumask *src1, const struct cpumask *src2) |
334 | { |
335 | return cpumask_subset(src1p: src1, src2p: src2); |
336 | } |
337 | |
338 | /** |
339 | * bpf_cpumask_empty() - Check if a cpumask is empty. |
340 | * @cpumask: The cpumask being checked. |
341 | * |
342 | * Return: |
343 | * * true - None of the bits in @cpumask are set. |
344 | * * false - At least one bit in @cpumask is set. |
345 | * |
346 | * A struct bpf_cpumask pointer may be safely passed to @cpumask. |
347 | */ |
348 | __bpf_kfunc bool bpf_cpumask_empty(const struct cpumask *cpumask) |
349 | { |
350 | return cpumask_empty(srcp: cpumask); |
351 | } |
352 | |
353 | /** |
354 | * bpf_cpumask_full() - Check if a cpumask has all bits set. |
355 | * @cpumask: The cpumask being checked. |
356 | * |
357 | * Return: |
358 | * * true - All of the bits in @cpumask are set. |
359 | * * false - At least one bit in @cpumask is cleared. |
360 | * |
361 | * A struct bpf_cpumask pointer may be safely passed to @cpumask. |
362 | */ |
363 | __bpf_kfunc bool bpf_cpumask_full(const struct cpumask *cpumask) |
364 | { |
365 | return cpumask_full(srcp: cpumask); |
366 | } |
367 | |
368 | /** |
369 | * bpf_cpumask_copy() - Copy the contents of a cpumask into a BPF cpumask. |
370 | * @dst: The BPF cpumask being copied into. |
371 | * @src: The cpumask being copied. |
372 | * |
373 | * A struct bpf_cpumask pointer may be safely passed to @src. |
374 | */ |
375 | __bpf_kfunc void bpf_cpumask_copy(struct bpf_cpumask *dst, const struct cpumask *src) |
376 | { |
377 | cpumask_copy(dstp: (struct cpumask *)dst, srcp: src); |
378 | } |
379 | |
380 | /** |
381 | * bpf_cpumask_any_distribute() - Return a random set CPU from a cpumask. |
382 | * @cpumask: The cpumask being queried. |
383 | * |
384 | * Return: |
385 | * * A random set bit within [0, num_cpus) if at least one bit is set. |
386 | * * >= num_cpus if no bit is set. |
387 | * |
388 | * A struct bpf_cpumask pointer may be safely passed to @src. |
389 | */ |
390 | __bpf_kfunc u32 bpf_cpumask_any_distribute(const struct cpumask *cpumask) |
391 | { |
392 | return cpumask_any_distribute(srcp: cpumask); |
393 | } |
394 | |
395 | /** |
396 | * bpf_cpumask_any_and_distribute() - Return a random set CPU from the AND of |
397 | * two cpumasks. |
398 | * @src1: The first cpumask. |
399 | * @src2: The second cpumask. |
400 | * |
401 | * Return: |
402 | * * A random set bit within [0, num_cpus) from the AND of two cpumasks, if at |
403 | * least one bit is set. |
404 | * * >= num_cpus if no bit is set. |
405 | * |
406 | * struct bpf_cpumask pointers may be safely passed to @src1 and @src2. |
407 | */ |
408 | __bpf_kfunc u32 bpf_cpumask_any_and_distribute(const struct cpumask *src1, |
409 | const struct cpumask *src2) |
410 | { |
411 | return cpumask_any_and_distribute(src1p: src1, src2p: src2); |
412 | } |
413 | |
414 | /** |
415 | * bpf_cpumask_weight() - Return the number of bits in @cpumask. |
416 | * @cpumask: The cpumask being queried. |
417 | * |
418 | * Count the number of set bits in the given cpumask. |
419 | */ |
420 | __bpf_kfunc u32 bpf_cpumask_weight(const struct cpumask *cpumask) |
421 | { |
422 | return cpumask_weight(srcp: cpumask); |
423 | } |
424 | |
425 | __bpf_kfunc_end_defs(); |
426 | |
427 | BTF_KFUNCS_START(cpumask_kfunc_btf_ids) |
428 | BTF_ID_FLAGS(func, bpf_cpumask_create, KF_ACQUIRE | KF_RET_NULL) |
429 | BTF_ID_FLAGS(func, bpf_cpumask_release, KF_RELEASE) |
430 | BTF_ID_FLAGS(func, bpf_cpumask_acquire, KF_ACQUIRE | KF_TRUSTED_ARGS) |
431 | BTF_ID_FLAGS(func, bpf_cpumask_first, KF_RCU) |
432 | BTF_ID_FLAGS(func, bpf_cpumask_first_zero, KF_RCU) |
433 | BTF_ID_FLAGS(func, bpf_cpumask_first_and, KF_RCU) |
434 | BTF_ID_FLAGS(func, bpf_cpumask_set_cpu, KF_RCU) |
435 | BTF_ID_FLAGS(func, bpf_cpumask_clear_cpu, KF_RCU) |
436 | BTF_ID_FLAGS(func, bpf_cpumask_test_cpu, KF_RCU) |
437 | BTF_ID_FLAGS(func, bpf_cpumask_test_and_set_cpu, KF_RCU) |
438 | BTF_ID_FLAGS(func, bpf_cpumask_test_and_clear_cpu, KF_RCU) |
439 | BTF_ID_FLAGS(func, bpf_cpumask_setall, KF_RCU) |
440 | BTF_ID_FLAGS(func, bpf_cpumask_clear, KF_RCU) |
441 | BTF_ID_FLAGS(func, bpf_cpumask_and, KF_RCU) |
442 | BTF_ID_FLAGS(func, bpf_cpumask_or, KF_RCU) |
443 | BTF_ID_FLAGS(func, bpf_cpumask_xor, KF_RCU) |
444 | BTF_ID_FLAGS(func, bpf_cpumask_equal, KF_RCU) |
445 | BTF_ID_FLAGS(func, bpf_cpumask_intersects, KF_RCU) |
446 | BTF_ID_FLAGS(func, bpf_cpumask_subset, KF_RCU) |
447 | BTF_ID_FLAGS(func, bpf_cpumask_empty, KF_RCU) |
448 | BTF_ID_FLAGS(func, bpf_cpumask_full, KF_RCU) |
449 | BTF_ID_FLAGS(func, bpf_cpumask_copy, KF_RCU) |
450 | BTF_ID_FLAGS(func, bpf_cpumask_any_distribute, KF_RCU) |
451 | BTF_ID_FLAGS(func, bpf_cpumask_any_and_distribute, KF_RCU) |
452 | BTF_ID_FLAGS(func, bpf_cpumask_weight, KF_RCU) |
453 | BTF_KFUNCS_END(cpumask_kfunc_btf_ids) |
454 | |
455 | static const struct btf_kfunc_id_set cpumask_kfunc_set = { |
456 | .owner = THIS_MODULE, |
457 | .set = &cpumask_kfunc_btf_ids, |
458 | }; |
459 | |
460 | BTF_ID_LIST(cpumask_dtor_ids) |
461 | BTF_ID(struct, bpf_cpumask) |
462 | BTF_ID(func, bpf_cpumask_release_dtor) |
463 | |
464 | static int __init cpumask_kfunc_init(void) |
465 | { |
466 | int ret; |
467 | const struct btf_id_dtor_kfunc cpumask_dtors[] = { |
468 | { |
469 | .btf_id = cpumask_dtor_ids[0], |
470 | .kfunc_btf_id = cpumask_dtor_ids[1] |
471 | }, |
472 | }; |
473 | |
474 | ret = bpf_mem_alloc_init(ma: &bpf_cpumask_ma, size: sizeof(struct bpf_cpumask), percpu: false); |
475 | ret = ret ?: register_btf_kfunc_id_set(prog_type: BPF_PROG_TYPE_TRACING, s: &cpumask_kfunc_set); |
476 | ret = ret ?: register_btf_kfunc_id_set(prog_type: BPF_PROG_TYPE_STRUCT_OPS, s: &cpumask_kfunc_set); |
477 | return ret ?: register_btf_id_dtor_kfuncs(dtors: cpumask_dtors, |
478 | ARRAY_SIZE(cpumask_dtors), |
479 | THIS_MODULE); |
480 | } |
481 | |
482 | late_initcall(cpumask_kfunc_init); |
483 | |