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 __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  */
  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  */
  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  */
  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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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  */
 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)
 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);