162306a36Sopenharmony_ci// SPDX-License-Identifier: GPL-2.0
262306a36Sopenharmony_ci/* Copyright (c) 2021 Facebook */
362306a36Sopenharmony_ci#include <linux/bpf.h>
462306a36Sopenharmony_ci#include <time.h>
562306a36Sopenharmony_ci#include <errno.h>
662306a36Sopenharmony_ci#include <bpf/bpf_helpers.h>
762306a36Sopenharmony_ci#include "bpf_tcp_helpers.h"
862306a36Sopenharmony_ci
962306a36Sopenharmony_cichar _license[] SEC("license") = "GPL";
1062306a36Sopenharmony_cistruct hmap_elem {
1162306a36Sopenharmony_ci	int counter;
1262306a36Sopenharmony_ci	struct bpf_timer timer;
1362306a36Sopenharmony_ci	struct bpf_spin_lock lock; /* unused */
1462306a36Sopenharmony_ci};
1562306a36Sopenharmony_ci
1662306a36Sopenharmony_cistruct {
1762306a36Sopenharmony_ci	__uint(type, BPF_MAP_TYPE_HASH);
1862306a36Sopenharmony_ci	__uint(max_entries, 1000);
1962306a36Sopenharmony_ci	__type(key, int);
2062306a36Sopenharmony_ci	__type(value, struct hmap_elem);
2162306a36Sopenharmony_ci} hmap SEC(".maps");
2262306a36Sopenharmony_ci
2362306a36Sopenharmony_cistruct {
2462306a36Sopenharmony_ci	__uint(type, BPF_MAP_TYPE_HASH);
2562306a36Sopenharmony_ci	__uint(map_flags, BPF_F_NO_PREALLOC);
2662306a36Sopenharmony_ci	__uint(max_entries, 1000);
2762306a36Sopenharmony_ci	__type(key, int);
2862306a36Sopenharmony_ci	__type(value, struct hmap_elem);
2962306a36Sopenharmony_ci} hmap_malloc SEC(".maps");
3062306a36Sopenharmony_ci
3162306a36Sopenharmony_cistruct elem {
3262306a36Sopenharmony_ci	struct bpf_timer t;
3362306a36Sopenharmony_ci};
3462306a36Sopenharmony_ci
3562306a36Sopenharmony_cistruct {
3662306a36Sopenharmony_ci	__uint(type, BPF_MAP_TYPE_ARRAY);
3762306a36Sopenharmony_ci	__uint(max_entries, 2);
3862306a36Sopenharmony_ci	__type(key, int);
3962306a36Sopenharmony_ci	__type(value, struct elem);
4062306a36Sopenharmony_ci} array SEC(".maps");
4162306a36Sopenharmony_ci
4262306a36Sopenharmony_cistruct {
4362306a36Sopenharmony_ci	__uint(type, BPF_MAP_TYPE_LRU_HASH);
4462306a36Sopenharmony_ci	__uint(max_entries, 4);
4562306a36Sopenharmony_ci	__type(key, int);
4662306a36Sopenharmony_ci	__type(value, struct elem);
4762306a36Sopenharmony_ci} lru SEC(".maps");
4862306a36Sopenharmony_ci
4962306a36Sopenharmony_cistruct {
5062306a36Sopenharmony_ci	__uint(type, BPF_MAP_TYPE_ARRAY);
5162306a36Sopenharmony_ci	__uint(max_entries, 1);
5262306a36Sopenharmony_ci	__type(key, int);
5362306a36Sopenharmony_ci	__type(value, struct elem);
5462306a36Sopenharmony_ci} abs_timer SEC(".maps");
5562306a36Sopenharmony_ci
5662306a36Sopenharmony_ci__u64 bss_data;
5762306a36Sopenharmony_ci__u64 abs_data;
5862306a36Sopenharmony_ci__u64 err;
5962306a36Sopenharmony_ci__u64 ok;
6062306a36Sopenharmony_ci__u64 callback_check = 52;
6162306a36Sopenharmony_ci__u64 callback2_check = 52;
6262306a36Sopenharmony_ci
6362306a36Sopenharmony_ci#define ARRAY 1
6462306a36Sopenharmony_ci#define HTAB 2
6562306a36Sopenharmony_ci#define HTAB_MALLOC 3
6662306a36Sopenharmony_ci#define LRU 4
6762306a36Sopenharmony_ci
6862306a36Sopenharmony_ci/* callback for array and lru timers */
6962306a36Sopenharmony_cistatic int timer_cb1(void *map, int *key, struct bpf_timer *timer)
7062306a36Sopenharmony_ci{
7162306a36Sopenharmony_ci	/* increment bss variable twice.
7262306a36Sopenharmony_ci	 * Once via array timer callback and once via lru timer callback
7362306a36Sopenharmony_ci	 */
7462306a36Sopenharmony_ci	bss_data += 5;
7562306a36Sopenharmony_ci
7662306a36Sopenharmony_ci	/* *key == 0 - the callback was called for array timer.
7762306a36Sopenharmony_ci	 * *key == 4 - the callback was called from lru timer.
7862306a36Sopenharmony_ci	 */
7962306a36Sopenharmony_ci	if (*key == ARRAY) {
8062306a36Sopenharmony_ci		struct bpf_timer *lru_timer;
8162306a36Sopenharmony_ci		int lru_key = LRU;
8262306a36Sopenharmony_ci
8362306a36Sopenharmony_ci		/* rearm array timer to be called again in ~35 seconds */
8462306a36Sopenharmony_ci		if (bpf_timer_start(timer, 1ull << 35, 0) != 0)
8562306a36Sopenharmony_ci			err |= 1;
8662306a36Sopenharmony_ci
8762306a36Sopenharmony_ci		lru_timer = bpf_map_lookup_elem(&lru, &lru_key);
8862306a36Sopenharmony_ci		if (!lru_timer)
8962306a36Sopenharmony_ci			return 0;
9062306a36Sopenharmony_ci		bpf_timer_set_callback(lru_timer, timer_cb1);
9162306a36Sopenharmony_ci		if (bpf_timer_start(lru_timer, 0, 0) != 0)
9262306a36Sopenharmony_ci			err |= 2;
9362306a36Sopenharmony_ci	} else if (*key == LRU) {
9462306a36Sopenharmony_ci		int lru_key, i;
9562306a36Sopenharmony_ci
9662306a36Sopenharmony_ci		for (i = LRU + 1;
9762306a36Sopenharmony_ci		     i <= 100  /* for current LRU eviction algorithm this number
9862306a36Sopenharmony_ci				* should be larger than ~ lru->max_entries * 2
9962306a36Sopenharmony_ci				*/;
10062306a36Sopenharmony_ci		     i++) {
10162306a36Sopenharmony_ci			struct elem init = {};
10262306a36Sopenharmony_ci
10362306a36Sopenharmony_ci			/* lru_key cannot be used as loop induction variable
10462306a36Sopenharmony_ci			 * otherwise the loop will be unbounded.
10562306a36Sopenharmony_ci			 */
10662306a36Sopenharmony_ci			lru_key = i;
10762306a36Sopenharmony_ci
10862306a36Sopenharmony_ci			/* add more elements into lru map to push out current
10962306a36Sopenharmony_ci			 * element and force deletion of this timer
11062306a36Sopenharmony_ci			 */
11162306a36Sopenharmony_ci			bpf_map_update_elem(map, &lru_key, &init, 0);
11262306a36Sopenharmony_ci			/* look it up to bump it into active list */
11362306a36Sopenharmony_ci			bpf_map_lookup_elem(map, &lru_key);
11462306a36Sopenharmony_ci
11562306a36Sopenharmony_ci			/* keep adding until *key changes underneath,
11662306a36Sopenharmony_ci			 * which means that key/timer memory was reused
11762306a36Sopenharmony_ci			 */
11862306a36Sopenharmony_ci			if (*key != LRU)
11962306a36Sopenharmony_ci				break;
12062306a36Sopenharmony_ci		}
12162306a36Sopenharmony_ci
12262306a36Sopenharmony_ci		/* check that the timer was removed */
12362306a36Sopenharmony_ci		if (bpf_timer_cancel(timer) != -EINVAL)
12462306a36Sopenharmony_ci			err |= 4;
12562306a36Sopenharmony_ci		ok |= 1;
12662306a36Sopenharmony_ci	}
12762306a36Sopenharmony_ci	return 0;
12862306a36Sopenharmony_ci}
12962306a36Sopenharmony_ci
13062306a36Sopenharmony_ciSEC("fentry/bpf_fentry_test1")
13162306a36Sopenharmony_ciint BPF_PROG2(test1, int, a)
13262306a36Sopenharmony_ci{
13362306a36Sopenharmony_ci	struct bpf_timer *arr_timer, *lru_timer;
13462306a36Sopenharmony_ci	struct elem init = {};
13562306a36Sopenharmony_ci	int lru_key = LRU;
13662306a36Sopenharmony_ci	int array_key = ARRAY;
13762306a36Sopenharmony_ci
13862306a36Sopenharmony_ci	arr_timer = bpf_map_lookup_elem(&array, &array_key);
13962306a36Sopenharmony_ci	if (!arr_timer)
14062306a36Sopenharmony_ci		return 0;
14162306a36Sopenharmony_ci	bpf_timer_init(arr_timer, &array, CLOCK_MONOTONIC);
14262306a36Sopenharmony_ci
14362306a36Sopenharmony_ci	bpf_map_update_elem(&lru, &lru_key, &init, 0);
14462306a36Sopenharmony_ci	lru_timer = bpf_map_lookup_elem(&lru, &lru_key);
14562306a36Sopenharmony_ci	if (!lru_timer)
14662306a36Sopenharmony_ci		return 0;
14762306a36Sopenharmony_ci	bpf_timer_init(lru_timer, &lru, CLOCK_MONOTONIC);
14862306a36Sopenharmony_ci
14962306a36Sopenharmony_ci	bpf_timer_set_callback(arr_timer, timer_cb1);
15062306a36Sopenharmony_ci	bpf_timer_start(arr_timer, 0 /* call timer_cb1 asap */, 0);
15162306a36Sopenharmony_ci
15262306a36Sopenharmony_ci	/* init more timers to check that array destruction
15362306a36Sopenharmony_ci	 * doesn't leak timer memory.
15462306a36Sopenharmony_ci	 */
15562306a36Sopenharmony_ci	array_key = 0;
15662306a36Sopenharmony_ci	arr_timer = bpf_map_lookup_elem(&array, &array_key);
15762306a36Sopenharmony_ci	if (!arr_timer)
15862306a36Sopenharmony_ci		return 0;
15962306a36Sopenharmony_ci	bpf_timer_init(arr_timer, &array, CLOCK_MONOTONIC);
16062306a36Sopenharmony_ci	return 0;
16162306a36Sopenharmony_ci}
16262306a36Sopenharmony_ci
16362306a36Sopenharmony_ci/* callback for prealloc and non-prealloca hashtab timers */
16462306a36Sopenharmony_cistatic int timer_cb2(void *map, int *key, struct hmap_elem *val)
16562306a36Sopenharmony_ci{
16662306a36Sopenharmony_ci	if (*key == HTAB)
16762306a36Sopenharmony_ci		callback_check--;
16862306a36Sopenharmony_ci	else
16962306a36Sopenharmony_ci		callback2_check--;
17062306a36Sopenharmony_ci	if (val->counter > 0 && --val->counter) {
17162306a36Sopenharmony_ci		/* re-arm the timer again to execute after 1 usec */
17262306a36Sopenharmony_ci		bpf_timer_start(&val->timer, 1000, 0);
17362306a36Sopenharmony_ci	} else if (*key == HTAB) {
17462306a36Sopenharmony_ci		struct bpf_timer *arr_timer;
17562306a36Sopenharmony_ci		int array_key = ARRAY;
17662306a36Sopenharmony_ci
17762306a36Sopenharmony_ci		/* cancel arr_timer otherwise bpf_fentry_test1 prog
17862306a36Sopenharmony_ci		 * will stay alive forever.
17962306a36Sopenharmony_ci		 */
18062306a36Sopenharmony_ci		arr_timer = bpf_map_lookup_elem(&array, &array_key);
18162306a36Sopenharmony_ci		if (!arr_timer)
18262306a36Sopenharmony_ci			return 0;
18362306a36Sopenharmony_ci		if (bpf_timer_cancel(arr_timer) != 1)
18462306a36Sopenharmony_ci			/* bpf_timer_cancel should return 1 to indicate
18562306a36Sopenharmony_ci			 * that arr_timer was active at this time
18662306a36Sopenharmony_ci			 */
18762306a36Sopenharmony_ci			err |= 8;
18862306a36Sopenharmony_ci
18962306a36Sopenharmony_ci		/* try to cancel ourself. It shouldn't deadlock. */
19062306a36Sopenharmony_ci		if (bpf_timer_cancel(&val->timer) != -EDEADLK)
19162306a36Sopenharmony_ci			err |= 16;
19262306a36Sopenharmony_ci
19362306a36Sopenharmony_ci		/* delete this key and this timer anyway.
19462306a36Sopenharmony_ci		 * It shouldn't deadlock either.
19562306a36Sopenharmony_ci		 */
19662306a36Sopenharmony_ci		bpf_map_delete_elem(map, key);
19762306a36Sopenharmony_ci
19862306a36Sopenharmony_ci		/* in preallocated hashmap both 'key' and 'val' could have been
19962306a36Sopenharmony_ci		 * reused to store another map element (like in LRU above),
20062306a36Sopenharmony_ci		 * but in controlled test environment the below test works.
20162306a36Sopenharmony_ci		 * It's not a use-after-free. The memory is owned by the map.
20262306a36Sopenharmony_ci		 */
20362306a36Sopenharmony_ci		if (bpf_timer_start(&val->timer, 1000, 0) != -EINVAL)
20462306a36Sopenharmony_ci			err |= 32;
20562306a36Sopenharmony_ci		ok |= 2;
20662306a36Sopenharmony_ci	} else {
20762306a36Sopenharmony_ci		if (*key != HTAB_MALLOC)
20862306a36Sopenharmony_ci			err |= 64;
20962306a36Sopenharmony_ci
21062306a36Sopenharmony_ci		/* try to cancel ourself. It shouldn't deadlock. */
21162306a36Sopenharmony_ci		if (bpf_timer_cancel(&val->timer) != -EDEADLK)
21262306a36Sopenharmony_ci			err |= 128;
21362306a36Sopenharmony_ci
21462306a36Sopenharmony_ci		/* delete this key and this timer anyway.
21562306a36Sopenharmony_ci		 * It shouldn't deadlock either.
21662306a36Sopenharmony_ci		 */
21762306a36Sopenharmony_ci		bpf_map_delete_elem(map, key);
21862306a36Sopenharmony_ci
21962306a36Sopenharmony_ci		ok |= 4;
22062306a36Sopenharmony_ci	}
22162306a36Sopenharmony_ci	return 0;
22262306a36Sopenharmony_ci}
22362306a36Sopenharmony_ci
22462306a36Sopenharmony_ciint bpf_timer_test(void)
22562306a36Sopenharmony_ci{
22662306a36Sopenharmony_ci	struct hmap_elem *val;
22762306a36Sopenharmony_ci	int key = HTAB, key_malloc = HTAB_MALLOC;
22862306a36Sopenharmony_ci
22962306a36Sopenharmony_ci	val = bpf_map_lookup_elem(&hmap, &key);
23062306a36Sopenharmony_ci	if (val) {
23162306a36Sopenharmony_ci		if (bpf_timer_init(&val->timer, &hmap, CLOCK_BOOTTIME) != 0)
23262306a36Sopenharmony_ci			err |= 512;
23362306a36Sopenharmony_ci		bpf_timer_set_callback(&val->timer, timer_cb2);
23462306a36Sopenharmony_ci		bpf_timer_start(&val->timer, 1000, 0);
23562306a36Sopenharmony_ci	}
23662306a36Sopenharmony_ci	val = bpf_map_lookup_elem(&hmap_malloc, &key_malloc);
23762306a36Sopenharmony_ci	if (val) {
23862306a36Sopenharmony_ci		if (bpf_timer_init(&val->timer, &hmap_malloc, CLOCK_BOOTTIME) != 0)
23962306a36Sopenharmony_ci			err |= 1024;
24062306a36Sopenharmony_ci		bpf_timer_set_callback(&val->timer, timer_cb2);
24162306a36Sopenharmony_ci		bpf_timer_start(&val->timer, 1000, 0);
24262306a36Sopenharmony_ci	}
24362306a36Sopenharmony_ci	return 0;
24462306a36Sopenharmony_ci}
24562306a36Sopenharmony_ci
24662306a36Sopenharmony_ciSEC("fentry/bpf_fentry_test2")
24762306a36Sopenharmony_ciint BPF_PROG2(test2, int, a, int, b)
24862306a36Sopenharmony_ci{
24962306a36Sopenharmony_ci	struct hmap_elem init = {}, *val;
25062306a36Sopenharmony_ci	int key = HTAB, key_malloc = HTAB_MALLOC;
25162306a36Sopenharmony_ci
25262306a36Sopenharmony_ci	init.counter = 10; /* number of times to trigger timer_cb2 */
25362306a36Sopenharmony_ci	bpf_map_update_elem(&hmap, &key, &init, 0);
25462306a36Sopenharmony_ci	val = bpf_map_lookup_elem(&hmap, &key);
25562306a36Sopenharmony_ci	if (val)
25662306a36Sopenharmony_ci		bpf_timer_init(&val->timer, &hmap, CLOCK_BOOTTIME);
25762306a36Sopenharmony_ci	/* update the same key to free the timer */
25862306a36Sopenharmony_ci	bpf_map_update_elem(&hmap, &key, &init, 0);
25962306a36Sopenharmony_ci
26062306a36Sopenharmony_ci	bpf_map_update_elem(&hmap_malloc, &key_malloc, &init, 0);
26162306a36Sopenharmony_ci	val = bpf_map_lookup_elem(&hmap_malloc, &key_malloc);
26262306a36Sopenharmony_ci	if (val)
26362306a36Sopenharmony_ci		bpf_timer_init(&val->timer, &hmap_malloc, CLOCK_BOOTTIME);
26462306a36Sopenharmony_ci	/* update the same key to free the timer */
26562306a36Sopenharmony_ci	bpf_map_update_elem(&hmap_malloc, &key_malloc, &init, 0);
26662306a36Sopenharmony_ci
26762306a36Sopenharmony_ci	/* init more timers to check that htab operations
26862306a36Sopenharmony_ci	 * don't leak timer memory.
26962306a36Sopenharmony_ci	 */
27062306a36Sopenharmony_ci	key = 0;
27162306a36Sopenharmony_ci	bpf_map_update_elem(&hmap, &key, &init, 0);
27262306a36Sopenharmony_ci	val = bpf_map_lookup_elem(&hmap, &key);
27362306a36Sopenharmony_ci	if (val)
27462306a36Sopenharmony_ci		bpf_timer_init(&val->timer, &hmap, CLOCK_BOOTTIME);
27562306a36Sopenharmony_ci	bpf_map_delete_elem(&hmap, &key);
27662306a36Sopenharmony_ci	bpf_map_update_elem(&hmap, &key, &init, 0);
27762306a36Sopenharmony_ci	val = bpf_map_lookup_elem(&hmap, &key);
27862306a36Sopenharmony_ci	if (val)
27962306a36Sopenharmony_ci		bpf_timer_init(&val->timer, &hmap, CLOCK_BOOTTIME);
28062306a36Sopenharmony_ci
28162306a36Sopenharmony_ci	/* and with non-prealloc htab */
28262306a36Sopenharmony_ci	key_malloc = 0;
28362306a36Sopenharmony_ci	bpf_map_update_elem(&hmap_malloc, &key_malloc, &init, 0);
28462306a36Sopenharmony_ci	val = bpf_map_lookup_elem(&hmap_malloc, &key_malloc);
28562306a36Sopenharmony_ci	if (val)
28662306a36Sopenharmony_ci		bpf_timer_init(&val->timer, &hmap_malloc, CLOCK_BOOTTIME);
28762306a36Sopenharmony_ci	bpf_map_delete_elem(&hmap_malloc, &key_malloc);
28862306a36Sopenharmony_ci	bpf_map_update_elem(&hmap_malloc, &key_malloc, &init, 0);
28962306a36Sopenharmony_ci	val = bpf_map_lookup_elem(&hmap_malloc, &key_malloc);
29062306a36Sopenharmony_ci	if (val)
29162306a36Sopenharmony_ci		bpf_timer_init(&val->timer, &hmap_malloc, CLOCK_BOOTTIME);
29262306a36Sopenharmony_ci
29362306a36Sopenharmony_ci	return bpf_timer_test();
29462306a36Sopenharmony_ci}
29562306a36Sopenharmony_ci
29662306a36Sopenharmony_ci/* callback for absolute timer */
29762306a36Sopenharmony_cistatic int timer_cb3(void *map, int *key, struct bpf_timer *timer)
29862306a36Sopenharmony_ci{
29962306a36Sopenharmony_ci	abs_data += 6;
30062306a36Sopenharmony_ci
30162306a36Sopenharmony_ci	if (abs_data < 12) {
30262306a36Sopenharmony_ci		bpf_timer_start(timer, bpf_ktime_get_boot_ns() + 1000,
30362306a36Sopenharmony_ci				BPF_F_TIMER_ABS);
30462306a36Sopenharmony_ci	} else {
30562306a36Sopenharmony_ci		/* Re-arm timer ~35 seconds in future */
30662306a36Sopenharmony_ci		bpf_timer_start(timer, bpf_ktime_get_boot_ns() + (1ull << 35),
30762306a36Sopenharmony_ci				BPF_F_TIMER_ABS);
30862306a36Sopenharmony_ci	}
30962306a36Sopenharmony_ci
31062306a36Sopenharmony_ci	return 0;
31162306a36Sopenharmony_ci}
31262306a36Sopenharmony_ci
31362306a36Sopenharmony_ciSEC("fentry/bpf_fentry_test3")
31462306a36Sopenharmony_ciint BPF_PROG2(test3, int, a)
31562306a36Sopenharmony_ci{
31662306a36Sopenharmony_ci	int key = 0;
31762306a36Sopenharmony_ci	struct bpf_timer *timer;
31862306a36Sopenharmony_ci
31962306a36Sopenharmony_ci	bpf_printk("test3");
32062306a36Sopenharmony_ci
32162306a36Sopenharmony_ci	timer = bpf_map_lookup_elem(&abs_timer, &key);
32262306a36Sopenharmony_ci	if (timer) {
32362306a36Sopenharmony_ci		if (bpf_timer_init(timer, &abs_timer, CLOCK_BOOTTIME) != 0)
32462306a36Sopenharmony_ci			err |= 2048;
32562306a36Sopenharmony_ci		bpf_timer_set_callback(timer, timer_cb3);
32662306a36Sopenharmony_ci		bpf_timer_start(timer, bpf_ktime_get_boot_ns() + 1000,
32762306a36Sopenharmony_ci				BPF_F_TIMER_ABS);
32862306a36Sopenharmony_ci	}
32962306a36Sopenharmony_ci
33062306a36Sopenharmony_ci	return 0;
33162306a36Sopenharmony_ci}
332