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