1 // SPDX-License-Identifier: GPL-2.0
2 #include <vmlinux.h>
3 #include <bpf/bpf_tracing.h>
4 #include <bpf/bpf_helpers.h>
5 #include "../test_kmods/bpf_testmod_kfunc.h"
6
7 struct map_uninit_value {
8 struct prog_test_ref_kfunc __kptr_untrusted *unref_ptr;
9 __u32 data;
10 } __attribute__((packed));
11
12 struct {
13 __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
14 __type(key, int);
15 __type(value, struct map_uninit_value);
16 __uint(max_entries, 1);
17 } pcpu_array SEC(".maps");
18
19 struct map_value {
20 struct prog_test_ref_kfunc __kptr_untrusted *unref_ptr;
21 struct prog_test_ref_kfunc __kptr *ref_ptr;
22 };
23
24 struct array_map {
25 __uint(type, BPF_MAP_TYPE_ARRAY);
26 __type(key, int);
27 __type(value, struct map_value);
28 __uint(max_entries, 1);
29 } array_map SEC(".maps");
30
31 struct pcpu_array_map {
32 __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
33 __type(key, int);
34 __type(value, struct map_value);
35 __uint(max_entries, 1);
36 } pcpu_array_map SEC(".maps");
37
38 struct hash_map {
39 __uint(type, BPF_MAP_TYPE_HASH);
40 __type(key, int);
41 __type(value, struct map_value);
42 __uint(max_entries, 1);
43 } hash_map SEC(".maps");
44
45 struct pcpu_hash_map {
46 __uint(type, BPF_MAP_TYPE_PERCPU_HASH);
47 __type(key, int);
48 __type(value, struct map_value);
49 __uint(max_entries, 1);
50 } pcpu_hash_map SEC(".maps");
51
52 struct hash_malloc_map {
53 __uint(type, BPF_MAP_TYPE_HASH);
54 __type(key, int);
55 __type(value, struct map_value);
56 __uint(max_entries, 1);
57 __uint(map_flags, BPF_F_NO_PREALLOC);
58 } hash_malloc_map SEC(".maps");
59
60 struct pcpu_hash_malloc_map {
61 __uint(type, BPF_MAP_TYPE_PERCPU_HASH);
62 __type(key, int);
63 __type(value, struct map_value);
64 __uint(max_entries, 1);
65 __uint(map_flags, BPF_F_NO_PREALLOC);
66 } pcpu_hash_malloc_map SEC(".maps");
67
68 struct lru_hash_map {
69 __uint(type, BPF_MAP_TYPE_LRU_HASH);
70 __type(key, int);
71 __type(value, struct map_value);
72 __uint(max_entries, 1);
73 } lru_hash_map SEC(".maps");
74
75 struct lru_pcpu_hash_map {
76 __uint(type, BPF_MAP_TYPE_LRU_PERCPU_HASH);
77 __type(key, int);
78 __type(value, struct map_value);
79 __uint(max_entries, 1);
80 } lru_pcpu_hash_map SEC(".maps");
81
82 struct cgrp_ls_map {
83 __uint(type, BPF_MAP_TYPE_CGRP_STORAGE);
84 __uint(map_flags, BPF_F_NO_PREALLOC);
85 __type(key, int);
86 __type(value, struct map_value);
87 } cgrp_ls_map SEC(".maps");
88
89 struct task_ls_map {
90 __uint(type, BPF_MAP_TYPE_TASK_STORAGE);
91 __uint(map_flags, BPF_F_NO_PREALLOC);
92 __type(key, int);
93 __type(value, struct map_value);
94 } task_ls_map SEC(".maps");
95
96 struct inode_ls_map {
97 __uint(type, BPF_MAP_TYPE_INODE_STORAGE);
98 __uint(map_flags, BPF_F_NO_PREALLOC);
99 __type(key, int);
100 __type(value, struct map_value);
101 } inode_ls_map SEC(".maps");
102
103 struct sk_ls_map {
104 __uint(type, BPF_MAP_TYPE_SK_STORAGE);
105 __uint(map_flags, BPF_F_NO_PREALLOC);
106 __type(key, int);
107 __type(value, struct map_value);
108 } sk_ls_map SEC(".maps");
109
110 #define DEFINE_MAP_OF_MAP(map_type, inner_map_type, name) \
111 struct { \
112 __uint(type, map_type); \
113 __uint(max_entries, 1); \
114 __uint(key_size, sizeof(int)); \
115 __uint(value_size, sizeof(int)); \
116 __array(values, struct inner_map_type); \
117 } name SEC(".maps") = { \
118 .values = { [0] = &inner_map_type }, \
119 }
120
121 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_ARRAY_OF_MAPS, array_map, array_of_array_maps);
122 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_ARRAY_OF_MAPS, hash_map, array_of_hash_maps);
123 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_ARRAY_OF_MAPS, hash_malloc_map, array_of_hash_malloc_maps);
124 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_ARRAY_OF_MAPS, lru_hash_map, array_of_lru_hash_maps);
125 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_ARRAY_OF_MAPS, pcpu_array_map, array_of_pcpu_array_maps);
126 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_ARRAY_OF_MAPS, pcpu_hash_map, array_of_pcpu_hash_maps);
127 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_HASH_OF_MAPS, array_map, hash_of_array_maps);
128 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_HASH_OF_MAPS, hash_map, hash_of_hash_maps);
129 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_HASH_OF_MAPS, hash_malloc_map, hash_of_hash_malloc_maps);
130 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_HASH_OF_MAPS, lru_hash_map, hash_of_lru_hash_maps);
131 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_HASH_OF_MAPS, pcpu_array_map, hash_of_pcpu_array_maps);
132 DEFINE_MAP_OF_MAP(BPF_MAP_TYPE_HASH_OF_MAPS, pcpu_hash_map, hash_of_pcpu_hash_maps);
133
134 #define WRITE_ONCE(x, val) ((*(volatile typeof(x) *) &(x)) = (val))
135
test_kptr_unref(struct map_value * v)136 static void test_kptr_unref(struct map_value *v)
137 {
138 struct prog_test_ref_kfunc *p;
139
140 p = v->unref_ptr;
141 /* store untrusted_ptr_or_null_ */
142 WRITE_ONCE(v->unref_ptr, p);
143 if (!p)
144 return;
145 if (p->a + p->b > 100)
146 return;
147 /* store untrusted_ptr_ */
148 WRITE_ONCE(v->unref_ptr, p);
149 /* store NULL */
150 WRITE_ONCE(v->unref_ptr, NULL);
151 }
152
test_kptr_ref(struct map_value * v)153 static void test_kptr_ref(struct map_value *v)
154 {
155 struct prog_test_ref_kfunc *p;
156
157 p = v->ref_ptr;
158 /* store ptr_or_null_ */
159 WRITE_ONCE(v->unref_ptr, p);
160 if (!p)
161 return;
162 /*
163 * p is rcu_ptr_prog_test_ref_kfunc,
164 * because bpf prog is non-sleepable and runs in RCU CS.
165 * p can be passed to kfunc that requires KF_RCU.
166 */
167 bpf_kfunc_call_test_ref(p);
168 if (p->a + p->b > 100)
169 return;
170 /* store NULL */
171 p = bpf_kptr_xchg(&v->ref_ptr, NULL);
172 if (!p)
173 return;
174 /*
175 * p is trusted_ptr_prog_test_ref_kfunc.
176 * p can be passed to kfunc that requires KF_RCU.
177 */
178 bpf_kfunc_call_test_ref(p);
179 if (p->a + p->b > 100) {
180 bpf_kfunc_call_test_release(p);
181 return;
182 }
183 /* store ptr_ */
184 WRITE_ONCE(v->unref_ptr, p);
185 bpf_kfunc_call_test_release(p);
186
187 p = bpf_kfunc_call_test_acquire(&(unsigned long){0});
188 if (!p)
189 return;
190 /* store ptr_ */
191 p = bpf_kptr_xchg(&v->ref_ptr, p);
192 if (!p)
193 return;
194 if (p->a + p->b > 100) {
195 bpf_kfunc_call_test_release(p);
196 return;
197 }
198 bpf_kfunc_call_test_release(p);
199 }
200
test_kptr(struct map_value * v)201 static void test_kptr(struct map_value *v)
202 {
203 test_kptr_unref(v);
204 test_kptr_ref(v);
205 }
206
207 SEC("tc")
test_map_kptr(struct __sk_buff * ctx)208 int test_map_kptr(struct __sk_buff *ctx)
209 {
210 struct map_value *v;
211 int key = 0;
212
213 #define TEST(map) \
214 v = bpf_map_lookup_elem(&map, &key); \
215 if (!v) \
216 return 0; \
217 test_kptr(v)
218
219 TEST(array_map);
220 TEST(hash_map);
221 TEST(hash_malloc_map);
222 TEST(lru_hash_map);
223 TEST(pcpu_array_map);
224 TEST(pcpu_hash_map);
225
226 #undef TEST
227 return 0;
228 }
229
230 SEC("tp_btf/cgroup_mkdir")
BPF_PROG(test_cgrp_map_kptr,struct cgroup * cgrp,const char * path)231 int BPF_PROG(test_cgrp_map_kptr, struct cgroup *cgrp, const char *path)
232 {
233 struct map_value *v;
234
235 v = bpf_cgrp_storage_get(&cgrp_ls_map, cgrp, NULL, BPF_LOCAL_STORAGE_GET_F_CREATE);
236 if (v)
237 test_kptr(v);
238 return 0;
239 }
240
241 SEC("lsm/inode_unlink")
BPF_PROG(test_task_map_kptr,struct inode * inode,struct dentry * victim)242 int BPF_PROG(test_task_map_kptr, struct inode *inode, struct dentry *victim)
243 {
244 struct task_struct *task;
245 struct map_value *v;
246
247 task = bpf_get_current_task_btf();
248 if (!task)
249 return 0;
250 v = bpf_task_storage_get(&task_ls_map, task, NULL, BPF_LOCAL_STORAGE_GET_F_CREATE);
251 if (v)
252 test_kptr(v);
253 return 0;
254 }
255
256 SEC("lsm/inode_unlink")
BPF_PROG(test_inode_map_kptr,struct inode * inode,struct dentry * victim)257 int BPF_PROG(test_inode_map_kptr, struct inode *inode, struct dentry *victim)
258 {
259 struct map_value *v;
260
261 v = bpf_inode_storage_get(&inode_ls_map, inode, NULL, BPF_LOCAL_STORAGE_GET_F_CREATE);
262 if (v)
263 test_kptr(v);
264 return 0;
265 }
266
267 SEC("tc")
test_sk_map_kptr(struct __sk_buff * ctx)268 int test_sk_map_kptr(struct __sk_buff *ctx)
269 {
270 struct map_value *v;
271 struct bpf_sock *sk;
272
273 sk = ctx->sk;
274 if (!sk)
275 return 0;
276 v = bpf_sk_storage_get(&sk_ls_map, sk, NULL, BPF_LOCAL_STORAGE_GET_F_CREATE);
277 if (v)
278 test_kptr(v);
279 return 0;
280 }
281
282 SEC("tc")
test_map_in_map_kptr(struct __sk_buff * ctx)283 int test_map_in_map_kptr(struct __sk_buff *ctx)
284 {
285 struct map_value *v;
286 int key = 0;
287 void *map;
288
289 #define TEST(map_in_map) \
290 map = bpf_map_lookup_elem(&map_in_map, &key); \
291 if (!map) \
292 return 0; \
293 v = bpf_map_lookup_elem(map, &key); \
294 if (!v) \
295 return 0; \
296 test_kptr(v)
297
298 TEST(array_of_array_maps);
299 TEST(array_of_hash_maps);
300 TEST(array_of_hash_malloc_maps);
301 TEST(array_of_lru_hash_maps);
302 TEST(array_of_pcpu_array_maps);
303 TEST(array_of_pcpu_hash_maps);
304 TEST(hash_of_array_maps);
305 TEST(hash_of_hash_maps);
306 TEST(hash_of_hash_malloc_maps);
307 TEST(hash_of_lru_hash_maps);
308 TEST(hash_of_pcpu_array_maps);
309 TEST(hash_of_pcpu_hash_maps);
310
311 #undef TEST
312 return 0;
313 }
314
315 int ref = 1;
316
317 static __always_inline
test_map_kptr_ref_pre(struct map_value * v)318 int test_map_kptr_ref_pre(struct map_value *v)
319 {
320 struct prog_test_ref_kfunc *p, *p_st;
321 unsigned long arg = 0;
322 int ret;
323
324 p = bpf_kfunc_call_test_acquire(&arg);
325 if (!p)
326 return 1;
327 ref++;
328
329 p_st = p->next;
330 if (p_st->cnt.refs.counter != ref) {
331 ret = 2;
332 goto end;
333 }
334
335 p = bpf_kptr_xchg(&v->ref_ptr, p);
336 if (p) {
337 ret = 3;
338 goto end;
339 }
340 if (p_st->cnt.refs.counter != ref)
341 return 4;
342
343 p = bpf_kptr_xchg(&v->ref_ptr, NULL);
344 if (!p)
345 return 5;
346 bpf_kfunc_call_test_release(p);
347 ref--;
348 if (p_st->cnt.refs.counter != ref)
349 return 6;
350
351 p = bpf_kfunc_call_test_acquire(&arg);
352 if (!p)
353 return 7;
354 ref++;
355 p = bpf_kptr_xchg(&v->ref_ptr, p);
356 if (p) {
357 ret = 8;
358 goto end;
359 }
360 if (p_st->cnt.refs.counter != ref)
361 return 9;
362 /* Leave in map */
363
364 return 0;
365 end:
366 ref--;
367 bpf_kfunc_call_test_release(p);
368 return ret;
369 }
370
371 static __always_inline
test_map_kptr_ref_post(struct map_value * v)372 int test_map_kptr_ref_post(struct map_value *v)
373 {
374 struct prog_test_ref_kfunc *p, *p_st;
375
376 p_st = v->ref_ptr;
377 if (!p_st || p_st->cnt.refs.counter != ref)
378 return 1;
379
380 p = bpf_kptr_xchg(&v->ref_ptr, NULL);
381 if (!p)
382 return 2;
383 if (p_st->cnt.refs.counter != ref) {
384 bpf_kfunc_call_test_release(p);
385 return 3;
386 }
387
388 p = bpf_kptr_xchg(&v->ref_ptr, p);
389 if (p) {
390 bpf_kfunc_call_test_release(p);
391 return 4;
392 }
393 if (p_st->cnt.refs.counter != ref)
394 return 5;
395
396 return 0;
397 }
398
399 #define TEST(map) \
400 v = bpf_map_lookup_elem(&map, &key); \
401 if (!v) \
402 return -1; \
403 ret = test_map_kptr_ref_pre(v); \
404 if (ret) \
405 return ret;
406
407 #define TEST_PCPU(map) \
408 v = bpf_map_lookup_percpu_elem(&map, &key, 0); \
409 if (!v) \
410 return -1; \
411 ret = test_map_kptr_ref_pre(v); \
412 if (ret) \
413 return ret;
414
415 SEC("tc")
test_map_kptr_ref1(struct __sk_buff * ctx)416 int test_map_kptr_ref1(struct __sk_buff *ctx)
417 {
418 struct map_value *v, val = {};
419 int key = 0, ret;
420
421 bpf_map_update_elem(&hash_map, &key, &val, 0);
422 bpf_map_update_elem(&hash_malloc_map, &key, &val, 0);
423 bpf_map_update_elem(&lru_hash_map, &key, &val, 0);
424
425 bpf_map_update_elem(&pcpu_hash_map, &key, &val, 0);
426 bpf_map_update_elem(&pcpu_hash_malloc_map, &key, &val, 0);
427 bpf_map_update_elem(&lru_pcpu_hash_map, &key, &val, 0);
428
429 TEST(array_map);
430 TEST(hash_map);
431 TEST(hash_malloc_map);
432 TEST(lru_hash_map);
433
434 TEST_PCPU(pcpu_array_map);
435 TEST_PCPU(pcpu_hash_map);
436 TEST_PCPU(pcpu_hash_malloc_map);
437 TEST_PCPU(lru_pcpu_hash_map);
438
439 return 0;
440 }
441
442 #undef TEST
443 #undef TEST_PCPU
444
445 #define TEST(map) \
446 v = bpf_map_lookup_elem(&map, &key); \
447 if (!v) \
448 return -1; \
449 ret = test_map_kptr_ref_post(v); \
450 if (ret) \
451 return ret;
452
453 #define TEST_PCPU(map) \
454 v = bpf_map_lookup_percpu_elem(&map, &key, 0); \
455 if (!v) \
456 return -1; \
457 ret = test_map_kptr_ref_post(v); \
458 if (ret) \
459 return ret;
460
461 SEC("tc")
test_map_kptr_ref2(struct __sk_buff * ctx)462 int test_map_kptr_ref2(struct __sk_buff *ctx)
463 {
464 struct map_value *v;
465 int key = 0, ret;
466
467 TEST(array_map);
468 TEST(hash_map);
469 TEST(hash_malloc_map);
470 TEST(lru_hash_map);
471
472 TEST_PCPU(pcpu_array_map);
473 TEST_PCPU(pcpu_hash_map);
474 TEST_PCPU(pcpu_hash_malloc_map);
475 TEST_PCPU(lru_pcpu_hash_map);
476
477 return 0;
478 }
479
480 #undef TEST
481 #undef TEST_PCPU
482
483 SEC("tc")
test_map_kptr_ref3(struct __sk_buff * ctx)484 int test_map_kptr_ref3(struct __sk_buff *ctx)
485 {
486 struct prog_test_ref_kfunc *p;
487 unsigned long sp = 0;
488
489 p = bpf_kfunc_call_test_acquire(&sp);
490 if (!p)
491 return 1;
492 ref++;
493 if (p->cnt.refs.counter != ref) {
494 bpf_kfunc_call_test_release(p);
495 return 2;
496 }
497 bpf_kfunc_call_test_release(p);
498 ref--;
499 return 0;
500 }
501
502 int num_of_refs;
503
read_ref_count(void)504 static __always_inline int read_ref_count(void)
505 {
506 struct prog_test_ref_kfunc *p;
507 unsigned long arg = 0;
508
509 p = bpf_kfunc_call_test_acquire(&arg);
510 if (!p)
511 return 1;
512
513 num_of_refs = p->cnt.refs.counter;
514 bpf_kfunc_call_test_release(p);
515 return 0;
516 }
517
518 SEC("syscall")
count_ref(void * ctx)519 int count_ref(void *ctx)
520 {
521 return read_ref_count();
522 }
523
stash_ref_ptr(struct map_value * v)524 static __always_inline int stash_ref_ptr(struct map_value *v)
525 {
526 struct prog_test_ref_kfunc *p, *old;
527 unsigned long arg = 0;
528
529 p = bpf_kfunc_call_test_acquire(&arg);
530 if (!p)
531 return 1;
532
533 old = bpf_kptr_xchg(&v->ref_ptr, p);
534 if (old) {
535 bpf_kfunc_call_test_release(old);
536 old = bpf_kptr_xchg(&v->ref_ptr, NULL);
537 if (old)
538 bpf_kfunc_call_test_release(old);
539 return 2;
540 }
541 return 0;
542 }
543
check_refs(int expected)544 static __always_inline int check_refs(int expected)
545 {
546 int ret;
547
548 ret = read_ref_count();
549 if (ret)
550 return ret;
551 return num_of_refs == expected ? 0 : 3;
552 }
553
554 SEC("syscall")
test_array_map_update_kptr(void * ctx)555 int test_array_map_update_kptr(void *ctx)
556 {
557 struct map_value init = {}, *v;
558 int key = 0, ret;
559
560 v = bpf_map_lookup_elem(&array_map, &key);
561 if (!v)
562 return 1;
563 ret = stash_ref_ptr(v);
564 if (ret)
565 return ret;
566 ret = check_refs(3);
567 if (ret)
568 return ret;
569 ret = bpf_map_update_elem(&array_map, &key, &init, BPF_EXIST);
570 if (ret)
571 return 4;
572 return check_refs(3);
573 }
574
575 #define DEFINE_HASH_UPDATE_KPTR_TEST(name, map) \
576 SEC("syscall") \
577 int name(void *ctx) \
578 { \
579 struct map_value init = {}, *v; \
580 int key = 0, ret; \
581 \
582 ret = bpf_map_update_elem(&map, &key, &init, BPF_NOEXIST); \
583 if (ret) \
584 return 1; \
585 v = bpf_map_lookup_elem(&map, &key); \
586 if (!v) \
587 return 2; \
588 ret = stash_ref_ptr(v); \
589 if (ret) \
590 return ret; \
591 ret = check_refs(3); \
592 if (ret) \
593 return ret; \
594 ret = bpf_map_update_elem(&map, &key, &init, BPF_EXIST); \
595 if (ret) \
596 return 4; \
597 return check_refs(3); \
598 }
599
DEFINE_HASH_UPDATE_KPTR_TEST(test_hash_map_update_kptr,hash_map)600 DEFINE_HASH_UPDATE_KPTR_TEST(test_hash_map_update_kptr, hash_map)
601 DEFINE_HASH_UPDATE_KPTR_TEST(test_hash_malloc_map_update_kptr, hash_malloc_map)
602
603 SEC("syscall")
604 int test_ls_map_kptr_ref1(void *ctx)
605 {
606 struct task_struct *current;
607 struct map_value *v;
608
609 current = bpf_get_current_task_btf();
610 if (!current)
611 return 100;
612 v = bpf_task_storage_get(&task_ls_map, current, NULL, 0);
613 if (v)
614 return 150;
615 v = bpf_task_storage_get(&task_ls_map, current, NULL, BPF_LOCAL_STORAGE_GET_F_CREATE);
616 if (!v)
617 return 200;
618 return test_map_kptr_ref_pre(v);
619 }
620
621 SEC("syscall")
test_ls_map_kptr_ref2(void * ctx)622 int test_ls_map_kptr_ref2(void *ctx)
623 {
624 struct task_struct *current;
625 struct map_value *v;
626
627 current = bpf_get_current_task_btf();
628 if (!current)
629 return 100;
630 v = bpf_task_storage_get(&task_ls_map, current, NULL, 0);
631 if (!v)
632 return 200;
633 return test_map_kptr_ref_post(v);
634 }
635
636 SEC("syscall")
test_ls_map_kptr_ref_del(void * ctx)637 int test_ls_map_kptr_ref_del(void *ctx)
638 {
639 struct task_struct *current;
640 struct map_value *v;
641
642 current = bpf_get_current_task_btf();
643 if (!current)
644 return 100;
645 v = bpf_task_storage_get(&task_ls_map, current, NULL, 0);
646 if (!v)
647 return 200;
648 if (!v->ref_ptr)
649 return 300;
650 return bpf_task_storage_delete(&task_ls_map, current);
651 }
652
653 char _license[] SEC("license") = "GPL";
654