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
9char _license[] SEC("license") = "GPL";
10struct hmap_elem {
11	int counter;
12	struct bpf_timer timer;
13	struct bpf_spin_lock lock; /* unused */
14};
15
16struct {
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
23struct {
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
31struct elem {
32	struct bpf_timer t;
33};
34
35struct {
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
42struct {
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
49struct {
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 */
69static 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
130SEC("fentry/bpf_fentry_test1")
131int 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 */
164static 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
224int 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
246SEC("fentry/bpf_fentry_test2")
247int 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 */
297static 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
313SEC("fentry/bpf_fentry_test3")
314int 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