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 long less_callback_ran = -1;
  17 long removed_key = -1;
  18 long first_data[2] = {-1, -1};
  19
  20 #define private(name) SEC(".data." #name) __hidden __attribute__((aligned(8)))
  21 private(A) struct bpf_spin_lock glock;
  22 private(A) struct bpf_rb_root groot __contains(node_data, node);
  23
  24 static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
  25 {
  26         struct node_data *node_a;
  27         struct node_data *node_b;
  28
  29         node_a = container_of(a, struct node_data, node);
  30         node_b = container_of(b, struct node_data, node);
  31         less_callback_ran = 1;
  32
  33         return node_a->key < node_b->key;
  34 }
  35
  36 static long __add_three(struct bpf_rb_root *root, struct bpf_spin_lock *lock)
  37 {
  38         struct node_data *n, *m;
  39
  40         n = bpf_obj_new(typeof(*n));
  41         if (!n)
  42                 return 1;
  43         n->key = 5;
  44
  45         m = bpf_obj_new(typeof(*m));
  46         if (!m) {
  47                 bpf_obj_drop(n);
  48                 return 2;
  49         }
  50         m->key = 1;
  51
  52         bpf_spin_lock(&glock);
  53         bpf_rbtree_add(&groot, &n->node, less);
  54         bpf_rbtree_add(&groot, &m->node, less);
  55         bpf_spin_unlock(&glock);
  56
  57         n = bpf_obj_new(typeof(*n));
  58         if (!n)
  59                 return 3;
  60         n->key = 3;
  61
  62         bpf_spin_lock(&glock);
  63         bpf_rbtree_add(&groot, &n->node, less);
  64         bpf_spin_unlock(&glock);
  65         return 0;
  66 }
  67
  68 SEC("tc")
  69 long rbtree_add_nodes(void *ctx)
  70 {
  71         return __add_three(&groot, &glock);
  72 }
  73
  74 SEC("tc")
  75 long rbtree_add_and_remove(void *ctx)
  76 {
  77         struct bpf_rb_node *res = NULL;
  78         struct node_data *n, *m = NULL;
  79
  80         n = bpf_obj_new(typeof(*n));
  81         if (!n)
  82                 goto err_out;
  83         n->key = 5;
  84
  85         m = bpf_obj_new(typeof(*m));
  86         if (!m)
  87                 goto err_out;
  88         m->key = 3;
  89
  90         bpf_spin_lock(&glock);
  91         bpf_rbtree_add(&groot, &n->node, less);
  92         bpf_rbtree_add(&groot, &m->node, less);
  93         res = bpf_rbtree_remove(&groot, &n->node);
  94         bpf_spin_unlock(&glock);
  95
  96         if (!res)
  97                 return 1;
  98
  99         n = container_of(res, struct node_data, node);
 100         removed_key = n->key;
 101         bpf_obj_drop(n);
 102
 103         return 0;
 104 err_out:
 105         if (n)
 106                 bpf_obj_drop(n);
 107         if (m)
 108                 bpf_obj_drop(m);
 109         return 1;
 110 }
 111
 112 SEC("tc")
 113 long rbtree_first_and_remove(void *ctx)
 114 {
 115         struct bpf_rb_node *res = NULL;
 116         struct node_data *n, *m, *o;
 117
 118         n = bpf_obj_new(typeof(*n));
 119         if (!n)
 120                 return 1;
 121         n->key = 3;
 122         n->data = 4;
 123
 124         m = bpf_obj_new(typeof(*m));
 125         if (!m)
 126                 goto err_out;
 127         m->key = 5;
 128         m->data = 6;
 129
 130         o = bpf_obj_new(typeof(*o));
 131         if (!o)
 132                 goto err_out;
 133         o->key = 1;
 134         o->data = 2;
 135
 136         bpf_spin_lock(&glock);
 137         bpf_rbtree_add(&groot, &n->node, less);
 138         bpf_rbtree_add(&groot, &m->node, less);
 139         bpf_rbtree_add(&groot, &o->node, less);
 140
 141         res = bpf_rbtree_first(&groot);
 142         if (!res) {
 143                 bpf_spin_unlock(&glock);
 144                 return 2;
 145         }
 146
 147         o = container_of(res, struct node_data, node);
 148         first_data[0] = o->data;
 149
 150         res = bpf_rbtree_remove(&groot, &o->node);
 151         bpf_spin_unlock(&glock);
 152
 153         if (!res)
 154                 return 5;
 155
 156         o = container_of(res, struct node_data, node);
 157         removed_key = o->key;
 158         bpf_obj_drop(o);
 159
 160         bpf_spin_lock(&glock);
 161         res = bpf_rbtree_first(&groot);
 162         if (!res) {
 163                 bpf_spin_unlock(&glock);
 164                 return 3;
 165         }
 166
 167         o = container_of(res, struct node_data, node);
 168         first_data[1] = o->data;
 169         bpf_spin_unlock(&glock);
 170
 171         return 0;
 172 err_out:
 173         if (n)
 174                 bpf_obj_drop(n);
 175         if (m)
 176                 bpf_obj_drop(m);
 177         return 1;
 178 }
 179
 180 SEC("tc")
 181 long rbtree_api_release_aliasing(void *ctx)
 182 {
 183         struct node_data *n, *m, *o;
 184         struct bpf_rb_node *res, *res2;
 185
 186         n = bpf_obj_new(typeof(*n));
 187         if (!n)
 188                 return 1;
 189         n->key = 41;
 190         n->data = 42;
 191
 192         bpf_spin_lock(&glock);
 193         bpf_rbtree_add(&groot, &n->node, less);
 194         bpf_spin_unlock(&glock);
 195
 196         bpf_spin_lock(&glock);
 197
 198         /* m and o point to the same node,
 199          * but verifier doesn't know this
 200          */
 201         res = bpf_rbtree_first(&groot);
 202         if (!res)
 203                 goto err_out;
 204         o = container_of(res, struct node_data, node);
 205
 206         res = bpf_rbtree_first(&groot);
 207         if (!res)
 208                 goto err_out;
 209         m = container_of(res, struct node_data, node);
 210
 211         res = bpf_rbtree_remove(&groot, &m->node);
 212         /* Retval of previous remove returns an owning reference to m,
 213          * which is the same node non-owning ref o is pointing at.
 214          * We can safely try to remove o as the second rbtree_remove will
 215          * return NULL since the node isn't in a tree.
 216          *
 217          * Previously we relied on the verifier type system + rbtree_remove
 218          * invalidating non-owning refs to ensure that rbtree_remove couldn't
 219          * fail, but now rbtree_remove does runtime checking so we no longer
 220          * invalidate non-owning refs after remove.
 221          */
 222         res2 = bpf_rbtree_remove(&groot, &o->node);
 223
 224         bpf_spin_unlock(&glock);
 225
 226         if (res) {
 227                 o = container_of(res, struct node_data, node);
 228                 first_data[0] = o->data;
 229                 bpf_obj_drop(o);
 230         }
 231         if (res2) {
 232                 /* The second remove fails, so res2 is null and this doesn't
 233                  * execute
 234                  */
 235                 m = container_of(res2, struct node_data, node);
 236                 first_data[1] = m->data;
 237                 bpf_obj_drop(m);
 238         }
 239         return 0;
 240
 241 err_out:
 242         bpf_spin_unlock(&glock);
 243         return 1;
 244 }
 245
 246 char _license[] SEC("license") = "GPL";