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