kernel/bpf/cpumask.c

   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(&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 }
 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(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
 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(&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         return  ret ?: register_btf_id_dtor_kfuncs(cpumask_dtors,
 478                                                    ARRAY_SIZE(cpumask_dtors),
 479                                                    THIS_MODULE);
 480 }
 481
 482 late_initcall(cpumask_kfunc_init);