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_experimental.h"
9 #include "../bpf_testmod/bpf_testmod_kfunc.h"
10
11 struct plain_local;
12
13 struct node_data {
14 long key;
15 long data;
16 struct plain_local __kptr * stashed_in_local_kptr;
17 struct bpf_rb_node node;
18 };
19
20 struct refcounted_node {
21 long data;
22 struct bpf_rb_node rb_node;
23 struct bpf_refcount refcount;
24 };
25
26 struct stash {
27 struct bpf_spin_lock l;
28 struct refcounted_node __kptr *stashed;
29 };
30
31 struct {
32 __uint(type, BPF_MAP_TYPE_ARRAY);
33 __type(key, int);
34 __type(value, struct stash);
35 __uint(max_entries, 10);
36 } refcounted_node_stash SEC(".maps");
37
38 struct plain_local {
39 long key;
40 long data;
41 };
42
43 struct local_with_root {
44 long key;
45 struct bpf_spin_lock l;
46 struct bpf_rb_root r __contains(node_data, node);
47 };
48
49 struct map_value {
50 struct prog_test_ref_kfunc *not_kptr;
51 struct prog_test_ref_kfunc __kptr *val;
52 struct node_data __kptr *node;
53 struct plain_local __kptr *plain;
54 struct local_with_root __kptr *local_root;
55 };
56
57 /* This is necessary so that LLVM generates BTF for node_data struct
58 * If it's not included, a fwd reference for node_data will be generated but
59 * no struct. Example BTF of "node" field in map_value when not included:
60 *
61 * [10] PTR '(anon)' type_id=35
62 * [34] FWD 'node_data' fwd_kind=struct
63 * [35] TYPE_TAG 'kptr_ref' type_id=34
64 *
65 * (with no node_data struct defined)
66 * Had to do the same w/ bpf_kfunc_call_test_release below
67 */
68 struct node_data *just_here_because_btf_bug;
69 struct refcounted_node *just_here_because_btf_bug2;
70
71 struct {
72 __uint(type, BPF_MAP_TYPE_ARRAY);
73 __type(key, int);
74 __type(value, struct map_value);
75 __uint(max_entries, 2);
76 } some_nodes SEC(".maps");
77
less(struct bpf_rb_node * a,const struct bpf_rb_node * b)78 static bool less(struct bpf_rb_node *a, const struct bpf_rb_node *b)
79 {
80 struct node_data *node_a;
81 struct node_data *node_b;
82
83 node_a = container_of(a, struct node_data, node);
84 node_b = container_of(b, struct node_data, node);
85
86 return node_a->key < node_b->key;
87 }
88
create_and_stash(int idx,int val)89 static int create_and_stash(int idx, int val)
90 {
91 struct plain_local *inner_local_kptr;
92 struct map_value *mapval;
93 struct node_data *res;
94
95 mapval = bpf_map_lookup_elem(&some_nodes, &idx);
96 if (!mapval)
97 return 1;
98
99 inner_local_kptr = bpf_obj_new(typeof(*inner_local_kptr));
100 if (!inner_local_kptr)
101 return 2;
102
103 res = bpf_obj_new(typeof(*res));
104 if (!res) {
105 bpf_obj_drop(inner_local_kptr);
106 return 3;
107 }
108 res->key = val;
109
110 inner_local_kptr = bpf_kptr_xchg(&res->stashed_in_local_kptr, inner_local_kptr);
111 if (inner_local_kptr) {
112 /* Should never happen, we just obj_new'd res */
113 bpf_obj_drop(inner_local_kptr);
114 bpf_obj_drop(res);
115 return 4;
116 }
117
118 res = bpf_kptr_xchg(&mapval->node, res);
119 if (res)
120 bpf_obj_drop(res);
121 return 0;
122 }
123
124 SEC("tc")
stash_rb_nodes(void * ctx)125 long stash_rb_nodes(void *ctx)
126 {
127 return create_and_stash(0, 41) ?: create_and_stash(1, 42);
128 }
129
130 SEC("tc")
stash_plain(void * ctx)131 long stash_plain(void *ctx)
132 {
133 struct map_value *mapval;
134 struct plain_local *res;
135 int idx = 0;
136
137 mapval = bpf_map_lookup_elem(&some_nodes, &idx);
138 if (!mapval)
139 return 1;
140
141 res = bpf_obj_new(typeof(*res));
142 if (!res)
143 return 1;
144 res->key = 41;
145
146 res = bpf_kptr_xchg(&mapval->plain, res);
147 if (res)
148 bpf_obj_drop(res);
149 return 0;
150 }
151
152 SEC("tc")
stash_local_with_root(void * ctx)153 long stash_local_with_root(void *ctx)
154 {
155 struct local_with_root *res;
156 struct map_value *mapval;
157 struct node_data *n;
158 int idx = 0;
159
160 mapval = bpf_map_lookup_elem(&some_nodes, &idx);
161 if (!mapval)
162 return 1;
163
164 res = bpf_obj_new(typeof(*res));
165 if (!res)
166 return 2;
167 res->key = 41;
168
169 n = bpf_obj_new(typeof(*n));
170 if (!n) {
171 bpf_obj_drop(res);
172 return 3;
173 }
174
175 bpf_spin_lock(&res->l);
176 bpf_rbtree_add(&res->r, &n->node, less);
177 bpf_spin_unlock(&res->l);
178
179 res = bpf_kptr_xchg(&mapval->local_root, res);
180 if (res) {
181 bpf_obj_drop(res);
182 return 4;
183 }
184 return 0;
185 }
186
187 SEC("tc")
unstash_rb_node(void * ctx)188 long unstash_rb_node(void *ctx)
189 {
190 struct plain_local *inner_local_kptr = NULL;
191 struct map_value *mapval;
192 struct node_data *res;
193 long retval;
194 int key = 1;
195
196 mapval = bpf_map_lookup_elem(&some_nodes, &key);
197 if (!mapval)
198 return 1;
199
200 res = bpf_kptr_xchg(&mapval->node, NULL);
201 if (res) {
202 inner_local_kptr = bpf_kptr_xchg(&res->stashed_in_local_kptr, inner_local_kptr);
203 if (!inner_local_kptr) {
204 bpf_obj_drop(res);
205 return 1;
206 }
207 bpf_obj_drop(inner_local_kptr);
208
209 retval = res->key;
210 bpf_obj_drop(res);
211 return retval;
212 }
213 return 1;
214 }
215
216 SEC("tc")
stash_test_ref_kfunc(void * ctx)217 long stash_test_ref_kfunc(void *ctx)
218 {
219 struct prog_test_ref_kfunc *res;
220 struct map_value *mapval;
221 int key = 0;
222
223 mapval = bpf_map_lookup_elem(&some_nodes, &key);
224 if (!mapval)
225 return 1;
226
227 res = bpf_kptr_xchg(&mapval->val, NULL);
228 if (res)
229 bpf_kfunc_call_test_release(res);
230 return 0;
231 }
232
233 SEC("tc")
refcount_acquire_without_unstash(void * ctx)234 long refcount_acquire_without_unstash(void *ctx)
235 {
236 struct refcounted_node *p;
237 struct stash *s;
238 int ret = 0;
239
240 s = bpf_map_lookup_elem(&refcounted_node_stash, &ret);
241 if (!s)
242 return 1;
243
244 if (!s->stashed)
245 /* refcount_acquire failure is expected when no refcounted_node
246 * has been stashed before this program executes
247 */
248 return 2;
249
250 p = bpf_refcount_acquire(s->stashed);
251 if (!p)
252 return 3;
253
254 ret = s->stashed ? s->stashed->data : -1;
255 bpf_obj_drop(p);
256 return ret;
257 }
258
259 /* Helper for refcount_acquire_without_unstash test */
260 SEC("tc")
stash_refcounted_node(void * ctx)261 long stash_refcounted_node(void *ctx)
262 {
263 struct refcounted_node *p;
264 struct stash *s;
265 int key = 0;
266
267 s = bpf_map_lookup_elem(&refcounted_node_stash, &key);
268 if (!s)
269 return 1;
270
271 p = bpf_obj_new(typeof(*p));
272 if (!p)
273 return 2;
274 p->data = 42;
275
276 p = bpf_kptr_xchg(&s->stashed, p);
277 if (p) {
278 bpf_obj_drop(p);
279 return 3;
280 }
281
282 return 0;
283 }
284
285 char _license[] SEC("license") = "GPL";
286