1 // SPDX-License-Identifier: GPL-2.0
2 /* Copyright (c) 2023 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_misc.h"
9 #include "bpf_experimental.h"
10 
11 struct node_data {
12 	long key;
13 	long list_data;
14 	struct bpf_rb_node r;
15 	struct bpf_list_node l;
16 	struct bpf_refcount ref;
17 };
18 
19 struct map_value {
20 	struct node_data __kptr *node;
21 };
22 
23 struct {
24 	__uint(type, BPF_MAP_TYPE_ARRAY);
25 	__type(key, int);
26 	__type(value, struct map_value);
27 	__uint(max_entries, 1);
28 } stashed_nodes SEC(".maps");
29 
30 struct node_acquire {
31 	long key;
32 	long data;
33 	struct bpf_rb_node node;
34 	struct bpf_refcount refcount;
35 };
36 
37 #define private(name) SEC(".bss." #name) __hidden __attribute__((aligned(8)))
38 private(A) struct bpf_spin_lock lock;
39 private(A) struct bpf_rb_root root __contains(node_data, r);
40 private(A) struct bpf_list_head head __contains(node_data, l);
41 
42 private(B) struct bpf_spin_lock alock;
43 private(B) struct bpf_rb_root aroot __contains(node_acquire, node);
44 
45 static bool less(struct bpf_rb_node *node_a, const struct bpf_rb_node *node_b)
46 {
47 	struct node_data *a;
48 	struct node_data *b;
49 
50 	a = container_of(node_a, struct node_data, r);
51 	b = container_of(node_b, struct node_data, r);
52 
53 	return a->key < b->key;
54 }
55 
56 static bool less_a(struct bpf_rb_node *a, const struct bpf_rb_node *b)
57 {
58 	struct node_acquire *node_a;
59 	struct node_acquire *node_b;
60 
61 	node_a = container_of(a, struct node_acquire, node);
62 	node_b = container_of(b, struct node_acquire, node);
63 
64 	return node_a->key < node_b->key;
65 }
66 
67 static long __insert_in_tree_and_list(struct bpf_list_head *head,
68 				      struct bpf_rb_root *root,
69 				      struct bpf_spin_lock *lock)
70 {
71 	struct node_data *n, *m;
72 
73 	n = bpf_obj_new(typeof(*n));
74 	if (!n)
75 		return -1;
76 
77 	m = bpf_refcount_acquire(n);
78 	m->key = 123;
79 	m->list_data = 456;
80 
81 	bpf_spin_lock(lock);
82 	if (bpf_rbtree_add(root, &n->r, less)) {
83 		/* Failure to insert - unexpected */
84 		bpf_spin_unlock(lock);
85 		bpf_obj_drop(m);
86 		return -2;
87 	}
88 	bpf_spin_unlock(lock);
89 
90 	bpf_spin_lock(lock);
91 	if (bpf_list_push_front(head, &m->l)) {
92 		/* Failure to insert - unexpected */
93 		bpf_spin_unlock(lock);
94 		return -3;
95 	}
96 	bpf_spin_unlock(lock);
97 	return 0;
98 }
99 
100 static long __stash_map_insert_tree(int idx, int val, struct bpf_rb_root *root,
101 				    struct bpf_spin_lock *lock)
102 {
103 	struct map_value *mapval;
104 	struct node_data *n, *m;
105 
106 	mapval = bpf_map_lookup_elem(&stashed_nodes, &idx);
107 	if (!mapval)
108 		return -1;
109 
110 	n = bpf_obj_new(typeof(*n));
111 	if (!n)
112 		return -2;
113 
114 	n->key = val;
115 	m = bpf_refcount_acquire(n);
116 
117 	n = bpf_kptr_xchg(&mapval->node, n);
118 	if (n) {
119 		bpf_obj_drop(n);
120 		bpf_obj_drop(m);
121 		return -3;
122 	}
123 
124 	bpf_spin_lock(lock);
125 	if (bpf_rbtree_add(root, &m->r, less)) {
126 		/* Failure to insert - unexpected */
127 		bpf_spin_unlock(lock);
128 		return -4;
129 	}
130 	bpf_spin_unlock(lock);
131 	return 0;
132 }
133 
134 static long __read_from_tree(struct bpf_rb_root *root,
135 			     struct bpf_spin_lock *lock,
136 			     bool remove_from_tree)
137 {
138 	struct bpf_rb_node *rb;
139 	struct node_data *n;
140 	long res = -99;
141 
142 	bpf_spin_lock(lock);
143 
144 	rb = bpf_rbtree_first(root);
145 	if (!rb) {
146 		bpf_spin_unlock(lock);
147 		return -1;
148 	}
149 
150 	n = container_of(rb, struct node_data, r);
151 	res = n->key;
152 
153 	if (!remove_from_tree) {
154 		bpf_spin_unlock(lock);
155 		return res;
156 	}
157 
158 	rb = bpf_rbtree_remove(root, rb);
159 	bpf_spin_unlock(lock);
160 	if (!rb)
161 		return -2;
162 	n = container_of(rb, struct node_data, r);
163 	bpf_obj_drop(n);
164 	return res;
165 }
166 
167 static long __read_from_list(struct bpf_list_head *head,
168 			     struct bpf_spin_lock *lock,
169 			     bool remove_from_list)
170 {
171 	struct bpf_list_node *l;
172 	struct node_data *n;
173 	long res = -99;
174 
175 	bpf_spin_lock(lock);
176 
177 	l = bpf_list_pop_front(head);
178 	if (!l) {
179 		bpf_spin_unlock(lock);
180 		return -1;
181 	}
182 
183 	n = container_of(l, struct node_data, l);
184 	res = n->list_data;
185 
186 	if (!remove_from_list) {
187 		if (bpf_list_push_back(head, &n->l)) {
188 			bpf_spin_unlock(lock);
189 			return -2;
190 		}
191 	}
192 
193 	bpf_spin_unlock(lock);
194 
195 	if (remove_from_list)
196 		bpf_obj_drop(n);
197 	return res;
198 }
199 
200 static long __read_from_unstash(int idx)
201 {
202 	struct node_data *n = NULL;
203 	struct map_value *mapval;
204 	long val = -99;
205 
206 	mapval = bpf_map_lookup_elem(&stashed_nodes, &idx);
207 	if (!mapval)
208 		return -1;
209 
210 	n = bpf_kptr_xchg(&mapval->node, n);
211 	if (!n)
212 		return -2;
213 
214 	val = n->key;
215 	bpf_obj_drop(n);
216 	return val;
217 }
218 
219 #define INSERT_READ_BOTH(rem_tree, rem_list, desc)			\
220 SEC("tc")								\
221 __description(desc)							\
222 __success __retval(579)							\
223 long insert_and_remove_tree_##rem_tree##_list_##rem_list(void *ctx)	\
224 {									\
225 	long err, tree_data, list_data;					\
226 									\
227 	err = __insert_in_tree_and_list(&head, &root, &lock);		\
228 	if (err)							\
229 		return err;						\
230 									\
231 	err = __read_from_tree(&root, &lock, rem_tree);			\
232 	if (err < 0)							\
233 		return err;						\
234 	else								\
235 		tree_data = err;					\
236 									\
237 	err = __read_from_list(&head, &lock, rem_list);			\
238 	if (err < 0)							\
239 		return err;						\
240 	else								\
241 		list_data = err;					\
242 									\
243 	return tree_data + list_data;					\
244 }
245 
246 /* After successful insert of struct node_data into both collections:
247  *   - it should have refcount = 2
248  *   - removing / not removing the node_data from a collection after
249  *     reading should have no effect on ability to read / remove from
250  *     the other collection
251  */
252 INSERT_READ_BOTH(true, true, "insert_read_both: remove from tree + list");
253 INSERT_READ_BOTH(false, false, "insert_read_both: remove from neither");
254 INSERT_READ_BOTH(true, false, "insert_read_both: remove from tree");
255 INSERT_READ_BOTH(false, true, "insert_read_both: remove from list");
256 
257 #undef INSERT_READ_BOTH
258 #define INSERT_READ_BOTH(rem_tree, rem_list, desc)			\
259 SEC("tc")								\
260 __description(desc)							\
261 __success __retval(579)							\
262 long insert_and_remove_lf_tree_##rem_tree##_list_##rem_list(void *ctx)	\
263 {									\
264 	long err, tree_data, list_data;					\
265 									\
266 	err = __insert_in_tree_and_list(&head, &root, &lock);		\
267 	if (err)							\
268 		return err;						\
269 									\
270 	err = __read_from_list(&head, &lock, rem_list);			\
271 	if (err < 0)							\
272 		return err;						\
273 	else								\
274 		list_data = err;					\
275 									\
276 	err = __read_from_tree(&root, &lock, rem_tree);			\
277 	if (err < 0)							\
278 		return err;						\
279 	else								\
280 		tree_data = err;					\
281 									\
282 	return tree_data + list_data;					\
283 }
284 
285 /* Similar to insert_read_both, but list data is read and possibly removed
286  * first
287  *
288  * Results should be no different than reading and possibly removing rbtree
289  * node first
290  */
291 INSERT_READ_BOTH(true, true, "insert_read_both_list_first: remove from tree + list");
292 INSERT_READ_BOTH(false, false, "insert_read_both_list_first: remove from neither");
293 INSERT_READ_BOTH(true, false, "insert_read_both_list_first: remove from tree");
294 INSERT_READ_BOTH(false, true, "insert_read_both_list_first: remove from list");
295 
296 #define INSERT_DOUBLE_READ_AND_DEL(read_fn, read_root, desc)		\
297 SEC("tc")								\
298 __description(desc)							\
299 __success __retval(-1)							\
300 long insert_double_##read_fn##_and_del_##read_root(void *ctx)		\
301 {									\
302 	long err, list_data;						\
303 									\
304 	err = __insert_in_tree_and_list(&head, &root, &lock);		\
305 	if (err)							\
306 		return err;						\
307 									\
308 	err = read_fn(&read_root, &lock, true);				\
309 	if (err < 0)							\
310 		return err;						\
311 	else								\
312 		list_data = err;					\
313 									\
314 	err = read_fn(&read_root, &lock, true);				\
315 	if (err < 0)							\
316 		return err;						\
317 									\
318 	return err + list_data;						\
319 }
320 
321 /* Insert into both tree and list, then try reading-and-removing from either twice
322  *
323  * The second read-and-remove should fail on read step since the node has
324  * already been removed
325  */
326 INSERT_DOUBLE_READ_AND_DEL(__read_from_tree, root, "insert_double_del: 2x read-and-del from tree");
327 INSERT_DOUBLE_READ_AND_DEL(__read_from_list, head, "insert_double_del: 2x read-and-del from list");
328 
329 #define INSERT_STASH_READ(rem_tree, desc)				\
330 SEC("tc")								\
331 __description(desc)							\
332 __success __retval(84)							\
333 long insert_rbtree_and_stash__del_tree_##rem_tree(void *ctx)		\
334 {									\
335 	long err, tree_data, map_data;					\
336 									\
337 	err = __stash_map_insert_tree(0, 42, &root, &lock);		\
338 	if (err)							\
339 		return err;						\
340 									\
341 	err = __read_from_tree(&root, &lock, rem_tree);			\
342 	if (err < 0)							\
343 		return err;						\
344 	else								\
345 		tree_data = err;					\
346 									\
347 	err = __read_from_unstash(0);					\
348 	if (err < 0)							\
349 		return err;						\
350 	else								\
351 		map_data = err;						\
352 									\
353 	return tree_data + map_data;					\
354 }
355 
356 /* Stash a refcounted node in map_val, insert same node into tree, then try
357  * reading data from tree then unstashed map_val, possibly removing from tree
358  *
359  * Removing from tree should have no effect on map_val kptr validity
360  */
361 INSERT_STASH_READ(true, "insert_stash_read: remove from tree");
362 INSERT_STASH_READ(false, "insert_stash_read: don't remove from tree");
363 
364 SEC("tc")
365 __success
366 long rbtree_refcounted_node_ref_escapes(void *ctx)
367 {
368 	struct node_acquire *n, *m;
369 
370 	n = bpf_obj_new(typeof(*n));
371 	if (!n)
372 		return 1;
373 
374 	bpf_spin_lock(&alock);
375 	bpf_rbtree_add(&aroot, &n->node, less_a);
376 	m = bpf_refcount_acquire(n);
377 	bpf_spin_unlock(&alock);
378 	if (!m)
379 		return 2;
380 
381 	m->key = 2;
382 	bpf_obj_drop(m);
383 	return 0;
384 }
385 
386 SEC("tc")
387 __success
388 long rbtree_refcounted_node_ref_escapes_owning_input(void *ctx)
389 {
390 	struct node_acquire *n, *m;
391 
392 	n = bpf_obj_new(typeof(*n));
393 	if (!n)
394 		return 1;
395 
396 	m = bpf_refcount_acquire(n);
397 	m->key = 2;
398 
399 	bpf_spin_lock(&alock);
400 	bpf_rbtree_add(&aroot, &n->node, less_a);
401 	bpf_spin_unlock(&alock);
402 
403 	bpf_obj_drop(m);
404 
405 	return 0;
406 }
407 
408 char _license[] SEC("license") = "GPL";
409