1// SPDX-License-Identifier: GPL-2.0
2#include <vmlinux.h>
3#include <bpf/bpf_tracing.h>
4#include <bpf/bpf_helpers.h>
5#include <bpf/bpf_core_read.h>
6#include "bpf_experimental.h"
7#include "bpf_misc.h"
8
9struct node_data {
10	long key;
11	long data;
12	struct bpf_rb_node node;
13};
14
15#define private(name) SEC(".data." #name) __hidden __attribute__((aligned(8)))
16private(A) struct bpf_spin_lock glock;
17private(A) struct bpf_rb_root groot __contains(node_data, node);
18private(A) struct bpf_rb_root groot2 __contains(node_data, node);
19
20static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
21{
22	struct node_data *node_a;
23	struct node_data *node_b;
24
25	node_a = container_of(a, struct node_data, node);
26	node_b = container_of(b, struct node_data, node);
27
28	return node_a->key < node_b->key;
29}
30
31SEC("?tc")
32__failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
33long rbtree_api_nolock_add(void *ctx)
34{
35	struct node_data *n;
36
37	n = bpf_obj_new(typeof(*n));
38	if (!n)
39		return 1;
40
41	bpf_rbtree_add(&groot, &n->node, less);
42	return 0;
43}
44
45SEC("?tc")
46__failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
47long rbtree_api_nolock_remove(void *ctx)
48{
49	struct node_data *n;
50
51	n = bpf_obj_new(typeof(*n));
52	if (!n)
53		return 1;
54
55	bpf_spin_lock(&glock);
56	bpf_rbtree_add(&groot, &n->node, less);
57	bpf_spin_unlock(&glock);
58
59	bpf_rbtree_remove(&groot, &n->node);
60	return 0;
61}
62
63SEC("?tc")
64__failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
65long rbtree_api_nolock_first(void *ctx)
66{
67	bpf_rbtree_first(&groot);
68	return 0;
69}
70
71SEC("?tc")
72__failure __msg("rbtree_remove node input must be non-owning ref")
73long rbtree_api_remove_unadded_node(void *ctx)
74{
75	struct node_data *n, *m;
76	struct bpf_rb_node *res;
77
78	n = bpf_obj_new(typeof(*n));
79	if (!n)
80		return 1;
81
82	m = bpf_obj_new(typeof(*m));
83	if (!m) {
84		bpf_obj_drop(n);
85		return 1;
86	}
87
88	bpf_spin_lock(&glock);
89	bpf_rbtree_add(&groot, &n->node, less);
90
91	/* This remove should pass verifier */
92	res = bpf_rbtree_remove(&groot, &n->node);
93	n = container_of(res, struct node_data, node);
94
95	/* This remove shouldn't, m isn't in an rbtree */
96	res = bpf_rbtree_remove(&groot, &m->node);
97	m = container_of(res, struct node_data, node);
98	bpf_spin_unlock(&glock);
99
100	if (n)
101		bpf_obj_drop(n);
102	if (m)
103		bpf_obj_drop(m);
104	return 0;
105}
106
107SEC("?tc")
108__failure __msg("Unreleased reference id=3 alloc_insn=10")
109long rbtree_api_remove_no_drop(void *ctx)
110{
111	struct bpf_rb_node *res;
112	struct node_data *n;
113
114	bpf_spin_lock(&glock);
115	res = bpf_rbtree_first(&groot);
116	if (!res)
117		goto unlock_err;
118
119	res = bpf_rbtree_remove(&groot, res);
120
121	if (res) {
122		n = container_of(res, struct node_data, node);
123		__sink(n);
124	}
125	bpf_spin_unlock(&glock);
126
127	/* if (res) { bpf_obj_drop(n); } is missing here */
128	return 0;
129
130unlock_err:
131	bpf_spin_unlock(&glock);
132	return 1;
133}
134
135SEC("?tc")
136__failure __msg("arg#1 expected pointer to allocated object")
137long rbtree_api_add_to_multiple_trees(void *ctx)
138{
139	struct node_data *n;
140
141	n = bpf_obj_new(typeof(*n));
142	if (!n)
143		return 1;
144
145	bpf_spin_lock(&glock);
146	bpf_rbtree_add(&groot, &n->node, less);
147
148	/* This add should fail since n already in groot's tree */
149	bpf_rbtree_add(&groot2, &n->node, less);
150	bpf_spin_unlock(&glock);
151	return 0;
152}
153
154SEC("?tc")
155__failure __msg("dereference of modified ptr_or_null_ ptr R2 off=16 disallowed")
156long rbtree_api_use_unchecked_remove_retval(void *ctx)
157{
158	struct bpf_rb_node *res;
159
160	bpf_spin_lock(&glock);
161
162	res = bpf_rbtree_first(&groot);
163	if (!res)
164		goto err_out;
165	res = bpf_rbtree_remove(&groot, res);
166
167	bpf_spin_unlock(&glock);
168
169	bpf_spin_lock(&glock);
170	/* Must check res for NULL before using in rbtree_add below */
171	bpf_rbtree_add(&groot, res, less);
172	bpf_spin_unlock(&glock);
173	return 0;
174
175err_out:
176	bpf_spin_unlock(&glock);
177	return 1;
178}
179
180SEC("?tc")
181__failure __msg("rbtree_remove node input must be non-owning ref")
182long rbtree_api_add_release_unlock_escape(void *ctx)
183{
184	struct node_data *n;
185
186	n = bpf_obj_new(typeof(*n));
187	if (!n)
188		return 1;
189
190	bpf_spin_lock(&glock);
191	bpf_rbtree_add(&groot, &n->node, less);
192	bpf_spin_unlock(&glock);
193
194	bpf_spin_lock(&glock);
195	/* After add() in previous critical section, n should be
196	 * release_on_unlock and released after previous spin_unlock,
197	 * so should not be possible to use it here
198	 */
199	bpf_rbtree_remove(&groot, &n->node);
200	bpf_spin_unlock(&glock);
201	return 0;
202}
203
204SEC("?tc")
205__failure __msg("rbtree_remove node input must be non-owning ref")
206long rbtree_api_first_release_unlock_escape(void *ctx)
207{
208	struct bpf_rb_node *res;
209	struct node_data *n;
210
211	bpf_spin_lock(&glock);
212	res = bpf_rbtree_first(&groot);
213	if (!res) {
214		bpf_spin_unlock(&glock);
215		return 1;
216	}
217	n = container_of(res, struct node_data, node);
218	bpf_spin_unlock(&glock);
219
220	bpf_spin_lock(&glock);
221	/* After first() in previous critical section, n should be
222	 * release_on_unlock and released after previous spin_unlock,
223	 * so should not be possible to use it here
224	 */
225	bpf_rbtree_remove(&groot, &n->node);
226	bpf_spin_unlock(&glock);
227	return 0;
228}
229
230static bool less__bad_fn_call_add(struct bpf_rb_node *a, const struct bpf_rb_node *b)
231{
232	struct node_data *node_a;
233	struct node_data *node_b;
234
235	node_a = container_of(a, struct node_data, node);
236	node_b = container_of(b, struct node_data, node);
237	bpf_rbtree_add(&groot, &node_a->node, less);
238
239	return node_a->key < node_b->key;
240}
241
242static bool less__bad_fn_call_remove(struct bpf_rb_node *a, const struct bpf_rb_node *b)
243{
244	struct node_data *node_a;
245	struct node_data *node_b;
246
247	node_a = container_of(a, struct node_data, node);
248	node_b = container_of(b, struct node_data, node);
249	bpf_rbtree_remove(&groot, &node_a->node);
250
251	return node_a->key < node_b->key;
252}
253
254static bool less__bad_fn_call_first_unlock_after(struct bpf_rb_node *a, const struct bpf_rb_node *b)
255{
256	struct node_data *node_a;
257	struct node_data *node_b;
258
259	node_a = container_of(a, struct node_data, node);
260	node_b = container_of(b, struct node_data, node);
261	bpf_rbtree_first(&groot);
262	bpf_spin_unlock(&glock);
263
264	return node_a->key < node_b->key;
265}
266
267static __always_inline
268long add_with_cb(bool (cb)(struct bpf_rb_node *a, const struct bpf_rb_node *b))
269{
270	struct node_data *n;
271
272	n = bpf_obj_new(typeof(*n));
273	if (!n)
274		return 1;
275
276	bpf_spin_lock(&glock);
277	bpf_rbtree_add(&groot, &n->node, cb);
278	bpf_spin_unlock(&glock);
279	return 0;
280}
281
282SEC("?tc")
283__failure __msg("arg#1 expected pointer to allocated object")
284long rbtree_api_add_bad_cb_bad_fn_call_add(void *ctx)
285{
286	return add_with_cb(less__bad_fn_call_add);
287}
288
289SEC("?tc")
290__failure __msg("rbtree_remove not allowed in rbtree cb")
291long rbtree_api_add_bad_cb_bad_fn_call_remove(void *ctx)
292{
293	return add_with_cb(less__bad_fn_call_remove);
294}
295
296SEC("?tc")
297__failure __msg("can't spin_{lock,unlock} in rbtree cb")
298long rbtree_api_add_bad_cb_bad_fn_call_first_unlock_after(void *ctx)
299{
300	return add_with_cb(less__bad_fn_call_first_unlock_after);
301}
302
303char _license[] SEC("license") = "GPL";
304