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 struct node_data { 12 long key; 13 long list_data; 14 struct bpf_rb_node r; 15 struct bpf_list_node l; 16 struct bpf_refcount ref; 17 }; 18 19 struct map_value { 20 struct node_data __kptr *node; 21 }; 22 23 struct { 24 __uint(type, BPF_MAP_TYPE_ARRAY); 25 __type(key, int); 26 __type(value, struct map_value); 27 __uint(max_entries, 1); 28 } stashed_nodes SEC(".maps"); 29 30 struct node_acquire { 31 long key; 32 long data; 33 struct bpf_rb_node node; 34 struct bpf_refcount refcount; 35 }; 36 37 #define private(name) SEC(".bss." #name) __hidden __attribute__((aligned(8))) 38 private(A) struct bpf_spin_lock lock; 39 private(A) struct bpf_rb_root root __contains(node_data, r); 40 private(A) struct bpf_list_head head __contains(node_data, l); 41 42 private(B) struct bpf_spin_lock alock; 43 private(B) struct bpf_rb_root aroot __contains(node_acquire, node); 44 45 static bool less(struct bpf_rb_node *node_a, const struct bpf_rb_node *node_b) 46 { 47 struct node_data *a; 48 struct node_data *b; 49 50 a = container_of(node_a, struct node_data, r); 51 b = container_of(node_b, struct node_data, r); 52 53 return a->key < b->key; 54 } 55 56 static bool less_a(struct bpf_rb_node *a, const struct bpf_rb_node *b) 57 { 58 struct node_acquire *node_a; 59 struct node_acquire *node_b; 60 61 node_a = container_of(a, struct node_acquire, node); 62 node_b = container_of(b, struct node_acquire, node); 63 64 return node_a->key < node_b->key; 65 } 66 67 static long __insert_in_tree_and_list(struct bpf_list_head *head, 68 struct bpf_rb_root *root, 69 struct bpf_spin_lock *lock) 70 { 71 struct node_data *n, *m; 72 73 n = bpf_obj_new(typeof(*n)); 74 if (!n) 75 return -1; 76 77 m = bpf_refcount_acquire(n); 78 m->key = 123; 79 m->list_data = 456; 80 81 bpf_spin_lock(lock); 82 if (bpf_rbtree_add(root, &n->r, less)) { 83 /* Failure to insert - unexpected */ 84 bpf_spin_unlock(lock); 85 bpf_obj_drop(m); 86 return -2; 87 } 88 bpf_spin_unlock(lock); 89 90 bpf_spin_lock(lock); 91 if (bpf_list_push_front(head, &m->l)) { 92 /* Failure to insert - unexpected */ 93 bpf_spin_unlock(lock); 94 return -3; 95 } 96 bpf_spin_unlock(lock); 97 return 0; 98 } 99 100 static long __stash_map_insert_tree(int idx, int val, struct bpf_rb_root *root, 101 struct bpf_spin_lock *lock) 102 { 103 struct map_value *mapval; 104 struct node_data *n, *m; 105 106 mapval = bpf_map_lookup_elem(&stashed_nodes, &idx); 107 if (!mapval) 108 return -1; 109 110 n = bpf_obj_new(typeof(*n)); 111 if (!n) 112 return -2; 113 114 n->key = val; 115 m = bpf_refcount_acquire(n); 116 117 n = bpf_kptr_xchg(&mapval->node, n); 118 if (n) { 119 bpf_obj_drop(n); 120 bpf_obj_drop(m); 121 return -3; 122 } 123 124 bpf_spin_lock(lock); 125 if (bpf_rbtree_add(root, &m->r, less)) { 126 /* Failure to insert - unexpected */ 127 bpf_spin_unlock(lock); 128 return -4; 129 } 130 bpf_spin_unlock(lock); 131 return 0; 132 } 133 134 static long __read_from_tree(struct bpf_rb_root *root, 135 struct bpf_spin_lock *lock, 136 bool remove_from_tree) 137 { 138 struct bpf_rb_node *rb; 139 struct node_data *n; 140 long res = -99; 141 142 bpf_spin_lock(lock); 143 144 rb = bpf_rbtree_first(root); 145 if (!rb) { 146 bpf_spin_unlock(lock); 147 return -1; 148 } 149 150 n = container_of(rb, struct node_data, r); 151 res = n->key; 152 153 if (!remove_from_tree) { 154 bpf_spin_unlock(lock); 155 return res; 156 } 157 158 rb = bpf_rbtree_remove(root, rb); 159 bpf_spin_unlock(lock); 160 if (!rb) 161 return -2; 162 n = container_of(rb, struct node_data, r); 163 bpf_obj_drop(n); 164 return res; 165 } 166 167 static long __read_from_list(struct bpf_list_head *head, 168 struct bpf_spin_lock *lock, 169 bool remove_from_list) 170 { 171 struct bpf_list_node *l; 172 struct node_data *n; 173 long res = -99; 174 175 bpf_spin_lock(lock); 176 177 l = bpf_list_pop_front(head); 178 if (!l) { 179 bpf_spin_unlock(lock); 180 return -1; 181 } 182 183 n = container_of(l, struct node_data, l); 184 res = n->list_data; 185 186 if (!remove_from_list) { 187 if (bpf_list_push_back(head, &n->l)) { 188 bpf_spin_unlock(lock); 189 return -2; 190 } 191 } 192 193 bpf_spin_unlock(lock); 194 195 if (remove_from_list) 196 bpf_obj_drop(n); 197 return res; 198 } 199 200 static long __read_from_unstash(int idx) 201 { 202 struct node_data *n = NULL; 203 struct map_value *mapval; 204 long val = -99; 205 206 mapval = bpf_map_lookup_elem(&stashed_nodes, &idx); 207 if (!mapval) 208 return -1; 209 210 n = bpf_kptr_xchg(&mapval->node, n); 211 if (!n) 212 return -2; 213 214 val = n->key; 215 bpf_obj_drop(n); 216 return val; 217 } 218 219 #define INSERT_READ_BOTH(rem_tree, rem_list, desc) \ 220 SEC("tc") \ 221 __description(desc) \ 222 __success __retval(579) \ 223 long insert_and_remove_tree_##rem_tree##_list_##rem_list(void *ctx) \ 224 { \ 225 long err, tree_data, list_data; \ 226 \ 227 err = __insert_in_tree_and_list(&head, &root, &lock); \ 228 if (err) \ 229 return err; \ 230 \ 231 err = __read_from_tree(&root, &lock, rem_tree); \ 232 if (err < 0) \ 233 return err; \ 234 else \ 235 tree_data = err; \ 236 \ 237 err = __read_from_list(&head, &lock, rem_list); \ 238 if (err < 0) \ 239 return err; \ 240 else \ 241 list_data = err; \ 242 \ 243 return tree_data + list_data; \ 244 } 245 246 /* After successful insert of struct node_data into both collections: 247 * - it should have refcount = 2 248 * - removing / not removing the node_data from a collection after 249 * reading should have no effect on ability to read / remove from 250 * the other collection 251 */ 252 INSERT_READ_BOTH(true, true, "insert_read_both: remove from tree + list"); 253 INSERT_READ_BOTH(false, false, "insert_read_both: remove from neither"); 254 INSERT_READ_BOTH(true, false, "insert_read_both: remove from tree"); 255 INSERT_READ_BOTH(false, true, "insert_read_both: remove from list"); 256 257 #undef INSERT_READ_BOTH 258 #define INSERT_READ_BOTH(rem_tree, rem_list, desc) \ 259 SEC("tc") \ 260 __description(desc) \ 261 __success __retval(579) \ 262 long insert_and_remove_lf_tree_##rem_tree##_list_##rem_list(void *ctx) \ 263 { \ 264 long err, tree_data, list_data; \ 265 \ 266 err = __insert_in_tree_and_list(&head, &root, &lock); \ 267 if (err) \ 268 return err; \ 269 \ 270 err = __read_from_list(&head, &lock, rem_list); \ 271 if (err < 0) \ 272 return err; \ 273 else \ 274 list_data = err; \ 275 \ 276 err = __read_from_tree(&root, &lock, rem_tree); \ 277 if (err < 0) \ 278 return err; \ 279 else \ 280 tree_data = err; \ 281 \ 282 return tree_data + list_data; \ 283 } 284 285 /* Similar to insert_read_both, but list data is read and possibly removed 286 * first 287 * 288 * Results should be no different than reading and possibly removing rbtree 289 * node first 290 */ 291 INSERT_READ_BOTH(true, true, "insert_read_both_list_first: remove from tree + list"); 292 INSERT_READ_BOTH(false, false, "insert_read_both_list_first: remove from neither"); 293 INSERT_READ_BOTH(true, false, "insert_read_both_list_first: remove from tree"); 294 INSERT_READ_BOTH(false, true, "insert_read_both_list_first: remove from list"); 295 296 #define INSERT_DOUBLE_READ_AND_DEL(read_fn, read_root, desc) \ 297 SEC("tc") \ 298 __description(desc) \ 299 __success __retval(-1) \ 300 long insert_double_##read_fn##_and_del_##read_root(void *ctx) \ 301 { \ 302 long err, list_data; \ 303 \ 304 err = __insert_in_tree_and_list(&head, &root, &lock); \ 305 if (err) \ 306 return err; \ 307 \ 308 err = read_fn(&read_root, &lock, true); \ 309 if (err < 0) \ 310 return err; \ 311 else \ 312 list_data = err; \ 313 \ 314 err = read_fn(&read_root, &lock, true); \ 315 if (err < 0) \ 316 return err; \ 317 \ 318 return err + list_data; \ 319 } 320 321 /* Insert into both tree and list, then try reading-and-removing from either twice 322 * 323 * The second read-and-remove should fail on read step since the node has 324 * already been removed 325 */ 326 INSERT_DOUBLE_READ_AND_DEL(__read_from_tree, root, "insert_double_del: 2x read-and-del from tree"); 327 INSERT_DOUBLE_READ_AND_DEL(__read_from_list, head, "insert_double_del: 2x read-and-del from list"); 328 329 #define INSERT_STASH_READ(rem_tree, desc) \ 330 SEC("tc") \ 331 __description(desc) \ 332 __success __retval(84) \ 333 long insert_rbtree_and_stash__del_tree_##rem_tree(void *ctx) \ 334 { \ 335 long err, tree_data, map_data; \ 336 \ 337 err = __stash_map_insert_tree(0, 42, &root, &lock); \ 338 if (err) \ 339 return err; \ 340 \ 341 err = __read_from_tree(&root, &lock, rem_tree); \ 342 if (err < 0) \ 343 return err; \ 344 else \ 345 tree_data = err; \ 346 \ 347 err = __read_from_unstash(0); \ 348 if (err < 0) \ 349 return err; \ 350 else \ 351 map_data = err; \ 352 \ 353 return tree_data + map_data; \ 354 } 355 356 /* Stash a refcounted node in map_val, insert same node into tree, then try 357 * reading data from tree then unstashed map_val, possibly removing from tree 358 * 359 * Removing from tree should have no effect on map_val kptr validity 360 */ 361 INSERT_STASH_READ(true, "insert_stash_read: remove from tree"); 362 INSERT_STASH_READ(false, "insert_stash_read: don't remove from tree"); 363 364 SEC("tc") 365 __success 366 long rbtree_refcounted_node_ref_escapes(void *ctx) 367 { 368 struct node_acquire *n, *m; 369 370 n = bpf_obj_new(typeof(*n)); 371 if (!n) 372 return 1; 373 374 bpf_spin_lock(&alock); 375 bpf_rbtree_add(&aroot, &n->node, less_a); 376 m = bpf_refcount_acquire(n); 377 bpf_spin_unlock(&alock); 378 379 m->key = 2; 380 bpf_obj_drop(m); 381 return 0; 382 } 383 384 SEC("tc") 385 __success 386 long rbtree_refcounted_node_ref_escapes_owning_input(void *ctx) 387 { 388 struct node_acquire *n, *m; 389 390 n = bpf_obj_new(typeof(*n)); 391 if (!n) 392 return 1; 393 394 m = bpf_refcount_acquire(n); 395 m->key = 2; 396 397 bpf_spin_lock(&alock); 398 bpf_rbtree_add(&aroot, &n->node, less_a); 399 bpf_spin_unlock(&alock); 400 401 bpf_obj_drop(m); 402 403 return 0; 404 } 405 406 char _license[] SEC("license") = "GPL"; 407