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 
9 struct 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)))
16 private(A) struct bpf_spin_lock glock;
17 private(A) struct bpf_rb_root groot __contains(node_data, node);
18 private(A) struct bpf_rb_root groot2 __contains(node_data, node);
19 
less(struct bpf_rb_node * a,const struct bpf_rb_node * b)20 static 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 
31 SEC("?tc")
32 __failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
rbtree_api_nolock_add(void * ctx)33 long 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 
45 SEC("?tc")
46 __failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
rbtree_api_nolock_remove(void * ctx)47 long 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 
63 SEC("?tc")
64 __failure __msg("bpf_spin_lock at off=16 must be held for bpf_rb_root")
rbtree_api_nolock_first(void * ctx)65 long rbtree_api_nolock_first(void *ctx)
66 {
67 	bpf_rbtree_first(&groot);
68 	return 0;
69 }
70 
71 SEC("?tc")
72 __failure __msg("rbtree_remove node input must be non-owning ref")
rbtree_api_remove_unadded_node(void * ctx)73 long 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 
107 SEC("?tc")
108 __failure __msg("Unreleased reference id=2 alloc_insn=11")
rbtree_api_remove_no_drop(void * ctx)109 long 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 	n = container_of(res, struct node_data, node);
122 	bpf_spin_unlock(&glock);
123 
124 	/* bpf_obj_drop(n) is missing here */
125 	return 0;
126 
127 unlock_err:
128 	bpf_spin_unlock(&glock);
129 	return 1;
130 }
131 
132 SEC("?tc")
133 __failure __msg("arg#1 expected pointer to allocated object")
rbtree_api_add_to_multiple_trees(void * ctx)134 long rbtree_api_add_to_multiple_trees(void *ctx)
135 {
136 	struct node_data *n;
137 
138 	n = bpf_obj_new(typeof(*n));
139 	if (!n)
140 		return 1;
141 
142 	bpf_spin_lock(&glock);
143 	bpf_rbtree_add(&groot, &n->node, less);
144 
145 	/* This add should fail since n already in groot's tree */
146 	bpf_rbtree_add(&groot2, &n->node, less);
147 	bpf_spin_unlock(&glock);
148 	return 0;
149 }
150 
151 SEC("?tc")
152 __failure __msg("rbtree_remove node input must be non-owning ref")
rbtree_api_add_release_unlock_escape(void * ctx)153 long rbtree_api_add_release_unlock_escape(void *ctx)
154 {
155 	struct node_data *n;
156 
157 	n = bpf_obj_new(typeof(*n));
158 	if (!n)
159 		return 1;
160 
161 	bpf_spin_lock(&glock);
162 	bpf_rbtree_add(&groot, &n->node, less);
163 	bpf_spin_unlock(&glock);
164 
165 	bpf_spin_lock(&glock);
166 	/* After add() in previous critical section, n should be
167 	 * release_on_unlock and released after previous spin_unlock,
168 	 * so should not be possible to use it here
169 	 */
170 	bpf_rbtree_remove(&groot, &n->node);
171 	bpf_spin_unlock(&glock);
172 	return 0;
173 }
174 
175 SEC("?tc")
176 __failure __msg("rbtree_remove node input must be non-owning ref")
rbtree_api_release_aliasing(void * ctx)177 long rbtree_api_release_aliasing(void *ctx)
178 {
179 	struct node_data *n, *m, *o;
180 	struct bpf_rb_node *res;
181 
182 	n = bpf_obj_new(typeof(*n));
183 	if (!n)
184 		return 1;
185 
186 	bpf_spin_lock(&glock);
187 	bpf_rbtree_add(&groot, &n->node, less);
188 	bpf_spin_unlock(&glock);
189 
190 	bpf_spin_lock(&glock);
191 
192 	/* m and o point to the same node,
193 	 * but verifier doesn't know this
194 	 */
195 	res = bpf_rbtree_first(&groot);
196 	if (!res)
197 		return 1;
198 	o = container_of(res, struct node_data, node);
199 
200 	res = bpf_rbtree_first(&groot);
201 	if (!res)
202 		return 1;
203 	m = container_of(res, struct node_data, node);
204 
205 	bpf_rbtree_remove(&groot, &m->node);
206 	/* This second remove shouldn't be possible. Retval of previous
207 	 * remove returns owning reference to m, which is the same
208 	 * node o's non-owning ref is pointing at
209 	 *
210 	 * In order to preserve property
211 	 *   * owning ref must not be in rbtree
212 	 *   * non-owning ref must be in rbtree
213 	 *
214 	 * o's ref must be invalidated after previous remove. Otherwise
215 	 * we'd have non-owning ref to node that isn't in rbtree, and
216 	 * verifier wouldn't be able to use type system to prevent remove
217 	 * of ref that already isn't in any tree. Would have to do runtime
218 	 * checks in that case.
219 	 */
220 	bpf_rbtree_remove(&groot, &o->node);
221 
222 	bpf_spin_unlock(&glock);
223 	return 0;
224 }
225 
226 SEC("?tc")
227 __failure __msg("rbtree_remove node input must be non-owning ref")
rbtree_api_first_release_unlock_escape(void * ctx)228 long rbtree_api_first_release_unlock_escape(void *ctx)
229 {
230 	struct bpf_rb_node *res;
231 	struct node_data *n;
232 
233 	bpf_spin_lock(&glock);
234 	res = bpf_rbtree_first(&groot);
235 	if (res)
236 		n = container_of(res, struct node_data, node);
237 	bpf_spin_unlock(&glock);
238 
239 	bpf_spin_lock(&glock);
240 	/* After first() in previous critical section, n should be
241 	 * release_on_unlock and released after previous spin_unlock,
242 	 * so should not be possible to use it here
243 	 */
244 	bpf_rbtree_remove(&groot, &n->node);
245 	bpf_spin_unlock(&glock);
246 	return 0;
247 }
248 
less__bad_fn_call_add(struct bpf_rb_node * a,const struct bpf_rb_node * b)249 static bool less__bad_fn_call_add(struct bpf_rb_node *a, const struct bpf_rb_node *b)
250 {
251 	struct node_data *node_a;
252 	struct node_data *node_b;
253 
254 	node_a = container_of(a, struct node_data, node);
255 	node_b = container_of(b, struct node_data, node);
256 	bpf_rbtree_add(&groot, &node_a->node, less);
257 
258 	return node_a->key < node_b->key;
259 }
260 
less__bad_fn_call_remove(struct bpf_rb_node * a,const struct bpf_rb_node * b)261 static bool less__bad_fn_call_remove(struct bpf_rb_node *a, const struct bpf_rb_node *b)
262 {
263 	struct node_data *node_a;
264 	struct node_data *node_b;
265 
266 	node_a = container_of(a, struct node_data, node);
267 	node_b = container_of(b, struct node_data, node);
268 	bpf_rbtree_remove(&groot, &node_a->node);
269 
270 	return node_a->key < node_b->key;
271 }
272 
less__bad_fn_call_first_unlock_after(struct bpf_rb_node * a,const struct bpf_rb_node * b)273 static bool less__bad_fn_call_first_unlock_after(struct bpf_rb_node *a, const struct bpf_rb_node *b)
274 {
275 	struct node_data *node_a;
276 	struct node_data *node_b;
277 
278 	node_a = container_of(a, struct node_data, node);
279 	node_b = container_of(b, struct node_data, node);
280 	bpf_rbtree_first(&groot);
281 	bpf_spin_unlock(&glock);
282 
283 	return node_a->key < node_b->key;
284 }
285 
286 static __always_inline
add_with_cb(bool (cb)(struct bpf_rb_node * a,const struct bpf_rb_node * b))287 long add_with_cb(bool (cb)(struct bpf_rb_node *a, const struct bpf_rb_node *b))
288 {
289 	struct node_data *n;
290 
291 	n = bpf_obj_new(typeof(*n));
292 	if (!n)
293 		return 1;
294 
295 	bpf_spin_lock(&glock);
296 	bpf_rbtree_add(&groot, &n->node, cb);
297 	bpf_spin_unlock(&glock);
298 	return 0;
299 }
300 
301 SEC("?tc")
302 __failure __msg("arg#1 expected pointer to allocated object")
rbtree_api_add_bad_cb_bad_fn_call_add(void * ctx)303 long rbtree_api_add_bad_cb_bad_fn_call_add(void *ctx)
304 {
305 	return add_with_cb(less__bad_fn_call_add);
306 }
307 
308 SEC("?tc")
309 __failure __msg("rbtree_remove not allowed in rbtree cb")
rbtree_api_add_bad_cb_bad_fn_call_remove(void * ctx)310 long rbtree_api_add_bad_cb_bad_fn_call_remove(void *ctx)
311 {
312 	return add_with_cb(less__bad_fn_call_remove);
313 }
314 
315 SEC("?tc")
316 __failure __msg("can't spin_{lock,unlock} in rbtree cb")
rbtree_api_add_bad_cb_bad_fn_call_first_unlock_after(void * ctx)317 long rbtree_api_add_bad_cb_bad_fn_call_first_unlock_after(void *ctx)
318 {
319 	return add_with_cb(less__bad_fn_call_first_unlock_after);
320 }
321 
322 char _license[] SEC("license") = "GPL";
323