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
 10struct node_data {
 11	long key;
 12	long data;
 13	struct bpf_rb_node node;
 14};
 15
 16long less_callback_ran = -1;
 17long removed_key = -1;
 18long first_data[2] = {-1, -1};
 19
 20#define private(name) SEC(".data." #name) __hidden __attribute__((aligned(8)))
 21private(A) struct bpf_spin_lock glock;
 22private(A) struct bpf_rb_root groot __contains(node_data, node);
 23
 24static 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
 36static 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
 68SEC("tc")
 69long rbtree_add_nodes(void *ctx)
 70{
 71	return __add_three(&groot, &glock);
 72}
 73
 74SEC("tc")
 75long 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;
104err_out:
105	if (n)
106		bpf_obj_drop(n);
107	if (m)
108		bpf_obj_drop(m);
109	return 1;
110}
111
112SEC("tc")
113long 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;
172err_out:
173	if (n)
174		bpf_obj_drop(n);
175	if (m)
176		bpf_obj_drop(m);
177	return 1;
178}
179
180SEC("tc")
181long 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
241err_out:
242	bpf_spin_unlock(&glock);
243	return 1;
244}
245
246char _license[] SEC("license") = "GPL";