1112158Sdas// SPDX-License-Identifier: GPL-2.0
2112158Sdas#include <vmlinux.h>
3112158Sdas#include <bpf/bpf_tracing.h>
4112158Sdas#include <bpf/bpf_helpers.h>
5112158Sdas#include <bpf/bpf_core_read.h>
6112158Sdas#include "bpf_experimental.h"
7112158Sdas#include "bpf_misc.h"
8112158Sdas
9112158Sdasstruct node_data {
10112158Sdas	long key;
11112158Sdas	long data;
12112158Sdas	struct bpf_rb_node node;
13112158Sdas};
14112158Sdas
15112158Sdas#define private(name) SEC(".data." #name) __hidden __attribute__((aligned(8)))
16112158Sdasprivate(A) struct bpf_spin_lock glock;
17112158Sdasprivate(A) struct bpf_rb_root groot __contains(node_data, node);
18112158Sdasprivate(A) struct bpf_rb_root groot2 __contains(node_data, node);
19112158Sdas
20112158Sdasstatic bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
21112158Sdas{
22112158Sdas	struct node_data *node_a;
23112158Sdas	struct node_data *node_b;
24112158Sdas
25112158Sdas	node_a = container_of(a, struct node_data, node);
26112158Sdas	node_b = container_of(b, struct node_data, node);
27112158Sdas
28112158Sdas	return node_a->key < node_b->key;
29165743Sdas}
30165743Sdas
31112158SdasSEC("?tc")
32112158Sdas__failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
33112158Sdaslong rbtree_api_nolock_add(void *ctx)
34112158Sdas{
35112158Sdas	struct node_data *n;
36112158Sdas
37112158Sdas	n = bpf_obj_new(typeof(*n));
38112158Sdas	if (!n)
39112158Sdas		return 1;
40112158Sdas
41112158Sdas	bpf_rbtree_add(&groot, &n->node, less);
42112158Sdas	return 0;
43112158Sdas}
44112158Sdas
45112158SdasSEC("?tc")
46112158Sdas__failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
47112158Sdaslong rbtree_api_nolock_remove(void *ctx)
48112158Sdas{
49112158Sdas	struct node_data *n;
50112158Sdas
51112158Sdas	n = bpf_obj_new(typeof(*n));
52112158Sdas	if (!n)
53112158Sdas		return 1;
54112158Sdas
55112158Sdas	bpf_spin_lock(&glock);
56112158Sdas	bpf_rbtree_add(&groot, &n->node, less);
57112158Sdas	bpf_spin_unlock(&glock);
58112158Sdas
59112158Sdas	bpf_rbtree_remove(&groot, &n->node);
60112158Sdas	return 0;
61112158Sdas}
62112158Sdas
63112158SdasSEC("?tc")
64112158Sdas__failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
65112158Sdaslong rbtree_api_nolock_first(void *ctx)
66112158Sdas{
67112158Sdas	bpf_rbtree_first(&groot);
68112158Sdas	return 0;
69112158Sdas}
70112158Sdas
71112158SdasSEC("?tc")
72112158Sdas__failure __msg("rbtree_remove node input must be non-owning ref")
73112158Sdaslong rbtree_api_remove_unadded_node(void *ctx)
74112158Sdas{
75112158Sdas	struct node_data *n, *m;
76112158Sdas	struct bpf_rb_node *res;
77112158Sdas
78112158Sdas	n = bpf_obj_new(typeof(*n));
79112158Sdas	if (!n)
80112158Sdas		return 1;
81112158Sdas
82112158Sdas	m = bpf_obj_new(typeof(*m));
83112158Sdas	if (!m) {
84112158Sdas		bpf_obj_drop(n);
85112158Sdas		return 1;
86112158Sdas	}
87112158Sdas
88112158Sdas	bpf_spin_lock(&glock);
89112158Sdas	bpf_rbtree_add(&groot, &n->node, less);
90112158Sdas
91112158Sdas	/* This remove should pass verifier */
92182709Sdas	res = bpf_rbtree_remove(&groot, &n->node);
93112158Sdas	n = container_of(res, struct node_data, node);
94112158Sdas
95112158Sdas	/* This remove shouldn't, m isn't in an rbtree */
96112158Sdas	res = bpf_rbtree_remove(&groot, &m->node);
97112158Sdas	m = container_of(res, struct node_data, node);
98112158Sdas	bpf_spin_unlock(&glock);
99112158Sdas
100112158Sdas	if (n)
101112158Sdas		bpf_obj_drop(n);
102112158Sdas	if (m)
103112158Sdas		bpf_obj_drop(m);
104112158Sdas	return 0;
105112158Sdas}
106112158Sdas
107112158SdasSEC("?tc")
108112158Sdas__failure __msg("Unreleased reference id=3 alloc_insn=10")
109112158Sdaslong rbtree_api_remove_no_drop(void *ctx)
110112158Sdas{
111112158Sdas	struct bpf_rb_node *res;
112112158Sdas	struct node_data *n;
113112158Sdas
114112158Sdas	bpf_spin_lock(&glock);
115112158Sdas	res = bpf_rbtree_first(&groot);
116112158Sdas	if (!res)
117112158Sdas		goto unlock_err;
118112158Sdas
119112158Sdas	res = bpf_rbtree_remove(&groot, res);
120112158Sdas
121112158Sdas	if (res) {
122112158Sdas		n = container_of(res, struct node_data, node);
123112158Sdas		__sink(n);
124112158Sdas	}
125112158Sdas	bpf_spin_unlock(&glock);
126112158Sdas
127112158Sdas	/* if (res) { bpf_obj_drop(n); } is missing here */
128112158Sdas	return 0;
129112158Sdas
130112158Sdasunlock_err:
131112158Sdas	bpf_spin_unlock(&glock);
132112158Sdas	return 1;
133112158Sdas}
134112158Sdas
135112158SdasSEC("?tc")
136112158Sdas__failure __msg("arg#1 expected pointer to allocated object")
137112158Sdaslong rbtree_api_add_to_multiple_trees(void *ctx)
138112158Sdas{
139112158Sdas	struct node_data *n;
140112158Sdas
141112158Sdas	n = bpf_obj_new(typeof(*n));
142112158Sdas	if (!n)
143112158Sdas		return 1;
144112158Sdas
145112158Sdas	bpf_spin_lock(&glock);
146112158Sdas	bpf_rbtree_add(&groot, &n->node, less);
147112158Sdas
148112158Sdas	/* This add should fail since n already in groot's tree */
149112158Sdas	bpf_rbtree_add(&groot2, &n->node, less);
150112158Sdas	bpf_spin_unlock(&glock);
151112158Sdas	return 0;
152112158Sdas}
153112158Sdas
154112158SdasSEC("?tc")
155112158Sdas__failure __msg("dereference of modified ptr_or_null_ ptr R2 off=16 disallowed")
156112158Sdaslong rbtree_api_use_unchecked_remove_retval(void *ctx)
157112158Sdas{
158112158Sdas	struct bpf_rb_node *res;
159112158Sdas
160112158Sdas	bpf_spin_lock(&glock);
161112158Sdas
162112158Sdas	res = bpf_rbtree_first(&groot);
163112158Sdas	if (!res)
164112158Sdas		goto err_out;
165112158Sdas	res = bpf_rbtree_remove(&groot, res);
166112158Sdas
167112158Sdas	bpf_spin_unlock(&glock);
168112158Sdas
169112158Sdas	bpf_spin_lock(&glock);
170112158Sdas	/* Must check res for NULL before using in rbtree_add below */
171112158Sdas	bpf_rbtree_add(&groot, res, less);
172112158Sdas	bpf_spin_unlock(&glock);
173112158Sdas	return 0;
174112158Sdas
175219557Sdaserr_out:
176112158Sdas	bpf_spin_unlock(&glock);
177219557Sdas	return 1;
178112158Sdas}
179112158Sdas
180112158SdasSEC("?tc")
181112158Sdas__failure __msg("rbtree_remove node input must be non-owning ref")
182112158Sdaslong rbtree_api_add_release_unlock_escape(void *ctx)
183112158Sdas{
184112158Sdas	struct node_data *n;
185219557Sdas
186112158Sdas	n = bpf_obj_new(typeof(*n));
187112158Sdas	if (!n)
188112158Sdas		return 1;
189112158Sdas
190112158Sdas	bpf_spin_lock(&glock);
191112158Sdas	bpf_rbtree_add(&groot, &n->node, less);
192112158Sdas	bpf_spin_unlock(&glock);
193112158Sdas
194112158Sdas	bpf_spin_lock(&glock);
195112158Sdas	/* After add() in previous critical section, n should be
196112158Sdas	 * release_on_unlock and released after previous spin_unlock,
197112158Sdas	 * so should not be possible to use it here
198112158Sdas	 */
199112158Sdas	bpf_rbtree_remove(&groot, &n->node);
200112158Sdas	bpf_spin_unlock(&glock);
201112158Sdas	return 0;
202112158Sdas}
203112158Sdas
204112158SdasSEC("?tc")
205112158Sdas__failure __msg("rbtree_remove node input must be non-owning ref")
206112158Sdaslong rbtree_api_first_release_unlock_escape(void *ctx)
207112158Sdas{
208182709Sdas	struct bpf_rb_node *res;
209112158Sdas	struct node_data *n;
210182709Sdas
211112158Sdas	bpf_spin_lock(&glock);
212112158Sdas	res = bpf_rbtree_first(&groot);
213112158Sdas	if (!res) {
214112158Sdas		bpf_spin_unlock(&glock);
215112158Sdas		return 1;
216112158Sdas	}
217112158Sdas	n = container_of(res, struct node_data, node);
218112158Sdas	bpf_spin_unlock(&glock);
219112158Sdas
220112158Sdas	bpf_spin_lock(&glock);
221112158Sdas	/* After first() in previous critical section, n should be
222112158Sdas	 * release_on_unlock and released after previous spin_unlock,
223112158Sdas	 * so should not be possible to use it here
224112158Sdas	 */
225112158Sdas	bpf_rbtree_remove(&groot, &n->node);
226112158Sdas	bpf_spin_unlock(&glock);
227112158Sdas	return 0;
228112158Sdas}
229112158Sdas
230112158Sdasstatic bool less__bad_fn_call_add(struct bpf_rb_node *a, const struct bpf_rb_node *b)
231112158Sdas{
232112158Sdas	struct node_data *node_a;
233112158Sdas	struct node_data *node_b;
234112158Sdas
235112158Sdas	node_a = container_of(a, struct node_data, node);
236112158Sdas	node_b = container_of(b, struct node_data, node);
237112158Sdas	bpf_rbtree_add(&groot, &node_a->node, less);
238112158Sdas
239165743Sdas	return node_a->key < node_b->key;
240112158Sdas}
241112158Sdas
242112158Sdasstatic bool less__bad_fn_call_remove(struct bpf_rb_node *a, const struct bpf_rb_node *b)
243112158Sdas{
244112158Sdas	struct node_data *node_a;
245112158Sdas	struct node_data *node_b;
246112158Sdas
247112158Sdas	node_a = container_of(a, struct node_data, node);
248112158Sdas	node_b = container_of(b, struct node_data, node);
249112158Sdas	bpf_rbtree_remove(&groot, &node_a->node);
250112158Sdas
251112158Sdas	return node_a->key < node_b->key;
252112158Sdas}
253112158Sdas
254112158Sdasstatic bool less__bad_fn_call_first_unlock_after(struct bpf_rb_node *a, const struct bpf_rb_node *b)
255112158Sdas{
256112158Sdas	struct node_data *node_a;
257112158Sdas	struct node_data *node_b;
258112158Sdas
259112158Sdas	node_a = container_of(a, struct node_data, node);
260112158Sdas	node_b = container_of(b, struct node_data, node);
261112158Sdas	bpf_rbtree_first(&groot);
262112158Sdas	bpf_spin_unlock(&glock);
263112158Sdas
264112158Sdas	return node_a->key < node_b->key;
265112158Sdas}
266112158Sdas
267112158Sdasstatic __always_inline
268112158Sdaslong add_with_cb(bool (cb)(struct bpf_rb_node *a, const struct bpf_rb_node *b))
269112158Sdas{
270112158Sdas	struct node_data *n;
271112158Sdas
272112158Sdas	n = bpf_obj_new(typeof(*n));
273112158Sdas	if (!n)
274112158Sdas		return 1;
275112158Sdas
276112158Sdas	bpf_spin_lock(&glock);
277112158Sdas	bpf_rbtree_add(&groot, &n->node, cb);
278112158Sdas	bpf_spin_unlock(&glock);
279112158Sdas	return 0;
280112158Sdas}
281112158Sdas
282112158SdasSEC("?tc")
283112158Sdas__failure __msg("arg#1 expected pointer to allocated object")
284112158Sdaslong rbtree_api_add_bad_cb_bad_fn_call_add(void *ctx)
285112158Sdas{
286112158Sdas	return add_with_cb(less__bad_fn_call_add);
287112158Sdas}
288112158Sdas
289112158SdasSEC("?tc")
290112158Sdas__failure __msg("rbtree_remove not allowed in rbtree cb")
291112158Sdaslong rbtree_api_add_bad_cb_bad_fn_call_remove(void *ctx)
292112158Sdas{
293112158Sdas	return add_with_cb(less__bad_fn_call_remove);
294219557Sdas}
295112158Sdas
296219557SdasSEC("?tc")
297112158Sdas__failure __msg("can't spin_{lock,unlock} in rbtree cb")
298112158Sdaslong rbtree_api_add_bad_cb_bad_fn_call_first_unlock_after(void *ctx)
299112158Sdas{
300112158Sdas	return add_with_cb(less__bad_fn_call_first_unlock_after);
301112158Sdas}
302112158Sdas
303112158Sdaschar _license[] SEC("license") = "GPL";
304112158Sdas