1 // SPDX-License-Identifier: GPL-2.0
2 #include <vmlinux.h>
3 #include <bpf/bpf_tracing.h>
4 #include <bpf/bpf_helpers.h>
5 
6 struct map_value {
7 	struct prog_test_ref_kfunc __kptr *unref_ptr;
8 	struct prog_test_ref_kfunc __kptr_ref *ref_ptr;
9 };
10 
11 struct array_map {
12 	__uint(type, BPF_MAP_TYPE_ARRAY);
13 	__type(key, int);
14 	__type(value, struct map_value);
15 	__uint(max_entries, 1);
16 } array_map SEC(".maps");
17 
18 struct hash_map {
19 	__uint(type, BPF_MAP_TYPE_HASH);
20 	__type(key, int);
21 	__type(value, struct map_value);
22 	__uint(max_entries, 1);
23 } hash_map SEC(".maps");
24 
25 struct hash_malloc_map {
26 	__uint(type, BPF_MAP_TYPE_HASH);
27 	__type(key, int);
28 	__type(value, struct map_value);
29 	__uint(max_entries, 1);
30 	__uint(map_flags, BPF_F_NO_PREALLOC);
31 } hash_malloc_map SEC(".maps");
32 
33 struct lru_hash_map {
34 	__uint(type, BPF_MAP_TYPE_LRU_HASH);
35 	__type(key, int);
36 	__type(value, struct map_value);
37 	__uint(max_entries, 1);
38 } lru_hash_map SEC(".maps");
39 
40 #define DEFINE_MAP_OF_MAP(map_type, inner_map_type, name)       \
41 	struct {                                                \
42 		__uint(type, map_type);                         \
43 		__uint(max_entries, 1);                         \
44 		__uint(key_size, sizeof(int));                  \
45 		__uint(value_size, sizeof(int));                \
46 		__array(values, struct inner_map_type);         \
47 	} name SEC(".maps") = {                                 \
48 		.values = { [0] = &inner_map_type },            \
49 	}
50 
51 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_ARRAY_OF_MAPS, array_map, array_of_array_maps);
52 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_ARRAY_OF_MAPS, hash_map, array_of_hash_maps);
53 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_ARRAY_OF_MAPS, hash_malloc_map, array_of_hash_malloc_maps);
54 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_ARRAY_OF_MAPS, lru_hash_map, array_of_lru_hash_maps);
55 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_HASH_OF_MAPS, array_map, hash_of_array_maps);
56 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_HASH_OF_MAPS, hash_map, hash_of_hash_maps);
57 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_HASH_OF_MAPS, hash_malloc_map, hash_of_hash_malloc_maps);
58 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_HASH_OF_MAPS, lru_hash_map, hash_of_lru_hash_maps);
59 
60 extern struct prog_test_ref_kfunc *bpf_kfunc_call_test_acquire(unsigned long *sp) __ksym;
61 extern struct prog_test_ref_kfunc *
62 bpf_kfunc_call_test_kptr_get(struct prog_test_ref_kfunc **p, int a, int b) __ksym;
63 extern void bpf_kfunc_call_test_release(struct prog_test_ref_kfunc *p) __ksym;
64 
65 static void test_kptr_unref(struct map_value *v)
66 {
67 	struct prog_test_ref_kfunc *p;
68 
69 	p = v->unref_ptr;
70 	/* store untrusted_ptr_or_null_ */
71 	v->unref_ptr = p;
72 	if (!p)
73 		return;
74 	if (p->a + p->b > 100)
75 		return;
76 	/* store untrusted_ptr_ */
77 	v->unref_ptr = p;
78 	/* store NULL */
79 	v->unref_ptr = NULL;
80 }
81 
82 static void test_kptr_ref(struct map_value *v)
83 {
84 	struct prog_test_ref_kfunc *p;
85 
86 	p = v->ref_ptr;
87 	/* store ptr_or_null_ */
88 	v->unref_ptr = p;
89 	if (!p)
90 		return;
91 	if (p->a + p->b > 100)
92 		return;
93 	/* store NULL */
94 	p = bpf_kptr_xchg(&v->ref_ptr, NULL);
95 	if (!p)
96 		return;
97 	if (p->a + p->b > 100) {
98 		bpf_kfunc_call_test_release(p);
99 		return;
100 	}
101 	/* store ptr_ */
102 	v->unref_ptr = p;
103 	bpf_kfunc_call_test_release(p);
104 
105 	p = bpf_kfunc_call_test_acquire(&(unsigned long){0});
106 	if (!p)
107 		return;
108 	/* store ptr_ */
109 	p = bpf_kptr_xchg(&v->ref_ptr, p);
110 	if (!p)
111 		return;
112 	if (p->a + p->b > 100) {
113 		bpf_kfunc_call_test_release(p);
114 		return;
115 	}
116 	bpf_kfunc_call_test_release(p);
117 }
118 
119 static void test_kptr_get(struct map_value *v)
120 {
121 	struct prog_test_ref_kfunc *p;
122 
123 	p = bpf_kfunc_call_test_kptr_get(&v->ref_ptr, 0, 0);
124 	if (!p)
125 		return;
126 	if (p->a + p->b > 100) {
127 		bpf_kfunc_call_test_release(p);
128 		return;
129 	}
130 	bpf_kfunc_call_test_release(p);
131 }
132 
133 static void test_kptr(struct map_value *v)
134 {
135 	test_kptr_unref(v);
136 	test_kptr_ref(v);
137 	test_kptr_get(v);
138 }
139 
140 SEC("tc")
141 int test_map_kptr(struct __sk_buff *ctx)
142 {
143 	struct map_value *v;
144 	int key = 0;
145 
146 #define TEST(map)					\
147 	v = bpf_map_lookup_elem(&map, &key);		\
148 	if (!v)						\
149 		return 0;				\
150 	test_kptr(v)
151 
152 	TEST(array_map);
153 	TEST(hash_map);
154 	TEST(hash_malloc_map);
155 	TEST(lru_hash_map);
156 
157 #undef TEST
158 	return 0;
159 }
160 
161 SEC("tc")
162 int test_map_in_map_kptr(struct __sk_buff *ctx)
163 {
164 	struct map_value *v;
165 	int key = 0;
166 	void *map;
167 
168 #define TEST(map_in_map)                                \
169 	map = bpf_map_lookup_elem(&map_in_map, &key);   \
170 	if (!map)                                       \
171 		return 0;                               \
172 	v = bpf_map_lookup_elem(map, &key);		\
173 	if (!v)						\
174 		return 0;				\
175 	test_kptr(v)
176 
177 	TEST(array_of_array_maps);
178 	TEST(array_of_hash_maps);
179 	TEST(array_of_hash_malloc_maps);
180 	TEST(array_of_lru_hash_maps);
181 	TEST(hash_of_array_maps);
182 	TEST(hash_of_hash_maps);
183 	TEST(hash_of_hash_malloc_maps);
184 	TEST(hash_of_lru_hash_maps);
185 
186 #undef TEST
187 	return 0;
188 }
189 
190 SEC("tc")
191 int test_map_kptr_ref(struct __sk_buff *ctx)
192 {
193 	struct prog_test_ref_kfunc *p, *p_st;
194 	unsigned long arg = 0;
195 	struct map_value *v;
196 	int key = 0, ret;
197 
198 	p = bpf_kfunc_call_test_acquire(&arg);
199 	if (!p)
200 		return 1;
201 
202 	p_st = p->next;
203 	if (p_st->cnt.refs.counter != 2) {
204 		ret = 2;
205 		goto end;
206 	}
207 
208 	v = bpf_map_lookup_elem(&array_map, &key);
209 	if (!v) {
210 		ret = 3;
211 		goto end;
212 	}
213 
214 	p = bpf_kptr_xchg(&v->ref_ptr, p);
215 	if (p) {
216 		ret = 4;
217 		goto end;
218 	}
219 	if (p_st->cnt.refs.counter != 2)
220 		return 5;
221 
222 	p = bpf_kfunc_call_test_kptr_get(&v->ref_ptr, 0, 0);
223 	if (!p)
224 		return 6;
225 	if (p_st->cnt.refs.counter != 3) {
226 		ret = 7;
227 		goto end;
228 	}
229 	bpf_kfunc_call_test_release(p);
230 	if (p_st->cnt.refs.counter != 2)
231 		return 8;
232 
233 	p = bpf_kptr_xchg(&v->ref_ptr, NULL);
234 	if (!p)
235 		return 9;
236 	bpf_kfunc_call_test_release(p);
237 	if (p_st->cnt.refs.counter != 1)
238 		return 10;
239 
240 	p = bpf_kfunc_call_test_acquire(&arg);
241 	if (!p)
242 		return 11;
243 	p = bpf_kptr_xchg(&v->ref_ptr, p);
244 	if (p) {
245 		ret = 12;
246 		goto end;
247 	}
248 	if (p_st->cnt.refs.counter != 2)
249 		return 13;
250 	/* Leave in map */
251 
252 	return 0;
253 end:
254 	bpf_kfunc_call_test_release(p);
255 	return ret;
256 }
257 
258 SEC("tc")
259 int test_map_kptr_ref2(struct __sk_buff *ctx)
260 {
261 	struct prog_test_ref_kfunc *p, *p_st;
262 	struct map_value *v;
263 	int key = 0;
264 
265 	v = bpf_map_lookup_elem(&array_map, &key);
266 	if (!v)
267 		return 1;
268 
269 	p_st = v->ref_ptr;
270 	if (!p_st || p_st->cnt.refs.counter != 2)
271 		return 2;
272 
273 	p = bpf_kptr_xchg(&v->ref_ptr, NULL);
274 	if (!p)
275 		return 3;
276 	if (p_st->cnt.refs.counter != 2) {
277 		bpf_kfunc_call_test_release(p);
278 		return 4;
279 	}
280 
281 	p = bpf_kptr_xchg(&v->ref_ptr, p);
282 	if (p) {
283 		bpf_kfunc_call_test_release(p);
284 		return 5;
285 	}
286 	if (p_st->cnt.refs.counter != 2)
287 		return 6;
288 
289 	return 0;
290 }
291 
292 char _license[] SEC("license") = "GPL";
293