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