1 // SPDX-License-Identifier: GPL-2.0
2 /* Copyright (c) 2021 Facebook */
3 #include <linux/bpf.h>
4 #include <time.h>
5 #include <errno.h>
6 #include <bpf/bpf_helpers.h>
7 #include "bpf_tcp_helpers.h"
8 
9 char _license[] SEC("license") = "GPL";
10 struct hmap_elem {
11 	int counter;
12 	struct bpf_timer timer;
13 	struct bpf_spin_lock lock; /* unused */
14 };
15 
16 struct {
17 	__uint(type, BPF_MAP_TYPE_HASH);
18 	__uint(max_entries, 1000);
19 	__type(key, int);
20 	__type(value, struct hmap_elem);
21 } hmap SEC(".maps");
22 
23 struct {
24 	__uint(type, BPF_MAP_TYPE_HASH);
25 	__uint(map_flags, BPF_F_NO_PREALLOC);
26 	__uint(max_entries, 1000);
27 	__type(key, int);
28 	__type(value, struct hmap_elem);
29 } hmap_malloc SEC(".maps");
30 
31 struct elem {
32 	struct bpf_timer t;
33 };
34 
35 struct {
36 	__uint(type, BPF_MAP_TYPE_ARRAY);
37 	__uint(max_entries, 2);
38 	__type(key, int);
39 	__type(value, struct elem);
40 } array SEC(".maps");
41 
42 struct {
43 	__uint(type, BPF_MAP_TYPE_LRU_HASH);
44 	__uint(max_entries, 4);
45 	__type(key, int);
46 	__type(value, struct elem);
47 } lru SEC(".maps");
48 
49 __u64 bss_data;
50 __u64 err;
51 __u64 ok;
52 __u64 callback_check = 52;
53 __u64 callback2_check = 52;
54 
55 #define ARRAY 1
56 #define HTAB 2
57 #define HTAB_MALLOC 3
58 #define LRU 4
59 
60 /* callback for array and lru timers */
61 static int timer_cb1(void *map, int *key, struct bpf_timer *timer)
62 {
63 	/* increment bss variable twice.
64 	 * Once via array timer callback and once via lru timer callback
65 	 */
66 	bss_data += 5;
67 
68 	/* *key == 0 - the callback was called for array timer.
69 	 * *key == 4 - the callback was called from lru timer.
70 	 */
71 	if (*key == ARRAY) {
72 		struct bpf_timer *lru_timer;
73 		int lru_key = LRU;
74 
75 		/* rearm array timer to be called again in ~35 seconds */
76 		if (bpf_timer_start(timer, 1ull << 35, 0) != 0)
77 			err |= 1;
78 
79 		lru_timer = bpf_map_lookup_elem(&lru, &lru_key);
80 		if (!lru_timer)
81 			return 0;
82 		bpf_timer_set_callback(lru_timer, timer_cb1);
83 		if (bpf_timer_start(lru_timer, 0, 0) != 0)
84 			err |= 2;
85 	} else if (*key == LRU) {
86 		int lru_key, i;
87 
88 		for (i = LRU + 1;
89 		     i <= 100  /* for current LRU eviction algorithm this number
90 				* should be larger than ~ lru->max_entries * 2
91 				*/;
92 		     i++) {
93 			struct elem init = {};
94 
95 			/* lru_key cannot be used as loop induction variable
96 			 * otherwise the loop will be unbounded.
97 			 */
98 			lru_key = i;
99 
100 			/* add more elements into lru map to push out current
101 			 * element and force deletion of this timer
102 			 */
103 			bpf_map_update_elem(map, &lru_key, &init, 0);
104 			/* look it up to bump it into active list */
105 			bpf_map_lookup_elem(map, &lru_key);
106 
107 			/* keep adding until *key changes underneath,
108 			 * which means that key/timer memory was reused
109 			 */
110 			if (*key != LRU)
111 				break;
112 		}
113 
114 		/* check that the timer was removed */
115 		if (bpf_timer_cancel(timer) != -EINVAL)
116 			err |= 4;
117 		ok |= 1;
118 	}
119 	return 0;
120 }
121 
122 SEC("fentry/bpf_fentry_test1")
123 int BPF_PROG(test1, int a)
124 {
125 	struct bpf_timer *arr_timer, *lru_timer;
126 	struct elem init = {};
127 	int lru_key = LRU;
128 	int array_key = ARRAY;
129 
130 	arr_timer = bpf_map_lookup_elem(&array, &array_key);
131 	if (!arr_timer)
132 		return 0;
133 	bpf_timer_init(arr_timer, &array, CLOCK_MONOTONIC);
134 
135 	bpf_map_update_elem(&lru, &lru_key, &init, 0);
136 	lru_timer = bpf_map_lookup_elem(&lru, &lru_key);
137 	if (!lru_timer)
138 		return 0;
139 	bpf_timer_init(lru_timer, &lru, CLOCK_MONOTONIC);
140 
141 	bpf_timer_set_callback(arr_timer, timer_cb1);
142 	bpf_timer_start(arr_timer, 0 /* call timer_cb1 asap */, 0);
143 
144 	/* init more timers to check that array destruction
145 	 * doesn't leak timer memory.
146 	 */
147 	array_key = 0;
148 	arr_timer = bpf_map_lookup_elem(&array, &array_key);
149 	if (!arr_timer)
150 		return 0;
151 	bpf_timer_init(arr_timer, &array, CLOCK_MONOTONIC);
152 	return 0;
153 }
154 
155 /* callback for prealloc and non-prealloca hashtab timers */
156 static int timer_cb2(void *map, int *key, struct hmap_elem *val)
157 {
158 	if (*key == HTAB)
159 		callback_check--;
160 	else
161 		callback2_check--;
162 	if (val->counter > 0 && --val->counter) {
163 		/* re-arm the timer again to execute after 1 usec */
164 		bpf_timer_start(&val->timer, 1000, 0);
165 	} else if (*key == HTAB) {
166 		struct bpf_timer *arr_timer;
167 		int array_key = ARRAY;
168 
169 		/* cancel arr_timer otherwise bpf_fentry_test1 prog
170 		 * will stay alive forever.
171 		 */
172 		arr_timer = bpf_map_lookup_elem(&array, &array_key);
173 		if (!arr_timer)
174 			return 0;
175 		if (bpf_timer_cancel(arr_timer) != 1)
176 			/* bpf_timer_cancel should return 1 to indicate
177 			 * that arr_timer was active at this time
178 			 */
179 			err |= 8;
180 
181 		/* try to cancel ourself. It shouldn't deadlock. */
182 		if (bpf_timer_cancel(&val->timer) != -EDEADLK)
183 			err |= 16;
184 
185 		/* delete this key and this timer anyway.
186 		 * It shouldn't deadlock either.
187 		 */
188 		bpf_map_delete_elem(map, key);
189 
190 		/* in preallocated hashmap both 'key' and 'val' could have been
191 		 * reused to store another map element (like in LRU above),
192 		 * but in controlled test environment the below test works.
193 		 * It's not a use-after-free. The memory is owned by the map.
194 		 */
195 		if (bpf_timer_start(&val->timer, 1000, 0) != -EINVAL)
196 			err |= 32;
197 		ok |= 2;
198 	} else {
199 		if (*key != HTAB_MALLOC)
200 			err |= 64;
201 
202 		/* try to cancel ourself. It shouldn't deadlock. */
203 		if (bpf_timer_cancel(&val->timer) != -EDEADLK)
204 			err |= 128;
205 
206 		/* delete this key and this timer anyway.
207 		 * It shouldn't deadlock either.
208 		 */
209 		bpf_map_delete_elem(map, key);
210 
211 		/* in non-preallocated hashmap both 'key' and 'val' are RCU
212 		 * protected and still valid though this element was deleted
213 		 * from the map. Arm this timer for ~35 seconds. When callback
214 		 * finishes the call_rcu will invoke:
215 		 * htab_elem_free_rcu
216 		 *   check_and_free_timer
217 		 *     bpf_timer_cancel_and_free
218 		 * to cancel this 35 second sleep and delete the timer for real.
219 		 */
220 		if (bpf_timer_start(&val->timer, 1ull << 35, 0) != 0)
221 			err |= 256;
222 		ok |= 4;
223 	}
224 	return 0;
225 }
226 
227 int bpf_timer_test(void)
228 {
229 	struct hmap_elem *val;
230 	int key = HTAB, key_malloc = HTAB_MALLOC;
231 
232 	val = bpf_map_lookup_elem(&hmap, &key);
233 	if (val) {
234 		if (bpf_timer_init(&val->timer, &hmap, CLOCK_BOOTTIME) != 0)
235 			err |= 512;
236 		bpf_timer_set_callback(&val->timer, timer_cb2);
237 		bpf_timer_start(&val->timer, 1000, 0);
238 	}
239 	val = bpf_map_lookup_elem(&hmap_malloc, &key_malloc);
240 	if (val) {
241 		if (bpf_timer_init(&val->timer, &hmap_malloc, CLOCK_BOOTTIME) != 0)
242 			err |= 1024;
243 		bpf_timer_set_callback(&val->timer, timer_cb2);
244 		bpf_timer_start(&val->timer, 1000, 0);
245 	}
246 	return 0;
247 }
248 
249 SEC("fentry/bpf_fentry_test2")
250 int BPF_PROG(test2, int a, int b)
251 {
252 	struct hmap_elem init = {}, *val;
253 	int key = HTAB, key_malloc = HTAB_MALLOC;
254 
255 	init.counter = 10; /* number of times to trigger timer_cb2 */
256 	bpf_map_update_elem(&hmap, &key, &init, 0);
257 	val = bpf_map_lookup_elem(&hmap, &key);
258 	if (val)
259 		bpf_timer_init(&val->timer, &hmap, CLOCK_BOOTTIME);
260 	/* update the same key to free the timer */
261 	bpf_map_update_elem(&hmap, &key, &init, 0);
262 
263 	bpf_map_update_elem(&hmap_malloc, &key_malloc, &init, 0);
264 	val = bpf_map_lookup_elem(&hmap_malloc, &key_malloc);
265 	if (val)
266 		bpf_timer_init(&val->timer, &hmap_malloc, CLOCK_BOOTTIME);
267 	/* update the same key to free the timer */
268 	bpf_map_update_elem(&hmap_malloc, &key_malloc, &init, 0);
269 
270 	/* init more timers to check that htab operations
271 	 * don't leak timer memory.
272 	 */
273 	key = 0;
274 	bpf_map_update_elem(&hmap, &key, &init, 0);
275 	val = bpf_map_lookup_elem(&hmap, &key);
276 	if (val)
277 		bpf_timer_init(&val->timer, &hmap, CLOCK_BOOTTIME);
278 	bpf_map_delete_elem(&hmap, &key);
279 	bpf_map_update_elem(&hmap, &key, &init, 0);
280 	val = bpf_map_lookup_elem(&hmap, &key);
281 	if (val)
282 		bpf_timer_init(&val->timer, &hmap, CLOCK_BOOTTIME);
283 
284 	/* and with non-prealloc htab */
285 	key_malloc = 0;
286 	bpf_map_update_elem(&hmap_malloc, &key_malloc, &init, 0);
287 	val = bpf_map_lookup_elem(&hmap_malloc, &key_malloc);
288 	if (val)
289 		bpf_timer_init(&val->timer, &hmap_malloc, CLOCK_BOOTTIME);
290 	bpf_map_delete_elem(&hmap_malloc, &key_malloc);
291 	bpf_map_update_elem(&hmap_malloc, &key_malloc, &init, 0);
292 	val = bpf_map_lookup_elem(&hmap_malloc, &key_malloc);
293 	if (val)
294 		bpf_timer_init(&val->timer, &hmap_malloc, CLOCK_BOOTTIME);
295 
296 	return bpf_timer_test();
297 }
298