1215249f6SDave Marchevsky // SPDX-License-Identifier: GPL-2.0
2215249f6SDave Marchevsky #include <vmlinux.h>
3215249f6SDave Marchevsky #include <bpf/bpf_tracing.h>
4215249f6SDave Marchevsky #include <bpf/bpf_helpers.h>
5215249f6SDave Marchevsky #include <bpf/bpf_core_read.h>
6215249f6SDave Marchevsky #include "bpf_experimental.h"
7215249f6SDave Marchevsky #include "bpf_misc.h"
8215249f6SDave Marchevsky 
9215249f6SDave Marchevsky struct node_data {
10215249f6SDave Marchevsky 	long key;
11215249f6SDave Marchevsky 	long data;
12215249f6SDave Marchevsky 	struct bpf_rb_node node;
13215249f6SDave Marchevsky };
14215249f6SDave Marchevsky 
15215249f6SDave Marchevsky #define private(name) SEC(".data." #name) __hidden __attribute__((aligned(8)))
16215249f6SDave Marchevsky private(A) struct bpf_spin_lock glock;
17215249f6SDave Marchevsky private(A) struct bpf_rb_root groot __contains(node_data, node);
18215249f6SDave Marchevsky private(A) struct bpf_rb_root groot2 __contains(node_data, node);
19215249f6SDave Marchevsky 
less(struct bpf_rb_node * a,const struct bpf_rb_node * b)20215249f6SDave Marchevsky static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
21215249f6SDave Marchevsky {
22215249f6SDave Marchevsky 	struct node_data *node_a;
23215249f6SDave Marchevsky 	struct node_data *node_b;
24215249f6SDave Marchevsky 
25215249f6SDave Marchevsky 	node_a = container_of(a, struct node_data, node);
26215249f6SDave Marchevsky 	node_b = container_of(b, struct node_data, node);
27215249f6SDave Marchevsky 
28215249f6SDave Marchevsky 	return node_a->key < node_b->key;
29215249f6SDave Marchevsky }
30215249f6SDave Marchevsky 
31215249f6SDave Marchevsky SEC("?tc")
32215249f6SDave Marchevsky __failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
rbtree_api_nolock_add(void * ctx)33215249f6SDave Marchevsky long rbtree_api_nolock_add(void *ctx)
34215249f6SDave Marchevsky {
35215249f6SDave Marchevsky 	struct node_data *n;
36215249f6SDave Marchevsky 
37215249f6SDave Marchevsky 	n = bpf_obj_new(typeof(*n));
38215249f6SDave Marchevsky 	if (!n)
39215249f6SDave Marchevsky 		return 1;
40215249f6SDave Marchevsky 
41215249f6SDave Marchevsky 	bpf_rbtree_add(&groot, &n->node, less);
42215249f6SDave Marchevsky 	return 0;
43215249f6SDave Marchevsky }
44215249f6SDave Marchevsky 
45215249f6SDave Marchevsky SEC("?tc")
46215249f6SDave Marchevsky __failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
rbtree_api_nolock_remove(void * ctx)47215249f6SDave Marchevsky long rbtree_api_nolock_remove(void *ctx)
48215249f6SDave Marchevsky {
49215249f6SDave Marchevsky 	struct node_data *n;
50215249f6SDave Marchevsky 
51215249f6SDave Marchevsky 	n = bpf_obj_new(typeof(*n));
52215249f6SDave Marchevsky 	if (!n)
53215249f6SDave Marchevsky 		return 1;
54215249f6SDave Marchevsky 
55215249f6SDave Marchevsky 	bpf_spin_lock(&glock);
56215249f6SDave Marchevsky 	bpf_rbtree_add(&groot, &n->node, less);
57215249f6SDave Marchevsky 	bpf_spin_unlock(&glock);
58215249f6SDave Marchevsky 
59215249f6SDave Marchevsky 	bpf_rbtree_remove(&groot, &n->node);
60215249f6SDave Marchevsky 	return 0;
61215249f6SDave Marchevsky }
62215249f6SDave Marchevsky 
63215249f6SDave Marchevsky SEC("?tc")
64215249f6SDave Marchevsky __failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
rbtree_api_nolock_first(void * ctx)65215249f6SDave Marchevsky long rbtree_api_nolock_first(void *ctx)
66215249f6SDave Marchevsky {
67215249f6SDave Marchevsky 	bpf_rbtree_first(&groot);
68215249f6SDave Marchevsky 	return 0;
69215249f6SDave Marchevsky }
70215249f6SDave Marchevsky 
71215249f6SDave Marchevsky SEC("?tc")
72215249f6SDave Marchevsky __failure __msg("rbtree_remove node input must be non-owning ref")
rbtree_api_remove_unadded_node(void * ctx)73215249f6SDave Marchevsky long rbtree_api_remove_unadded_node(void *ctx)
74215249f6SDave Marchevsky {
75215249f6SDave Marchevsky 	struct node_data *n, *m;
76215249f6SDave Marchevsky 	struct bpf_rb_node *res;
77215249f6SDave Marchevsky 
78215249f6SDave Marchevsky 	n = bpf_obj_new(typeof(*n));
79215249f6SDave Marchevsky 	if (!n)
80215249f6SDave Marchevsky 		return 1;
81215249f6SDave Marchevsky 
82215249f6SDave Marchevsky 	m = bpf_obj_new(typeof(*m));
83215249f6SDave Marchevsky 	if (!m) {
84215249f6SDave Marchevsky 		bpf_obj_drop(n);
85215249f6SDave Marchevsky 		return 1;
86215249f6SDave Marchevsky 	}
87215249f6SDave Marchevsky 
88215249f6SDave Marchevsky 	bpf_spin_lock(&glock);
89215249f6SDave Marchevsky 	bpf_rbtree_add(&groot, &n->node, less);
90215249f6SDave Marchevsky 
91215249f6SDave Marchevsky 	/* This remove should pass verifier */
92215249f6SDave Marchevsky 	res = bpf_rbtree_remove(&groot, &n->node);
93215249f6SDave Marchevsky 	n = container_of(res, struct node_data, node);
94215249f6SDave Marchevsky 
95215249f6SDave Marchevsky 	/* This remove shouldn't, m isn't in an rbtree */
96215249f6SDave Marchevsky 	res = bpf_rbtree_remove(&groot, &m->node);
97215249f6SDave Marchevsky 	m = container_of(res, struct node_data, node);
98215249f6SDave Marchevsky 	bpf_spin_unlock(&glock);
99215249f6SDave Marchevsky 
100215249f6SDave Marchevsky 	if (n)
101215249f6SDave Marchevsky 		bpf_obj_drop(n);
102215249f6SDave Marchevsky 	if (m)
103215249f6SDave Marchevsky 		bpf_obj_drop(m);
104215249f6SDave Marchevsky 	return 0;
105215249f6SDave Marchevsky }
106215249f6SDave Marchevsky 
107215249f6SDave Marchevsky SEC("?tc")
108*404ad75aSDave Marchevsky __failure __msg("Unreleased reference id=3 alloc_insn=10")
rbtree_api_remove_no_drop(void * ctx)109215249f6SDave Marchevsky long rbtree_api_remove_no_drop(void *ctx)
110215249f6SDave Marchevsky {
111215249f6SDave Marchevsky 	struct bpf_rb_node *res;
112215249f6SDave Marchevsky 	struct node_data *n;
113215249f6SDave Marchevsky 
114215249f6SDave Marchevsky 	bpf_spin_lock(&glock);
115215249f6SDave Marchevsky 	res = bpf_rbtree_first(&groot);
116215249f6SDave Marchevsky 	if (!res)
117215249f6SDave Marchevsky 		goto unlock_err;
118215249f6SDave Marchevsky 
119215249f6SDave Marchevsky 	res = bpf_rbtree_remove(&groot, res);
120215249f6SDave Marchevsky 
121*404ad75aSDave Marchevsky 	if (res) {
122215249f6SDave Marchevsky 		n = container_of(res, struct node_data, node);
123c8ed6685SAndrii Nakryiko 		__sink(n);
124*404ad75aSDave Marchevsky 	}
125215249f6SDave Marchevsky 	bpf_spin_unlock(&glock);
126215249f6SDave Marchevsky 
127*404ad75aSDave Marchevsky 	/* if (res) { bpf_obj_drop(n); } is missing here */
128215249f6SDave Marchevsky 	return 0;
129215249f6SDave Marchevsky 
130215249f6SDave Marchevsky unlock_err:
131215249f6SDave Marchevsky 	bpf_spin_unlock(&glock);
132215249f6SDave Marchevsky 	return 1;
133215249f6SDave Marchevsky }
134215249f6SDave Marchevsky 
135215249f6SDave Marchevsky SEC("?tc")
136215249f6SDave Marchevsky __failure __msg("arg#1 expected pointer to allocated object")
rbtree_api_add_to_multiple_trees(void * ctx)137215249f6SDave Marchevsky long rbtree_api_add_to_multiple_trees(void *ctx)
138215249f6SDave Marchevsky {
139215249f6SDave Marchevsky 	struct node_data *n;
140215249f6SDave Marchevsky 
141215249f6SDave Marchevsky 	n = bpf_obj_new(typeof(*n));
142215249f6SDave Marchevsky 	if (!n)
143215249f6SDave Marchevsky 		return 1;
144215249f6SDave Marchevsky 
145215249f6SDave Marchevsky 	bpf_spin_lock(&glock);
146215249f6SDave Marchevsky 	bpf_rbtree_add(&groot, &n->node, less);
147215249f6SDave Marchevsky 
148215249f6SDave Marchevsky 	/* This add should fail since n already in groot's tree */
149215249f6SDave Marchevsky 	bpf_rbtree_add(&groot2, &n->node, less);
150215249f6SDave Marchevsky 	bpf_spin_unlock(&glock);
151215249f6SDave Marchevsky 	return 0;
152215249f6SDave Marchevsky }
153215249f6SDave Marchevsky 
154215249f6SDave Marchevsky SEC("?tc")
155*404ad75aSDave Marchevsky __failure __msg("dereference of modified ptr_or_null_ ptr R2 off=16 disallowed")
rbtree_api_use_unchecked_remove_retval(void * ctx)156*404ad75aSDave Marchevsky long rbtree_api_use_unchecked_remove_retval(void *ctx)
157*404ad75aSDave Marchevsky {
158*404ad75aSDave Marchevsky 	struct bpf_rb_node *res;
159*404ad75aSDave Marchevsky 
160*404ad75aSDave Marchevsky 	bpf_spin_lock(&glock);
161*404ad75aSDave Marchevsky 
162*404ad75aSDave Marchevsky 	res = bpf_rbtree_first(&groot);
163*404ad75aSDave Marchevsky 	if (!res)
164*404ad75aSDave Marchevsky 		goto err_out;
165*404ad75aSDave Marchevsky 	res = bpf_rbtree_remove(&groot, res);
166*404ad75aSDave Marchevsky 
167*404ad75aSDave Marchevsky 	bpf_spin_unlock(&glock);
168*404ad75aSDave Marchevsky 
169*404ad75aSDave Marchevsky 	bpf_spin_lock(&glock);
170*404ad75aSDave Marchevsky 	/* Must check res for NULL before using in rbtree_add below */
171*404ad75aSDave Marchevsky 	bpf_rbtree_add(&groot, res, less);
172*404ad75aSDave Marchevsky 	bpf_spin_unlock(&glock);
173*404ad75aSDave Marchevsky 	return 0;
174*404ad75aSDave Marchevsky 
175*404ad75aSDave Marchevsky err_out:
176*404ad75aSDave Marchevsky 	bpf_spin_unlock(&glock);
177*404ad75aSDave Marchevsky 	return 1;
178*404ad75aSDave Marchevsky }
179*404ad75aSDave Marchevsky 
180*404ad75aSDave Marchevsky SEC("?tc")
181215249f6SDave Marchevsky __failure __msg("rbtree_remove node input must be non-owning ref")
rbtree_api_add_release_unlock_escape(void * ctx)182215249f6SDave Marchevsky long rbtree_api_add_release_unlock_escape(void *ctx)
183215249f6SDave Marchevsky {
184215249f6SDave Marchevsky 	struct node_data *n;
185215249f6SDave Marchevsky 
186215249f6SDave Marchevsky 	n = bpf_obj_new(typeof(*n));
187215249f6SDave Marchevsky 	if (!n)
188215249f6SDave Marchevsky 		return 1;
189215249f6SDave Marchevsky 
190215249f6SDave Marchevsky 	bpf_spin_lock(&glock);
191215249f6SDave Marchevsky 	bpf_rbtree_add(&groot, &n->node, less);
192215249f6SDave Marchevsky 	bpf_spin_unlock(&glock);
193215249f6SDave Marchevsky 
194215249f6SDave Marchevsky 	bpf_spin_lock(&glock);
195215249f6SDave Marchevsky 	/* After add() in previous critical section, n should be
196215249f6SDave Marchevsky 	 * release_on_unlock and released after previous spin_unlock,
197215249f6SDave Marchevsky 	 * so should not be possible to use it here
198215249f6SDave Marchevsky 	 */
199215249f6SDave Marchevsky 	bpf_rbtree_remove(&groot, &n->node);
200215249f6SDave Marchevsky 	bpf_spin_unlock(&glock);
201215249f6SDave Marchevsky 	return 0;
202215249f6SDave Marchevsky }
203215249f6SDave Marchevsky 
204215249f6SDave Marchevsky SEC("?tc")
205215249f6SDave Marchevsky __failure __msg("rbtree_remove node input must be non-owning ref")
rbtree_api_first_release_unlock_escape(void * ctx)206215249f6SDave Marchevsky long rbtree_api_first_release_unlock_escape(void *ctx)
207215249f6SDave Marchevsky {
208215249f6SDave Marchevsky 	struct bpf_rb_node *res;
209215249f6SDave Marchevsky 	struct node_data *n;
210215249f6SDave Marchevsky 
211215249f6SDave Marchevsky 	bpf_spin_lock(&glock);
212215249f6SDave Marchevsky 	res = bpf_rbtree_first(&groot);
213ec97a76fSDave Marchevsky 	if (!res) {
214ec97a76fSDave Marchevsky 		bpf_spin_unlock(&glock);
215ec97a76fSDave Marchevsky 		return 1;
216ec97a76fSDave Marchevsky 	}
217215249f6SDave Marchevsky 	n = container_of(res, struct node_data, node);
218215249f6SDave Marchevsky 	bpf_spin_unlock(&glock);
219215249f6SDave Marchevsky 
220215249f6SDave Marchevsky 	bpf_spin_lock(&glock);
221215249f6SDave Marchevsky 	/* After first() in previous critical section, n should be
222215249f6SDave Marchevsky 	 * release_on_unlock and released after previous spin_unlock,
223215249f6SDave Marchevsky 	 * so should not be possible to use it here
224215249f6SDave Marchevsky 	 */
225215249f6SDave Marchevsky 	bpf_rbtree_remove(&groot, &n->node);
226215249f6SDave Marchevsky 	bpf_spin_unlock(&glock);
227215249f6SDave Marchevsky 	return 0;
228215249f6SDave Marchevsky }
229215249f6SDave Marchevsky 
less__bad_fn_call_add(struct bpf_rb_node * a,const struct bpf_rb_node * b)230215249f6SDave Marchevsky static bool less__bad_fn_call_add(struct bpf_rb_node *a, const struct bpf_rb_node *b)
231215249f6SDave Marchevsky {
232215249f6SDave Marchevsky 	struct node_data *node_a;
233215249f6SDave Marchevsky 	struct node_data *node_b;
234215249f6SDave Marchevsky 
235215249f6SDave Marchevsky 	node_a = container_of(a, struct node_data, node);
236215249f6SDave Marchevsky 	node_b = container_of(b, struct node_data, node);
237215249f6SDave Marchevsky 	bpf_rbtree_add(&groot, &node_a->node, less);
238215249f6SDave Marchevsky 
239215249f6SDave Marchevsky 	return node_a->key < node_b->key;
240215249f6SDave Marchevsky }
241215249f6SDave Marchevsky 
less__bad_fn_call_remove(struct bpf_rb_node * a,const struct bpf_rb_node * b)242215249f6SDave Marchevsky static bool less__bad_fn_call_remove(struct bpf_rb_node *a, const struct bpf_rb_node *b)
243215249f6SDave Marchevsky {
244215249f6SDave Marchevsky 	struct node_data *node_a;
245215249f6SDave Marchevsky 	struct node_data *node_b;
246215249f6SDave Marchevsky 
247215249f6SDave Marchevsky 	node_a = container_of(a, struct node_data, node);
248215249f6SDave Marchevsky 	node_b = container_of(b, struct node_data, node);
249215249f6SDave Marchevsky 	bpf_rbtree_remove(&groot, &node_a->node);
250215249f6SDave Marchevsky 
251215249f6SDave Marchevsky 	return node_a->key < node_b->key;
252215249f6SDave Marchevsky }
253215249f6SDave Marchevsky 
less__bad_fn_call_first_unlock_after(struct bpf_rb_node * a,const struct bpf_rb_node * b)254215249f6SDave Marchevsky static bool less__bad_fn_call_first_unlock_after(struct bpf_rb_node *a, const struct bpf_rb_node *b)
255215249f6SDave Marchevsky {
256215249f6SDave Marchevsky 	struct node_data *node_a;
257215249f6SDave Marchevsky 	struct node_data *node_b;
258215249f6SDave Marchevsky 
259215249f6SDave Marchevsky 	node_a = container_of(a, struct node_data, node);
260215249f6SDave Marchevsky 	node_b = container_of(b, struct node_data, node);
261215249f6SDave Marchevsky 	bpf_rbtree_first(&groot);
262215249f6SDave Marchevsky 	bpf_spin_unlock(&glock);
263215249f6SDave Marchevsky 
264215249f6SDave Marchevsky 	return node_a->key < node_b->key;
265215249f6SDave Marchevsky }
266215249f6SDave Marchevsky 
267215249f6SDave Marchevsky static __always_inline
add_with_cb(bool (cb)(struct bpf_rb_node * a,const struct bpf_rb_node * b))268215249f6SDave Marchevsky long add_with_cb(bool (cb)(struct bpf_rb_node *a, const struct bpf_rb_node *b))
269215249f6SDave Marchevsky {
270215249f6SDave Marchevsky 	struct node_data *n;
271215249f6SDave Marchevsky 
272215249f6SDave Marchevsky 	n = bpf_obj_new(typeof(*n));
273215249f6SDave Marchevsky 	if (!n)
274215249f6SDave Marchevsky 		return 1;
275215249f6SDave Marchevsky 
276215249f6SDave Marchevsky 	bpf_spin_lock(&glock);
277215249f6SDave Marchevsky 	bpf_rbtree_add(&groot, &n->node, cb);
278215249f6SDave Marchevsky 	bpf_spin_unlock(&glock);
279215249f6SDave Marchevsky 	return 0;
280215249f6SDave Marchevsky }
281215249f6SDave Marchevsky 
282215249f6SDave Marchevsky SEC("?tc")
283215249f6SDave Marchevsky __failure __msg("arg#1 expected pointer to allocated object")
rbtree_api_add_bad_cb_bad_fn_call_add(void * ctx)284215249f6SDave Marchevsky long rbtree_api_add_bad_cb_bad_fn_call_add(void *ctx)
285215249f6SDave Marchevsky {
286215249f6SDave Marchevsky 	return add_with_cb(less__bad_fn_call_add);
287215249f6SDave Marchevsky }
288215249f6SDave Marchevsky 
289215249f6SDave Marchevsky SEC("?tc")
290215249f6SDave Marchevsky __failure __msg("rbtree_remove not allowed in rbtree cb")
rbtree_api_add_bad_cb_bad_fn_call_remove(void * ctx)291215249f6SDave Marchevsky long rbtree_api_add_bad_cb_bad_fn_call_remove(void *ctx)
292215249f6SDave Marchevsky {
293215249f6SDave Marchevsky 	return add_with_cb(less__bad_fn_call_remove);
294215249f6SDave Marchevsky }
295215249f6SDave Marchevsky 
296215249f6SDave Marchevsky SEC("?tc")
297215249f6SDave Marchevsky __failure __msg("can't spin_{lock,unlock} in rbtree cb")
rbtree_api_add_bad_cb_bad_fn_call_first_unlock_after(void * ctx)298215249f6SDave Marchevsky long rbtree_api_add_bad_cb_bad_fn_call_first_unlock_after(void *ctx)
299215249f6SDave Marchevsky {
300215249f6SDave Marchevsky 	return add_with_cb(less__bad_fn_call_first_unlock_after);
301215249f6SDave Marchevsky }
302215249f6SDave Marchevsky 
303215249f6SDave Marchevsky char _license[] SEC("license") = "GPL";
304