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 */
25struct bpf_cpumask {
26 cpumask_t cpumask;
27 refcount_t usage;
28};
29
30static struct bpf_mem_alloc bpf_cpumask_ma;
31
32static 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}
103CFI_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
427BTF_KFUNCS_START(cpumask_kfunc_btf_ids)
428BTF_ID_FLAGS(func, bpf_cpumask_create, KF_ACQUIRE | KF_RET_NULL)
429BTF_ID_FLAGS(func, bpf_cpumask_release, KF_RELEASE)
430BTF_ID_FLAGS(func, bpf_cpumask_acquire, KF_ACQUIRE | KF_TRUSTED_ARGS)
431BTF_ID_FLAGS(func, bpf_cpumask_first, KF_RCU)
432BTF_ID_FLAGS(func, bpf_cpumask_first_zero, KF_RCU)
433BTF_ID_FLAGS(func, bpf_cpumask_first_and, KF_RCU)
434BTF_ID_FLAGS(func, bpf_cpumask_set_cpu, KF_RCU)
435BTF_ID_FLAGS(func, bpf_cpumask_clear_cpu, KF_RCU)
436BTF_ID_FLAGS(func, bpf_cpumask_test_cpu, KF_RCU)
437BTF_ID_FLAGS(func, bpf_cpumask_test_and_set_cpu, KF_RCU)
438BTF_ID_FLAGS(func, bpf_cpumask_test_and_clear_cpu, KF_RCU)
439BTF_ID_FLAGS(func, bpf_cpumask_setall, KF_RCU)
440BTF_ID_FLAGS(func, bpf_cpumask_clear, KF_RCU)
441BTF_ID_FLAGS(func, bpf_cpumask_and, KF_RCU)
442BTF_ID_FLAGS(func, bpf_cpumask_or, KF_RCU)
443BTF_ID_FLAGS(func, bpf_cpumask_xor, KF_RCU)
444BTF_ID_FLAGS(func, bpf_cpumask_equal, KF_RCU)
445BTF_ID_FLAGS(func, bpf_cpumask_intersects, KF_RCU)
446BTF_ID_FLAGS(func, bpf_cpumask_subset, KF_RCU)
447BTF_ID_FLAGS(func, bpf_cpumask_empty, KF_RCU)
448BTF_ID_FLAGS(func, bpf_cpumask_full, KF_RCU)
449BTF_ID_FLAGS(func, bpf_cpumask_copy, KF_RCU)
450BTF_ID_FLAGS(func, bpf_cpumask_any_distribute, KF_RCU)
451BTF_ID_FLAGS(func, bpf_cpumask_any_and_distribute, KF_RCU)
452BTF_ID_FLAGS(func, bpf_cpumask_weight, KF_RCU)
453BTF_KFUNCS_END(cpumask_kfunc_btf_ids)
454
455static const struct btf_kfunc_id_set cpumask_kfunc_set = {
456 .owner = THIS_MODULE,
457 .set = &cpumask_kfunc_btf_ids,
458};
459
460BTF_ID_LIST(cpumask_dtor_ids)
461BTF_ID(struct, bpf_cpumask)
462BTF_ID(func, bpf_cpumask_release_dtor)
463
464static 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
482late_initcall(cpumask_kfunc_init);
483

source code of linux/kernel/bpf/cpumask.c