tools/testing/selftests/bpf/progs/rbtree.c

   1 // SPDX-License-Identifier: GPL-2.0
   2 /* Copyright (c) 2022 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_experimental.h"
   9
  10 struct node_data {
  11         long key;
  12         long data;
  13         struct bpf_rb_node node;
  14 };
  15
  16 struct root_nested_inner {
  17         struct bpf_spin_lock glock;
  18         struct bpf_rb_root root __contains(node_data, node);
  19 };
  20
  21 struct root_nested {
  22         struct root_nested_inner inner;
  23 };
  24
  25 long less_callback_ran = -1;
  26 long removed_key = -1;
  27 long first_data[2] = {-1, -1};
  28
  29 #define private(name) SEC(".data." #name) __hidden __attribute__((aligned(8)))
  30 private(A) struct bpf_spin_lock glock;
  31 private(A) struct bpf_rb_root groot __contains(node_data, node);
  32 private(A) struct bpf_rb_root groot_array[2] __contains(node_data, node);
  33 private(A) struct bpf_rb_root groot_array_one[1] __contains(node_data, node);
  34 private(B) struct root_nested groot_nested;
  35
  36 static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
  37 {
  38         struct node_data *node_a;
  39         struct node_data *node_b;
  40
  41         node_a = container_of(a, struct node_data, node);
  42         node_b = container_of(b, struct node_data, node);
  43         less_callback_ran = 1;
  44
  45         return node_a->key < node_b->key;
  46 }
  47
  48 static long __add_three(struct bpf_rb_root *root, struct bpf_spin_lock *lock)
  49 {
  50         struct node_data *n, *m;
  51
  52         n = bpf_obj_new(typeof(*n));
  53         if (!n)
  54                 return 1;
  55         n->key = 5;
  56
  57         m = bpf_obj_new(typeof(*m));
  58         if (!m) {
  59                 bpf_obj_drop(n);
  60                 return 2;
  61         }
  62         m->key = 1;
  63
  64         bpf_spin_lock(&glock);
  65         bpf_rbtree_add(&groot, &n->node, less);
  66         bpf_rbtree_add(&groot, &m->node, less);
  67         bpf_spin_unlock(&glock);
  68
  69         n = bpf_obj_new(typeof(*n));
  70         if (!n)
  71                 return 3;
  72         n->key = 3;
  73
  74         bpf_spin_lock(&glock);
  75         bpf_rbtree_add(&groot, &n->node, less);
  76         bpf_spin_unlock(&glock);
  77         return 0;
  78 }
  79
  80 SEC("tc")
  81 long rbtree_add_nodes(void *ctx)
  82 {
  83         return __add_three(&groot, &glock);
  84 }
  85
  86 SEC("tc")
  87 long rbtree_add_nodes_nested(void *ctx)
  88 {
  89         return __add_three(&groot_nested.inner.root, &groot_nested.inner.glock);
  90 }
  91
  92 SEC("tc")
  93 long rbtree_add_and_remove(void *ctx)
  94 {
  95         struct bpf_rb_node *res = NULL;
  96         struct node_data *n, *m = NULL;
  97
  98         n = bpf_obj_new(typeof(*n));
  99         if (!n)
 100                 goto err_out;
 101         n->key = 5;
 102
 103         m = bpf_obj_new(typeof(*m));
 104         if (!m)
 105                 goto err_out;
 106         m->key = 3;
 107
 108         bpf_spin_lock(&glock);
 109         bpf_rbtree_add(&groot, &n->node, less);
 110         bpf_rbtree_add(&groot, &m->node, less);
 111         res = bpf_rbtree_remove(&groot, &n->node);
 112         bpf_spin_unlock(&glock);
 113
 114         if (!res)
 115                 return 1;
 116
 117         n = container_of(res, struct node_data, node);
 118         removed_key = n->key;
 119         bpf_obj_drop(n);
 120
 121         return 0;
 122 err_out:
 123         if (n)
 124                 bpf_obj_drop(n);
 125         if (m)
 126                 bpf_obj_drop(m);
 127         return 1;
 128 }
 129
 130 SEC("tc")
 131 long rbtree_add_and_remove_array(void *ctx)
 132 {
 133         struct bpf_rb_node *res1 = NULL, *res2 = NULL, *res3 = NULL;
 134         struct node_data *nodes[3][2] = {{NULL, NULL}, {NULL, NULL}, {NULL, NULL}};
 135         struct node_data *n;
 136         long k1 = -1, k2 = -1, k3 = -1;
 137         int i, j;
 138
 139         for (i = 0; i < 3; i++) {
 140                 for (j = 0; j < 2; j++) {
 141                         nodes[i][j] = bpf_obj_new(typeof(*nodes[i][j]));
 142                         if (!nodes[i][j])
 143                                 goto err_out;
 144                         nodes[i][j]->key = i * 2 + j;
 145                 }
 146         }
 147
 148         bpf_spin_lock(&glock);
 149         for (i = 0; i < 2; i++)
 150                 for (j = 0; j < 2; j++)
 151                         bpf_rbtree_add(&groot_array[i], &nodes[i][j]->node, less);
 152         for (j = 0; j < 2; j++)
 153                 bpf_rbtree_add(&groot_array_one[0], &nodes[2][j]->node, less);
 154         res1 = bpf_rbtree_remove(&groot_array[0], &nodes[0][0]->node);
 155         res2 = bpf_rbtree_remove(&groot_array[1], &nodes[1][0]->node);
 156         res3 = bpf_rbtree_remove(&groot_array_one[0], &nodes[2][0]->node);
 157         bpf_spin_unlock(&glock);
 158
 159         if (res1) {
 160                 n = container_of(res1, struct node_data, node);
 161                 k1 = n->key;
 162                 bpf_obj_drop(n);
 163         }
 164         if (res2) {
 165                 n = container_of(res2, struct node_data, node);
 166                 k2 = n->key;
 167                 bpf_obj_drop(n);
 168         }
 169         if (res3) {
 170                 n = container_of(res3, struct node_data, node);
 171                 k3 = n->key;
 172                 bpf_obj_drop(n);
 173         }
 174         if (k1 != 0 || k2 != 2 || k3 != 4)
 175                 return 2;
 176
 177         return 0;
 178
 179 err_out:
 180         for (i = 0; i < 3; i++) {
 181                 for (j = 0; j < 2; j++) {
 182                         if (nodes[i][j])
 183                                 bpf_obj_drop(nodes[i][j]);
 184                 }
 185         }
 186         return 1;
 187 }
 188
 189 SEC("tc")
 190 long rbtree_first_and_remove(void *ctx)
 191 {
 192         struct bpf_rb_node *res = NULL;
 193         struct node_data *n, *m, *o;
 194
 195         n = bpf_obj_new(typeof(*n));
 196         if (!n)
 197                 return 1;
 198         n->key = 3;
 199         n->data = 4;
 200
 201         m = bpf_obj_new(typeof(*m));
 202         if (!m)
 203                 goto err_out;
 204         m->key = 5;
 205         m->data = 6;
 206
 207         o = bpf_obj_new(typeof(*o));
 208         if (!o)
 209                 goto err_out;
 210         o->key = 1;
 211         o->data = 2;
 212
 213         bpf_spin_lock(&glock);
 214         bpf_rbtree_add(&groot, &n->node, less);
 215         bpf_rbtree_add(&groot, &m->node, less);
 216         bpf_rbtree_add(&groot, &o->node, less);
 217
 218         res = bpf_rbtree_first(&groot);
 219         if (!res) {
 220                 bpf_spin_unlock(&glock);
 221                 return 2;
 222         }
 223
 224         o = container_of(res, struct node_data, node);
 225         first_data[0] = o->data;
 226
 227         res = bpf_rbtree_remove(&groot, &o->node);
 228         bpf_spin_unlock(&glock);
 229
 230         if (!res)
 231                 return 5;
 232
 233         o = container_of(res, struct node_data, node);
 234         removed_key = o->key;
 235         bpf_obj_drop(o);
 236
 237         bpf_spin_lock(&glock);
 238         res = bpf_rbtree_first(&groot);
 239         if (!res) {
 240                 bpf_spin_unlock(&glock);
 241                 return 3;
 242         }
 243
 244         o = container_of(res, struct node_data, node);
 245         first_data[1] = o->data;
 246         bpf_spin_unlock(&glock);
 247
 248         return 0;
 249 err_out:
 250         if (n)
 251                 bpf_obj_drop(n);
 252         if (m)
 253                 bpf_obj_drop(m);
 254         return 1;
 255 }
 256
 257 SEC("tc")
 258 long rbtree_api_release_aliasing(void *ctx)
 259 {
 260         struct node_data *n, *m, *o;
 261         struct bpf_rb_node *res, *res2;
 262
 263         n = bpf_obj_new(typeof(*n));
 264         if (!n)
 265                 return 1;
 266         n->key = 41;
 267         n->data = 42;
 268
 269         bpf_spin_lock(&glock);
 270         bpf_rbtree_add(&groot, &n->node, less);
 271         bpf_spin_unlock(&glock);
 272
 273         bpf_spin_lock(&glock);
 274
 275         /* m and o point to the same node,
 276          * but verifier doesn't know this
 277          */
 278         res = bpf_rbtree_first(&groot);
 279         if (!res)
 280                 goto err_out;
 281         o = container_of(res, struct node_data, node);
 282
 283         res = bpf_rbtree_first(&groot);
 284         if (!res)
 285                 goto err_out;
 286         m = container_of(res, struct node_data, node);
 287
 288         res = bpf_rbtree_remove(&groot, &m->node);
 289         /* Retval of previous remove returns an owning reference to m,
 290          * which is the same node non-owning ref o is pointing at.
 291          * We can safely try to remove o as the second rbtree_remove will
 292          * return NULL since the node isn't in a tree.
 293          *
 294          * Previously we relied on the verifier type system + rbtree_remove
 295          * invalidating non-owning refs to ensure that rbtree_remove couldn't
 296          * fail, but now rbtree_remove does runtime checking so we no longer
 297          * invalidate non-owning refs after remove.
 298          */
 299         res2 = bpf_rbtree_remove(&groot, &o->node);
 300
 301         bpf_spin_unlock(&glock);
 302
 303         if (res) {
 304                 o = container_of(res, struct node_data, node);
 305                 first_data[0] = o->data;
 306                 bpf_obj_drop(o);
 307         }
 308         if (res2) {
 309                 /* The second remove fails, so res2 is null and this doesn't
 310                  * execute
 311                  */
 312                 m = container_of(res2, struct node_data, node);
 313                 first_data[1] = m->data;
 314                 bpf_obj_drop(m);
 315         }
 316         return 0;
 317
 318 err_out:
 319         bpf_spin_unlock(&glock);
 320         return 1;
 321 }
 322
 323 char _license[] SEC("license") = "GPL";