// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */
#include <vmlinux.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>
#include "../test_kmods/bpf_testmod_kfunc.h"

struct map_value {
	struct prog_test_ref_kfunc __kptr *ref_ptr;
};

struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__uint(map_flags, BPF_F_NO_PREALLOC);
	__type(key, int);
	__type(value, struct map_value);
	__uint(max_entries, 1);
} race_hash_map SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_PERCPU_HASH);
	__uint(map_flags, BPF_F_NO_PREALLOC);
	__type(key, int);
	__type(value, struct map_value);
	__uint(max_entries, 1);
} race_percpu_hash_map SEC(".maps");

struct {
	__uint(type, BPF_MAP_TYPE_SK_STORAGE);
	__uint(map_flags, BPF_F_NO_PREALLOC);
	__type(key, int);
	__type(value, struct map_value);
} race_sk_ls_map SEC(".maps");

int num_of_refs;
int sk_ls_leak_done;
int target_map_id;
int map_freed;
const volatile int nr_cpus;

SEC("tc")
int test_htab_leak(struct __sk_buff *skb)
{
	struct prog_test_ref_kfunc *p, *old;
	struct map_value val = {};
	struct map_value *v;
	int key = 0;

	if (bpf_map_update_elem(&race_hash_map, &key, &val, BPF_ANY))
		return 1;

	v = bpf_map_lookup_elem(&race_hash_map, &key);
	if (!v)
		return 2;

	p = bpf_kfunc_call_test_acquire(&(unsigned long){0});
	if (!p)
		return 3;
	old = bpf_kptr_xchg(&v->ref_ptr, p);
	if (old)
		bpf_kfunc_call_test_release(old);

	bpf_map_delete_elem(&race_hash_map, &key);

	p = bpf_kfunc_call_test_acquire(&(unsigned long){0});
	if (!p)
		return 4;
	old = bpf_kptr_xchg(&v->ref_ptr, p);
	if (old)
		bpf_kfunc_call_test_release(old);

	return 0;
}

static int fill_percpu_kptr(struct map_value *v)
{
	struct prog_test_ref_kfunc *p, *old;

	p = bpf_kfunc_call_test_acquire(&(unsigned long){0});
	if (!p)
		return 1;
	old = bpf_kptr_xchg(&v->ref_ptr, p);
	if (old)
		bpf_kfunc_call_test_release(old);
	return 0;
}

SEC("tc")
int test_percpu_htab_leak(struct __sk_buff *skb)
{
	struct map_value *v, *arr[16] = {};
	struct map_value val = {};
	int key = 0;
	int err = 0;

	if (bpf_map_update_elem(&race_percpu_hash_map, &key, &val, BPF_ANY))
		return 1;

	for (int i = 0; i < nr_cpus; i++) {
		v = bpf_map_lookup_percpu_elem(&race_percpu_hash_map, &key, i);
		if (!v)
			return 2;
		arr[i] = v;
	}

	bpf_map_delete_elem(&race_percpu_hash_map, &key);

	for (int i = 0; i < nr_cpus; i++) {
		v = arr[i];
		err = fill_percpu_kptr(v);
		if (err)
			return 3;
	}

	return 0;
}

SEC("tp_btf/inet_sock_set_state")
int BPF_PROG(test_sk_ls_leak, struct sock *sk, int oldstate, int newstate)
{
	struct prog_test_ref_kfunc *p, *old;
	struct map_value *v;

	if (newstate != BPF_TCP_SYN_SENT)
		return 0;

	if (sk_ls_leak_done)
		return 0;

	v = bpf_sk_storage_get(&race_sk_ls_map, sk, NULL,
				BPF_SK_STORAGE_GET_F_CREATE);
	if (!v)
		return 0;

	p = bpf_kfunc_call_test_acquire(&(unsigned long){0});
	if (!p)
		return 0;
	old = bpf_kptr_xchg(&v->ref_ptr, p);
	if (old)
		bpf_kfunc_call_test_release(old);

	bpf_sk_storage_delete(&race_sk_ls_map, sk);

	p = bpf_kfunc_call_test_acquire(&(unsigned long){0});
	if (!p)
		return 0;
	old = bpf_kptr_xchg(&v->ref_ptr, p);
	if (old)
		bpf_kfunc_call_test_release(old);

	sk_ls_leak_done = 1;
	return 0;
}

long target_map_ptr;

SEC("fentry/bpf_map_put")
int BPF_PROG(map_put, struct bpf_map *map)
{
	if (target_map_id && map->id == (u32)target_map_id)
		target_map_ptr = (long)map;
	return 0;
}

SEC("fexit/htab_map_free")
int BPF_PROG(htab_map_free, struct bpf_map *map)
{
	if (target_map_ptr && (long)map == target_map_ptr)
		map_freed = 1;
	return 0;
}

SEC("fexit/bpf_sk_storage_map_free")
int BPF_PROG(sk_map_free, struct bpf_map *map)
{
	if (target_map_ptr && (long)map == target_map_ptr)
		map_freed = 1;
	return 0;
}

SEC("syscall")
int count_ref(void *ctx)
{
	struct prog_test_ref_kfunc *p;
	unsigned long arg = 0;

	p = bpf_kfunc_call_test_acquire(&arg);
	if (!p)
		return 1;

	num_of_refs = p->cnt.refs.counter;

	bpf_kfunc_call_test_release(p);
	return 0;
}

char _license[] SEC("license") = "GPL";