162306a36Sopenharmony_ci// SPDX-License-Identifier: GPL-2.0
262306a36Sopenharmony_ci#include <vmlinux.h>
362306a36Sopenharmony_ci#include <bpf/bpf_tracing.h>
462306a36Sopenharmony_ci#include <bpf/bpf_helpers.h>
562306a36Sopenharmony_ci#include "../bpf_testmod/bpf_testmod_kfunc.h"
662306a36Sopenharmony_ci
762306a36Sopenharmony_cistruct map_value {
862306a36Sopenharmony_ci	struct prog_test_ref_kfunc __kptr *ptr;
962306a36Sopenharmony_ci};
1062306a36Sopenharmony_ci
1162306a36Sopenharmony_cistruct {
1262306a36Sopenharmony_ci	__uint(type, BPF_MAP_TYPE_ARRAY);
1362306a36Sopenharmony_ci	__type(key, int);
1462306a36Sopenharmony_ci	__type(value, struct map_value);
1562306a36Sopenharmony_ci	__uint(max_entries, 16);
1662306a36Sopenharmony_ci} array_map SEC(".maps");
1762306a36Sopenharmony_ci
1862306a36Sopenharmony_cistatic __noinline int cb1(void *map, void *key, void *value, void *ctx)
1962306a36Sopenharmony_ci{
2062306a36Sopenharmony_ci	void *p = *(void **)ctx;
2162306a36Sopenharmony_ci	bpf_kfunc_call_test_release(p);
2262306a36Sopenharmony_ci	/* Without the fix this would cause underflow */
2362306a36Sopenharmony_ci	return 0;
2462306a36Sopenharmony_ci}
2562306a36Sopenharmony_ci
2662306a36Sopenharmony_ciSEC("?tc")
2762306a36Sopenharmony_ciint underflow_prog(void *ctx)
2862306a36Sopenharmony_ci{
2962306a36Sopenharmony_ci	struct prog_test_ref_kfunc *p;
3062306a36Sopenharmony_ci	unsigned long sl = 0;
3162306a36Sopenharmony_ci
3262306a36Sopenharmony_ci	p = bpf_kfunc_call_test_acquire(&sl);
3362306a36Sopenharmony_ci	if (!p)
3462306a36Sopenharmony_ci		return 0;
3562306a36Sopenharmony_ci	bpf_for_each_map_elem(&array_map, cb1, &p, 0);
3662306a36Sopenharmony_ci	bpf_kfunc_call_test_release(p);
3762306a36Sopenharmony_ci	return 0;
3862306a36Sopenharmony_ci}
3962306a36Sopenharmony_ci
4062306a36Sopenharmony_cistatic __always_inline int cb2(void *map, void *key, void *value, void *ctx)
4162306a36Sopenharmony_ci{
4262306a36Sopenharmony_ci	unsigned long sl = 0;
4362306a36Sopenharmony_ci
4462306a36Sopenharmony_ci	*(void **)ctx = bpf_kfunc_call_test_acquire(&sl);
4562306a36Sopenharmony_ci	/* Without the fix this would leak memory */
4662306a36Sopenharmony_ci	return 0;
4762306a36Sopenharmony_ci}
4862306a36Sopenharmony_ci
4962306a36Sopenharmony_ciSEC("?tc")
5062306a36Sopenharmony_ciint leak_prog(void *ctx)
5162306a36Sopenharmony_ci{
5262306a36Sopenharmony_ci	struct prog_test_ref_kfunc *p;
5362306a36Sopenharmony_ci	struct map_value *v;
5462306a36Sopenharmony_ci
5562306a36Sopenharmony_ci	v = bpf_map_lookup_elem(&array_map, &(int){0});
5662306a36Sopenharmony_ci	if (!v)
5762306a36Sopenharmony_ci		return 0;
5862306a36Sopenharmony_ci
5962306a36Sopenharmony_ci	p = NULL;
6062306a36Sopenharmony_ci	bpf_for_each_map_elem(&array_map, cb2, &p, 0);
6162306a36Sopenharmony_ci	p = bpf_kptr_xchg(&v->ptr, p);
6262306a36Sopenharmony_ci	if (p)
6362306a36Sopenharmony_ci		bpf_kfunc_call_test_release(p);
6462306a36Sopenharmony_ci	return 0;
6562306a36Sopenharmony_ci}
6662306a36Sopenharmony_ci
6762306a36Sopenharmony_cistatic __always_inline int cb(void *map, void *key, void *value, void *ctx)
6862306a36Sopenharmony_ci{
6962306a36Sopenharmony_ci	return 0;
7062306a36Sopenharmony_ci}
7162306a36Sopenharmony_ci
7262306a36Sopenharmony_cistatic __always_inline int cb3(void *map, void *key, void *value, void *ctx)
7362306a36Sopenharmony_ci{
7462306a36Sopenharmony_ci	unsigned long sl = 0;
7562306a36Sopenharmony_ci	void *p;
7662306a36Sopenharmony_ci
7762306a36Sopenharmony_ci	bpf_kfunc_call_test_acquire(&sl);
7862306a36Sopenharmony_ci	bpf_for_each_map_elem(&array_map, cb, &p, 0);
7962306a36Sopenharmony_ci	/* It should only complain here, not in cb. This is why we need
8062306a36Sopenharmony_ci	 * callback_ref to be set to frameno.
8162306a36Sopenharmony_ci	 */
8262306a36Sopenharmony_ci	return 0;
8362306a36Sopenharmony_ci}
8462306a36Sopenharmony_ci
8562306a36Sopenharmony_ciSEC("?tc")
8662306a36Sopenharmony_ciint nested_cb(void *ctx)
8762306a36Sopenharmony_ci{
8862306a36Sopenharmony_ci	struct prog_test_ref_kfunc *p;
8962306a36Sopenharmony_ci	unsigned long sl = 0;
9062306a36Sopenharmony_ci	int sp = 0;
9162306a36Sopenharmony_ci
9262306a36Sopenharmony_ci	p = bpf_kfunc_call_test_acquire(&sl);
9362306a36Sopenharmony_ci	if (!p)
9462306a36Sopenharmony_ci		return 0;
9562306a36Sopenharmony_ci	bpf_for_each_map_elem(&array_map, cb3, &sp, 0);
9662306a36Sopenharmony_ci	bpf_kfunc_call_test_release(p);
9762306a36Sopenharmony_ci	return 0;
9862306a36Sopenharmony_ci}
9962306a36Sopenharmony_ci
10062306a36Sopenharmony_ciSEC("?tc")
10162306a36Sopenharmony_ciint non_cb_transfer_ref(void *ctx)
10262306a36Sopenharmony_ci{
10362306a36Sopenharmony_ci	struct prog_test_ref_kfunc *p;
10462306a36Sopenharmony_ci	unsigned long sl = 0;
10562306a36Sopenharmony_ci
10662306a36Sopenharmony_ci	p = bpf_kfunc_call_test_acquire(&sl);
10762306a36Sopenharmony_ci	if (!p)
10862306a36Sopenharmony_ci		return 0;
10962306a36Sopenharmony_ci	cb1(NULL, NULL, NULL, &p);
11062306a36Sopenharmony_ci	bpf_kfunc_call_test_acquire(&sl);
11162306a36Sopenharmony_ci	return 0;
11262306a36Sopenharmony_ci}
11362306a36Sopenharmony_ci
11462306a36Sopenharmony_cichar _license[] SEC("license") = "GPL";
115