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(&bpf_cpumask_ma); 57 if (!cpumask) 58 return NULL; 59 60 memset(cpumask, 0, sizeof(*cpumask)); 61 refcount_set(&cpumask->usage, 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(&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(&cpumask->usage)) 92 return; 93 94 migrate_disable(); 95 bpf_mem_cache_free_rcu(&bpf_cpumask_ma, 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(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(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(src1, 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, (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, (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, (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, (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, (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((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((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((struct cpumask *)dst, src1, 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((struct cpumask *)dst, src1, 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((struct cpumask *)dst, src1, 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(src1, 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(src1, 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(src1, 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(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(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((struct cpumask *)dst, 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(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(src1, 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(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(&bpf_cpumask_ma, sizeof(struct bpf_cpumask), false); 475 ret = ret ?: register_btf_kfunc_id_set(BPF_PROG_TYPE_TRACING, &cpumask_kfunc_set); 476 ret = ret ?: register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS, &cpumask_kfunc_set); 477 ret = ret ?: register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL, &cpumask_kfunc_set); 478 return ret ?: register_btf_id_dtor_kfuncs(cpumask_dtors, 479 ARRAY_SIZE(cpumask_dtors), 480 THIS_MODULE); 481} 482 483late_initcall(cpumask_kfunc_init); 484