1 // SPDX-License-Identifier: GPL-2.0 2 /* Copyright (c) 2023 Meta Platforms, Inc. and affiliates. */ 3 4 #include <vmlinux.h> 5 #include <bpf/bpf_tracing.h> 6 #include <bpf/bpf_helpers.h> 7 #include <bpf/bpf_core_read.h> 8 #include "bpf_misc.h" 9 #include "bpf_experimental.h" 10 11 extern void bpf_rcu_read_lock(void) __ksym; 12 extern void bpf_rcu_read_unlock(void) __ksym; 13 14 struct node_data { 15 long key; 16 long list_data; 17 struct bpf_rb_node r; 18 struct bpf_list_node l; 19 struct bpf_refcount ref; 20 }; 21 22 struct map_value { 23 struct node_data __kptr *node; 24 }; 25 26 struct { 27 __uint(type, BPF_MAP_TYPE_ARRAY); 28 __type(key, int); 29 __type(value, struct map_value); 30 __uint(max_entries, 2); 31 } stashed_nodes SEC(".maps"); 32 33 struct node_acquire { 34 long key; 35 long data; 36 struct bpf_rb_node node; 37 struct bpf_refcount refcount; 38 }; 39 40 #define private(name) SEC(".bss." #name) __hidden __attribute__((aligned(8))) 41 private(A) struct bpf_spin_lock lock; 42 private(A) struct bpf_rb_root root __contains(node_data, r); 43 private(A) struct bpf_list_head head __contains(node_data, l); 44 45 private(B) struct bpf_spin_lock alock; 46 private(B) struct bpf_rb_root aroot __contains(node_acquire, node); 47 48 private(C) struct bpf_spin_lock block; 49 private(C) struct bpf_rb_root broot __contains(node_data, r); 50 51 static bool less(struct bpf_rb_node *node_a, const struct bpf_rb_node *node_b) 52 { 53 struct node_data *a; 54 struct node_data *b; 55 56 a = container_of(node_a, struct node_data, r); 57 b = container_of(node_b, struct node_data, r); 58 59 return a->key < b->key; 60 } 61 62 static bool less_a(struct bpf_rb_node *a, const struct bpf_rb_node *b) 63 { 64 struct node_acquire *node_a; 65 struct node_acquire *node_b; 66 67 node_a = container_of(a, struct node_acquire, node); 68 node_b = container_of(b, struct node_acquire, node); 69 70 return node_a->key < node_b->key; 71 } 72 73 static long __insert_in_tree_and_list(struct bpf_list_head *head, 74 struct bpf_rb_root *root, 75 struct bpf_spin_lock *lock) 76 { 77 struct node_data *n, *m; 78 79 n = bpf_obj_new(typeof(*n)); 80 if (!n) 81 return -1; 82 83 m = bpf_refcount_acquire(n); 84 m->key = 123; 85 m->list_data = 456; 86 87 bpf_spin_lock(lock); 88 if (bpf_rbtree_add(root, &n->r, less)) { 89 /* Failure to insert - unexpected */ 90 bpf_spin_unlock(lock); 91 bpf_obj_drop(m); 92 return -2; 93 } 94 bpf_spin_unlock(lock); 95 96 bpf_spin_lock(lock); 97 if (bpf_list_push_front(head, &m->l)) { 98 /* Failure to insert - unexpected */ 99 bpf_spin_unlock(lock); 100 return -3; 101 } 102 bpf_spin_unlock(lock); 103 return 0; 104 } 105 106 static long __stash_map_insert_tree(int idx, int val, struct bpf_rb_root *root, 107 struct bpf_spin_lock *lock) 108 { 109 struct map_value *mapval; 110 struct node_data *n, *m; 111 112 mapval = bpf_map_lookup_elem(&stashed_nodes, &idx); 113 if (!mapval) 114 return -1; 115 116 n = bpf_obj_new(typeof(*n)); 117 if (!n) 118 return -2; 119 120 n->key = val; 121 m = bpf_refcount_acquire(n); 122 123 n = bpf_kptr_xchg(&mapval->node, n); 124 if (n) { 125 bpf_obj_drop(n); 126 bpf_obj_drop(m); 127 return -3; 128 } 129 130 bpf_spin_lock(lock); 131 if (bpf_rbtree_add(root, &m->r, less)) { 132 /* Failure to insert - unexpected */ 133 bpf_spin_unlock(lock); 134 return -4; 135 } 136 bpf_spin_unlock(lock); 137 return 0; 138 } 139 140 static long __read_from_tree(struct bpf_rb_root *root, 141 struct bpf_spin_lock *lock, 142 bool remove_from_tree) 143 { 144 struct bpf_rb_node *rb; 145 struct node_data *n; 146 long res = -99; 147 148 bpf_spin_lock(lock); 149 150 rb = bpf_rbtree_first(root); 151 if (!rb) { 152 bpf_spin_unlock(lock); 153 return -1; 154 } 155 156 n = container_of(rb, struct node_data, r); 157 res = n->key; 158 159 if (!remove_from_tree) { 160 bpf_spin_unlock(lock); 161 return res; 162 } 163 164 rb = bpf_rbtree_remove(root, rb); 165 bpf_spin_unlock(lock); 166 if (!rb) 167 return -2; 168 n = container_of(rb, struct node_data, r); 169 bpf_obj_drop(n); 170 return res; 171 } 172 173 static long __read_from_list(struct bpf_list_head *head, 174 struct bpf_spin_lock *lock, 175 bool remove_from_list) 176 { 177 struct bpf_list_node *l; 178 struct node_data *n; 179 long res = -99; 180 181 bpf_spin_lock(lock); 182 183 l = bpf_list_pop_front(head); 184 if (!l) { 185 bpf_spin_unlock(lock); 186 return -1; 187 } 188 189 n = container_of(l, struct node_data, l); 190 res = n->list_data; 191 192 if (!remove_from_list) { 193 if (bpf_list_push_back(head, &n->l)) { 194 bpf_spin_unlock(lock); 195 return -2; 196 } 197 } 198 199 bpf_spin_unlock(lock); 200 201 if (remove_from_list) 202 bpf_obj_drop(n); 203 return res; 204 } 205 206 static long __read_from_unstash(int idx) 207 { 208 struct node_data *n = NULL; 209 struct map_value *mapval; 210 long val = -99; 211 212 mapval = bpf_map_lookup_elem(&stashed_nodes, &idx); 213 if (!mapval) 214 return -1; 215 216 n = bpf_kptr_xchg(&mapval->node, n); 217 if (!n) 218 return -2; 219 220 val = n->key; 221 bpf_obj_drop(n); 222 return val; 223 } 224 225 #define INSERT_READ_BOTH(rem_tree, rem_list, desc) \ 226 SEC("tc") \ 227 __description(desc) \ 228 __success __retval(579) \ 229 long insert_and_remove_tree_##rem_tree##_list_##rem_list(void *ctx) \ 230 { \ 231 long err, tree_data, list_data; \ 232 \ 233 err = __insert_in_tree_and_list(&head, &root, &lock); \ 234 if (err) \ 235 return err; \ 236 \ 237 err = __read_from_tree(&root, &lock, rem_tree); \ 238 if (err < 0) \ 239 return err; \ 240 else \ 241 tree_data = err; \ 242 \ 243 err = __read_from_list(&head, &lock, rem_list); \ 244 if (err < 0) \ 245 return err; \ 246 else \ 247 list_data = err; \ 248 \ 249 return tree_data + list_data; \ 250 } 251 252 /* After successful insert of struct node_data into both collections: 253 * - it should have refcount = 2 254 * - removing / not removing the node_data from a collection after 255 * reading should have no effect on ability to read / remove from 256 * the other collection 257 */ 258 INSERT_READ_BOTH(true, true, "insert_read_both: remove from tree + list"); 259 INSERT_READ_BOTH(false, false, "insert_read_both: remove from neither"); 260 INSERT_READ_BOTH(true, false, "insert_read_both: remove from tree"); 261 INSERT_READ_BOTH(false, true, "insert_read_both: remove from list"); 262 263 #undef INSERT_READ_BOTH 264 #define INSERT_READ_BOTH(rem_tree, rem_list, desc) \ 265 SEC("tc") \ 266 __description(desc) \ 267 __success __retval(579) \ 268 long insert_and_remove_lf_tree_##rem_tree##_list_##rem_list(void *ctx) \ 269 { \ 270 long err, tree_data, list_data; \ 271 \ 272 err = __insert_in_tree_and_list(&head, &root, &lock); \ 273 if (err) \ 274 return err; \ 275 \ 276 err = __read_from_list(&head, &lock, rem_list); \ 277 if (err < 0) \ 278 return err; \ 279 else \ 280 list_data = err; \ 281 \ 282 err = __read_from_tree(&root, &lock, rem_tree); \ 283 if (err < 0) \ 284 return err; \ 285 else \ 286 tree_data = err; \ 287 \ 288 return tree_data + list_data; \ 289 } 290 291 /* Similar to insert_read_both, but list data is read and possibly removed 292 * first 293 * 294 * Results should be no different than reading and possibly removing rbtree 295 * node first 296 */ 297 INSERT_READ_BOTH(true, true, "insert_read_both_list_first: remove from tree + list"); 298 INSERT_READ_BOTH(false, false, "insert_read_both_list_first: remove from neither"); 299 INSERT_READ_BOTH(true, false, "insert_read_both_list_first: remove from tree"); 300 INSERT_READ_BOTH(false, true, "insert_read_both_list_first: remove from list"); 301 302 #define INSERT_DOUBLE_READ_AND_DEL(read_fn, read_root, desc) \ 303 SEC("tc") \ 304 __description(desc) \ 305 __success __retval(-1) \ 306 long insert_double_##read_fn##_and_del_##read_root(void *ctx) \ 307 { \ 308 long err, list_data; \ 309 \ 310 err = __insert_in_tree_and_list(&head, &root, &lock); \ 311 if (err) \ 312 return err; \ 313 \ 314 err = read_fn(&read_root, &lock, true); \ 315 if (err < 0) \ 316 return err; \ 317 else \ 318 list_data = err; \ 319 \ 320 err = read_fn(&read_root, &lock, true); \ 321 if (err < 0) \ 322 return err; \ 323 \ 324 return err + list_data; \ 325 } 326 327 /* Insert into both tree and list, then try reading-and-removing from either twice 328 * 329 * The second read-and-remove should fail on read step since the node has 330 * already been removed 331 */ 332 INSERT_DOUBLE_READ_AND_DEL(__read_from_tree, root, "insert_double_del: 2x read-and-del from tree"); 333 INSERT_DOUBLE_READ_AND_DEL(__read_from_list, head, "insert_double_del: 2x read-and-del from list"); 334 335 #define INSERT_STASH_READ(rem_tree, desc) \ 336 SEC("tc") \ 337 __description(desc) \ 338 __success __retval(84) \ 339 long insert_rbtree_and_stash__del_tree_##rem_tree(void *ctx) \ 340 { \ 341 long err, tree_data, map_data; \ 342 \ 343 err = __stash_map_insert_tree(0, 42, &root, &lock); \ 344 if (err) \ 345 return err; \ 346 \ 347 err = __read_from_tree(&root, &lock, rem_tree); \ 348 if (err < 0) \ 349 return err; \ 350 else \ 351 tree_data = err; \ 352 \ 353 err = __read_from_unstash(0); \ 354 if (err < 0) \ 355 return err; \ 356 else \ 357 map_data = err; \ 358 \ 359 return tree_data + map_data; \ 360 } 361 362 /* Stash a refcounted node in map_val, insert same node into tree, then try 363 * reading data from tree then unstashed map_val, possibly removing from tree 364 * 365 * Removing from tree should have no effect on map_val kptr validity 366 */ 367 INSERT_STASH_READ(true, "insert_stash_read: remove from tree"); 368 INSERT_STASH_READ(false, "insert_stash_read: don't remove from tree"); 369 370 SEC("tc") 371 __success 372 long rbtree_refcounted_node_ref_escapes(void *ctx) 373 { 374 struct node_acquire *n, *m; 375 376 n = bpf_obj_new(typeof(*n)); 377 if (!n) 378 return 1; 379 380 bpf_spin_lock(&alock); 381 bpf_rbtree_add(&aroot, &n->node, less_a); 382 m = bpf_refcount_acquire(n); 383 bpf_spin_unlock(&alock); 384 if (!m) 385 return 2; 386 387 m->key = 2; 388 bpf_obj_drop(m); 389 return 0; 390 } 391 392 SEC("tc") 393 __success 394 long rbtree_refcounted_node_ref_escapes_owning_input(void *ctx) 395 { 396 struct node_acquire *n, *m; 397 398 n = bpf_obj_new(typeof(*n)); 399 if (!n) 400 return 1; 401 402 m = bpf_refcount_acquire(n); 403 m->key = 2; 404 405 bpf_spin_lock(&alock); 406 bpf_rbtree_add(&aroot, &n->node, less_a); 407 bpf_spin_unlock(&alock); 408 409 bpf_obj_drop(m); 410 411 return 0; 412 } 413 414 static long __stash_map_empty_xchg(struct node_data *n, int idx) 415 { 416 struct map_value *mapval = bpf_map_lookup_elem(&stashed_nodes, &idx); 417 418 if (!mapval) { 419 bpf_obj_drop(n); 420 return 1; 421 } 422 n = bpf_kptr_xchg(&mapval->node, n); 423 if (n) { 424 bpf_obj_drop(n); 425 return 2; 426 } 427 return 0; 428 } 429 430 SEC("tc") 431 long rbtree_wrong_owner_remove_fail_a1(void *ctx) 432 { 433 struct node_data *n, *m; 434 435 n = bpf_obj_new(typeof(*n)); 436 if (!n) 437 return 1; 438 m = bpf_refcount_acquire(n); 439 440 if (__stash_map_empty_xchg(n, 0)) { 441 bpf_obj_drop(m); 442 return 2; 443 } 444 445 if (__stash_map_empty_xchg(m, 1)) 446 return 3; 447 448 return 0; 449 } 450 451 SEC("tc") 452 long rbtree_wrong_owner_remove_fail_b(void *ctx) 453 { 454 struct map_value *mapval; 455 struct node_data *n; 456 int idx = 0; 457 458 mapval = bpf_map_lookup_elem(&stashed_nodes, &idx); 459 if (!mapval) 460 return 1; 461 462 n = bpf_kptr_xchg(&mapval->node, NULL); 463 if (!n) 464 return 2; 465 466 bpf_spin_lock(&block); 467 468 bpf_rbtree_add(&broot, &n->r, less); 469 470 bpf_spin_unlock(&block); 471 return 0; 472 } 473 474 SEC("tc") 475 long rbtree_wrong_owner_remove_fail_a2(void *ctx) 476 { 477 struct map_value *mapval; 478 struct bpf_rb_node *res; 479 struct node_data *m; 480 int idx = 1; 481 482 mapval = bpf_map_lookup_elem(&stashed_nodes, &idx); 483 if (!mapval) 484 return 1; 485 486 m = bpf_kptr_xchg(&mapval->node, NULL); 487 if (!m) 488 return 2; 489 bpf_spin_lock(&lock); 490 491 /* make m non-owning ref */ 492 bpf_list_push_back(&head, &m->l); 493 res = bpf_rbtree_remove(&root, &m->r); 494 495 bpf_spin_unlock(&lock); 496 if (res) { 497 bpf_obj_drop(container_of(res, struct node_data, r)); 498 return 3; 499 } 500 return 0; 501 } 502 503 SEC("?fentry.s/bpf_testmod_test_read") 504 __success 505 int BPF_PROG(rbtree_sleepable_rcu, 506 struct file *file, struct kobject *kobj, 507 struct bin_attribute *bin_attr, char *buf, loff_t off, size_t len) 508 { 509 struct bpf_rb_node *rb; 510 struct node_data *n, *m = NULL; 511 512 n = bpf_obj_new(typeof(*n)); 513 if (!n) 514 return 0; 515 516 bpf_rcu_read_lock(); 517 bpf_spin_lock(&lock); 518 bpf_rbtree_add(&root, &n->r, less); 519 rb = bpf_rbtree_first(&root); 520 if (!rb) 521 goto err_out; 522 523 rb = bpf_rbtree_remove(&root, rb); 524 if (!rb) 525 goto err_out; 526 527 m = container_of(rb, struct node_data, r); 528 529 err_out: 530 bpf_spin_unlock(&lock); 531 bpf_rcu_read_unlock(); 532 if (m) 533 bpf_obj_drop(m); 534 return 0; 535 } 536 537 SEC("?fentry.s/bpf_testmod_test_read") 538 __success 539 int BPF_PROG(rbtree_sleepable_rcu_no_explicit_rcu_lock, 540 struct file *file, struct kobject *kobj, 541 struct bin_attribute *bin_attr, char *buf, loff_t off, size_t len) 542 { 543 struct bpf_rb_node *rb; 544 struct node_data *n, *m = NULL; 545 546 n = bpf_obj_new(typeof(*n)); 547 if (!n) 548 return 0; 549 550 /* No explicit bpf_rcu_read_lock */ 551 bpf_spin_lock(&lock); 552 bpf_rbtree_add(&root, &n->r, less); 553 rb = bpf_rbtree_first(&root); 554 if (!rb) 555 goto err_out; 556 557 rb = bpf_rbtree_remove(&root, rb); 558 if (!rb) 559 goto err_out; 560 561 m = container_of(rb, struct node_data, r); 562 563 err_out: 564 bpf_spin_unlock(&lock); 565 /* No explicit bpf_rcu_read_unlock */ 566 if (m) 567 bpf_obj_drop(m); 568 return 0; 569 } 570 571 char _license[] SEC("license") = "GPL"; 572