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 struct {
50 	__uint(type, BPF_MAP_TYPE_ARRAY);
51 	__uint(max_entries, 1);
52 	__type(key, int);
53 	__type(value, struct elem);
54 } abs_timer SEC(".maps");
55 
56 __u64 bss_data;
57 __u64 abs_data;
58 __u64 err;
59 __u64 ok;
60 __u64 callback_check = 52;
61 __u64 callback2_check = 52;
62 
63 #define ARRAY 1
64 #define HTAB 2
65 #define HTAB_MALLOC 3
66 #define LRU 4
67 
68 /* callback for array and lru timers */
timer_cb1(void * map,int * key,struct bpf_timer * timer)69 static int timer_cb1(void *map, int *key, struct bpf_timer *timer)
70 {
71 	/* increment bss variable twice.
72 	 * Once via array timer callback and once via lru timer callback
73 	 */
74 	bss_data += 5;
75 
76 	/* *key == 0 - the callback was called for array timer.
77 	 * *key == 4 - the callback was called from lru timer.
78 	 */
79 	if (*key == ARRAY) {
80 		struct bpf_timer *lru_timer;
81 		int lru_key = LRU;
82 
83 		/* rearm array timer to be called again in ~35 seconds */
84 		if (bpf_timer_start(timer, 1ull << 35, 0) != 0)
85 			err |= 1;
86 
87 		lru_timer = bpf_map_lookup_elem(&lru, &lru_key);
88 		if (!lru_timer)
89 			return 0;
90 		bpf_timer_set_callback(lru_timer, timer_cb1);
91 		if (bpf_timer_start(lru_timer, 0, 0) != 0)
92 			err |= 2;
93 	} else if (*key == LRU) {
94 		int lru_key, i;
95 
96 		for (i = LRU + 1;
97 		     i <= 100  /* for current LRU eviction algorithm this number
98 				* should be larger than ~ lru->max_entries * 2
99 				*/;
100 		     i++) {
101 			struct elem init = {};
102 
103 			/* lru_key cannot be used as loop induction variable
104 			 * otherwise the loop will be unbounded.
105 			 */
106 			lru_key = i;
107 
108 			/* add more elements into lru map to push out current
109 			 * element and force deletion of this timer
110 			 */
111 			bpf_map_update_elem(map, &lru_key, &init, 0);
112 			/* look it up to bump it into active list */
113 			bpf_map_lookup_elem(map, &lru_key);
114 
115 			/* keep adding until *key changes underneath,
116 			 * which means that key/timer memory was reused
117 			 */
118 			if (*key != LRU)
119 				break;
120 		}
121 
122 		/* check that the timer was removed */
123 		if (bpf_timer_cancel(timer) != -EINVAL)
124 			err |= 4;
125 		ok |= 1;
126 	}
127 	return 0;
128 }
129 
130 SEC("fentry/bpf_fentry_test1")
BPF_PROG2(test1,int,a)131 int BPF_PROG2(test1, int, a)
132 {
133 	struct bpf_timer *arr_timer, *lru_timer;
134 	struct elem init = {};
135 	int lru_key = LRU;
136 	int array_key = ARRAY;
137 
138 	arr_timer = bpf_map_lookup_elem(&array, &array_key);
139 	if (!arr_timer)
140 		return 0;
141 	bpf_timer_init(arr_timer, &array, CLOCK_MONOTONIC);
142 
143 	bpf_map_update_elem(&lru, &lru_key, &init, 0);
144 	lru_timer = bpf_map_lookup_elem(&lru, &lru_key);
145 	if (!lru_timer)
146 		return 0;
147 	bpf_timer_init(lru_timer, &lru, CLOCK_MONOTONIC);
148 
149 	bpf_timer_set_callback(arr_timer, timer_cb1);
150 	bpf_timer_start(arr_timer, 0 /* call timer_cb1 asap */, 0);
151 
152 	/* init more timers to check that array destruction
153 	 * doesn't leak timer memory.
154 	 */
155 	array_key = 0;
156 	arr_timer = bpf_map_lookup_elem(&array, &array_key);
157 	if (!arr_timer)
158 		return 0;
159 	bpf_timer_init(arr_timer, &array, CLOCK_MONOTONIC);
160 	return 0;
161 }
162 
163 /* callback for prealloc and non-prealloca hashtab timers */
timer_cb2(void * map,int * key,struct hmap_elem * val)164 static int timer_cb2(void *map, int *key, struct hmap_elem *val)
165 {
166 	if (*key == HTAB)
167 		callback_check--;
168 	else
169 		callback2_check--;
170 	if (val->counter > 0 && --val->counter) {
171 		/* re-arm the timer again to execute after 1 usec */
172 		bpf_timer_start(&val->timer, 1000, 0);
173 	} else if (*key == HTAB) {
174 		struct bpf_timer *arr_timer;
175 		int array_key = ARRAY;
176 
177 		/* cancel arr_timer otherwise bpf_fentry_test1 prog
178 		 * will stay alive forever.
179 		 */
180 		arr_timer = bpf_map_lookup_elem(&array, &array_key);
181 		if (!arr_timer)
182 			return 0;
183 		if (bpf_timer_cancel(arr_timer) != 1)
184 			/* bpf_timer_cancel should return 1 to indicate
185 			 * that arr_timer was active at this time
186 			 */
187 			err |= 8;
188 
189 		/* try to cancel ourself. It shouldn't deadlock. */
190 		if (bpf_timer_cancel(&val->timer) != -EDEADLK)
191 			err |= 16;
192 
193 		/* delete this key and this timer anyway.
194 		 * It shouldn't deadlock either.
195 		 */
196 		bpf_map_delete_elem(map, key);
197 
198 		/* in preallocated hashmap both 'key' and 'val' could have been
199 		 * reused to store another map element (like in LRU above),
200 		 * but in controlled test environment the below test works.
201 		 * It's not a use-after-free. The memory is owned by the map.
202 		 */
203 		if (bpf_timer_start(&val->timer, 1000, 0) != -EINVAL)
204 			err |= 32;
205 		ok |= 2;
206 	} else {
207 		if (*key != HTAB_MALLOC)
208 			err |= 64;
209 
210 		/* try to cancel ourself. It shouldn't deadlock. */
211 		if (bpf_timer_cancel(&val->timer) != -EDEADLK)
212 			err |= 128;
213 
214 		/* delete this key and this timer anyway.
215 		 * It shouldn't deadlock either.
216 		 */
217 		bpf_map_delete_elem(map, key);
218 
219 		ok |= 4;
220 	}
221 	return 0;
222 }
223 
bpf_timer_test(void)224 int bpf_timer_test(void)
225 {
226 	struct hmap_elem *val;
227 	int key = HTAB, key_malloc = HTAB_MALLOC;
228 
229 	val = bpf_map_lookup_elem(&hmap, &key);
230 	if (val) {
231 		if (bpf_timer_init(&val->timer, &hmap, CLOCK_BOOTTIME) != 0)
232 			err |= 512;
233 		bpf_timer_set_callback(&val->timer, timer_cb2);
234 		bpf_timer_start(&val->timer, 1000, 0);
235 	}
236 	val = bpf_map_lookup_elem(&hmap_malloc, &key_malloc);
237 	if (val) {
238 		if (bpf_timer_init(&val->timer, &hmap_malloc, CLOCK_BOOTTIME) != 0)
239 			err |= 1024;
240 		bpf_timer_set_callback(&val->timer, timer_cb2);
241 		bpf_timer_start(&val->timer, 1000, 0);
242 	}
243 	return 0;
244 }
245 
246 SEC("fentry/bpf_fentry_test2")
BPF_PROG2(test2,int,a,int,b)247 int BPF_PROG2(test2, int, a, int, b)
248 {
249 	struct hmap_elem init = {}, *val;
250 	int key = HTAB, key_malloc = HTAB_MALLOC;
251 
252 	init.counter = 10; /* number of times to trigger timer_cb2 */
253 	bpf_map_update_elem(&hmap, &key, &init, 0);
254 	val = bpf_map_lookup_elem(&hmap, &key);
255 	if (val)
256 		bpf_timer_init(&val->timer, &hmap, CLOCK_BOOTTIME);
257 	/* update the same key to free the timer */
258 	bpf_map_update_elem(&hmap, &key, &init, 0);
259 
260 	bpf_map_update_elem(&hmap_malloc, &key_malloc, &init, 0);
261 	val = bpf_map_lookup_elem(&hmap_malloc, &key_malloc);
262 	if (val)
263 		bpf_timer_init(&val->timer, &hmap_malloc, CLOCK_BOOTTIME);
264 	/* update the same key to free the timer */
265 	bpf_map_update_elem(&hmap_malloc, &key_malloc, &init, 0);
266 
267 	/* init more timers to check that htab operations
268 	 * don't leak timer memory.
269 	 */
270 	key = 0;
271 	bpf_map_update_elem(&hmap, &key, &init, 0);
272 	val = bpf_map_lookup_elem(&hmap, &key);
273 	if (val)
274 		bpf_timer_init(&val->timer, &hmap, CLOCK_BOOTTIME);
275 	bpf_map_delete_elem(&hmap, &key);
276 	bpf_map_update_elem(&hmap, &key, &init, 0);
277 	val = bpf_map_lookup_elem(&hmap, &key);
278 	if (val)
279 		bpf_timer_init(&val->timer, &hmap, CLOCK_BOOTTIME);
280 
281 	/* and with non-prealloc htab */
282 	key_malloc = 0;
283 	bpf_map_update_elem(&hmap_malloc, &key_malloc, &init, 0);
284 	val = bpf_map_lookup_elem(&hmap_malloc, &key_malloc);
285 	if (val)
286 		bpf_timer_init(&val->timer, &hmap_malloc, CLOCK_BOOTTIME);
287 	bpf_map_delete_elem(&hmap_malloc, &key_malloc);
288 	bpf_map_update_elem(&hmap_malloc, &key_malloc, &init, 0);
289 	val = bpf_map_lookup_elem(&hmap_malloc, &key_malloc);
290 	if (val)
291 		bpf_timer_init(&val->timer, &hmap_malloc, CLOCK_BOOTTIME);
292 
293 	return bpf_timer_test();
294 }
295 
296 /* callback for absolute timer */
timer_cb3(void * map,int * key,struct bpf_timer * timer)297 static int timer_cb3(void *map, int *key, struct bpf_timer *timer)
298 {
299 	abs_data += 6;
300 
301 	if (abs_data < 12) {
302 		bpf_timer_start(timer, bpf_ktime_get_boot_ns() + 1000,
303 				BPF_F_TIMER_ABS);
304 	} else {
305 		/* Re-arm timer ~35 seconds in future */
306 		bpf_timer_start(timer, bpf_ktime_get_boot_ns() + (1ull << 35),
307 				BPF_F_TIMER_ABS);
308 	}
309 
310 	return 0;
311 }
312 
313 SEC("fentry/bpf_fentry_test3")
BPF_PROG2(test3,int,a)314 int BPF_PROG2(test3, int, a)
315 {
316 	int key = 0;
317 	struct bpf_timer *timer;
318 
319 	bpf_printk("test3");
320 
321 	timer = bpf_map_lookup_elem(&abs_timer, &key);
322 	if (timer) {
323 		if (bpf_timer_init(timer, &abs_timer, CLOCK_BOOTTIME) != 0)
324 			err |= 2048;
325 		bpf_timer_set_callback(timer, timer_cb3);
326 		bpf_timer_start(timer, bpf_ktime_get_boot_ns() + 1000,
327 				BPF_F_TIMER_ABS);
328 	}
329 
330 	return 0;
331 }
332