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