1 // SPDX-License-Identifier: GPL-2.0
2 /* Copyright (C) 2023. Huawei Technologies Co., Ltd */
3 #include <vmlinux.h>
4 #include <bpf/bpf_tracing.h>
5 #include <bpf/bpf_helpers.h>
6 
7 #include "bpf_experimental.h"
8 #include "bpf_misc.h"
9 
10 #ifndef ARRAY_SIZE
11 #define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0]))
12 #endif
13 
14 struct generic_map_value {
15 	void *data;
16 };
17 
18 char _license[] SEC("license") = "GPL";
19 
20 const unsigned int data_sizes[] = {8, 16, 32, 64, 96, 128, 192, 256, 512, 1024, 2048, 4096};
21 const volatile unsigned int data_btf_ids[ARRAY_SIZE(data_sizes)] = {};
22 
23 int err = 0;
24 int pid = 0;
25 
26 #define DEFINE_ARRAY_WITH_KPTR(_size) \
27 	struct bin_data_##_size { \
28 		char data[_size - sizeof(void *)]; \
29 	}; \
30 	struct map_value_##_size { \
31 		struct bin_data_##_size __kptr * data; \
32 		/* To emit BTF info for bin_data_xx */ \
33 		struct bin_data_##_size not_used; \
34 	}; \
35 	struct { \
36 		__uint(type, BPF_MAP_TYPE_ARRAY); \
37 		__type(key, int); \
38 		__type(value, struct map_value_##_size); \
39 		__uint(max_entries, 128); \
40 	} array_##_size SEC(".maps");
41 
42 static __always_inline void batch_alloc_free(struct bpf_map *map, unsigned int batch,
43 					     unsigned int idx)
44 {
45 	struct generic_map_value *value;
46 	unsigned int i, key;
47 	void *old, *new;
48 
49 	for (i = 0; i < batch; i++) {
50 		key = i;
51 		value = bpf_map_lookup_elem(map, &key);
52 		if (!value) {
53 			err = 1;
54 			return;
55 		}
56 		new = bpf_obj_new_impl(data_btf_ids[idx], NULL);
57 		if (!new) {
58 			err = 2;
59 			return;
60 		}
61 		old = bpf_kptr_xchg(&value->data, new);
62 		if (old) {
63 			bpf_obj_drop(old);
64 			err = 3;
65 			return;
66 		}
67 	}
68 	for (i = 0; i < batch; i++) {
69 		key = i;
70 		value = bpf_map_lookup_elem(map, &key);
71 		if (!value) {
72 			err = 4;
73 			return;
74 		}
75 		old = bpf_kptr_xchg(&value->data, NULL);
76 		if (!old) {
77 			err = 5;
78 			return;
79 		}
80 		bpf_obj_drop(old);
81 	}
82 }
83 
84 #define CALL_BATCH_ALLOC_FREE(size, batch, idx) \
85 	batch_alloc_free((struct bpf_map *)(&array_##size), batch, idx)
86 
87 DEFINE_ARRAY_WITH_KPTR(8);
88 DEFINE_ARRAY_WITH_KPTR(16);
89 DEFINE_ARRAY_WITH_KPTR(32);
90 DEFINE_ARRAY_WITH_KPTR(64);
91 DEFINE_ARRAY_WITH_KPTR(96);
92 DEFINE_ARRAY_WITH_KPTR(128);
93 DEFINE_ARRAY_WITH_KPTR(192);
94 DEFINE_ARRAY_WITH_KPTR(256);
95 DEFINE_ARRAY_WITH_KPTR(512);
96 DEFINE_ARRAY_WITH_KPTR(1024);
97 DEFINE_ARRAY_WITH_KPTR(2048);
98 DEFINE_ARRAY_WITH_KPTR(4096);
99 
100 SEC("fentry/" SYS_PREFIX "sys_nanosleep")
101 int test_bpf_mem_alloc_free(void *ctx)
102 {
103 	if ((u32)bpf_get_current_pid_tgid() != pid)
104 		return 0;
105 
106 	/* Alloc 128 8-bytes objects in batch to trigger refilling,
107 	 * then free 128 8-bytes objects in batch to trigger freeing.
108 	 */
109 	CALL_BATCH_ALLOC_FREE(8, 128, 0);
110 	CALL_BATCH_ALLOC_FREE(16, 128, 1);
111 	CALL_BATCH_ALLOC_FREE(32, 128, 2);
112 	CALL_BATCH_ALLOC_FREE(64, 128, 3);
113 	CALL_BATCH_ALLOC_FREE(96, 128, 4);
114 	CALL_BATCH_ALLOC_FREE(128, 128, 5);
115 	CALL_BATCH_ALLOC_FREE(192, 128, 6);
116 	CALL_BATCH_ALLOC_FREE(256, 128, 7);
117 	CALL_BATCH_ALLOC_FREE(512, 64, 8);
118 	CALL_BATCH_ALLOC_FREE(1024, 32, 9);
119 	CALL_BATCH_ALLOC_FREE(2048, 16, 10);
120 	CALL_BATCH_ALLOC_FREE(4096, 8, 11);
121 
122 	return 0;
123 }
124