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(0, 41) ?: create_and_stash(1, 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(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
Linux® is a registered trademark of Linus Torvalds in the United States and other countries.
TOMOYO® is a registered trademark of NTT DATA CORPORATION.