// SPDX-License-Identifier: GPL-2.0
#include <test_progs.h>
#include "stacktrace_ips.skel.h"

#ifdef __x86_64__
static int check_stacktrace_ips(int fd, __u32 key, int cnt, ...)
{
	__u64 ips[PERF_MAX_STACK_DEPTH];
	struct ksyms *ksyms = NULL;
	int i, err = 0;
	va_list args;

	/* sorted by addr */
	ksyms = load_kallsyms_local();
	if (!ASSERT_OK_PTR(ksyms, "load_kallsyms_local"))
		return -1;

	/* unlikely, but... */
	if (!ASSERT_LT(cnt, PERF_MAX_STACK_DEPTH, "check_max"))
		return -1;

	err = bpf_map_lookup_elem(fd, &key, ips);
	if (err)
		goto out;

	/*
	 * Compare all symbols provided via arguments with stacktrace ips,
	 * and their related symbol addresses.t
	 */
	va_start(args, cnt);

	for (i = 0; i < cnt; i++) {
		unsigned long val;
		struct ksym *ksym;

		val = va_arg(args, unsigned long);
		ksym = ksym_search_local(ksyms, ips[i]);
		if (!ASSERT_OK_PTR(ksym, "ksym_search_local"))
			break;
		ASSERT_EQ(ksym->addr, val, "stack_cmp");
	}

	va_end(args);

out:
	free_kallsyms_local(ksyms);
	return err;
}

static void test_stacktrace_ips_kprobe_multi(bool retprobe)
{
	LIBBPF_OPTS(bpf_kprobe_multi_opts, opts,
		.retprobe = retprobe
	);
	LIBBPF_OPTS(bpf_test_run_opts, topts);
	struct stacktrace_ips *skel;

	skel = stacktrace_ips__open_and_load();
	if (!ASSERT_OK_PTR(skel, "stacktrace_ips__open_and_load"))
		return;

	if (!skel->kconfig->CONFIG_UNWINDER_ORC) {
		test__skip();
		goto cleanup;
	}

	skel->links.kprobe_multi_test = bpf_program__attach_kprobe_multi_opts(
							skel->progs.kprobe_multi_test,
							"bpf_testmod_stacktrace_test", &opts);
	if (!ASSERT_OK_PTR(skel->links.kprobe_multi_test, "bpf_program__attach_kprobe_multi_opts"))
		goto cleanup;

	trigger_module_test_read(1);

	load_kallsyms();

	if (retprobe) {
		check_stacktrace_ips(bpf_map__fd(skel->maps.stackmap), skel->bss->stack_key, 4,
				     ksym_get_addr("bpf_testmod_stacktrace_test_3"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_2"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_1"),
				     ksym_get_addr("bpf_testmod_test_read"));
	} else {
		check_stacktrace_ips(bpf_map__fd(skel->maps.stackmap), skel->bss->stack_key, 5,
				     ksym_get_addr("bpf_testmod_stacktrace_test"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_3"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_2"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_1"),
				     ksym_get_addr("bpf_testmod_test_read"));
	}

cleanup:
	stacktrace_ips__destroy(skel);
}

static void test_stacktrace_ips_raw_tp(void)
{
	__u32 info_len = sizeof(struct bpf_prog_info);
	LIBBPF_OPTS(bpf_test_run_opts, topts);
	struct bpf_prog_info info = {};
	struct stacktrace_ips *skel;
	__u64 bpf_prog_ksym = 0;
	int err;

	skel = stacktrace_ips__open_and_load();
	if (!ASSERT_OK_PTR(skel, "stacktrace_ips__open_and_load"))
		return;

	if (!skel->kconfig->CONFIG_UNWINDER_ORC) {
		test__skip();
		goto cleanup;
	}

	skel->links.rawtp_test = bpf_program__attach_raw_tracepoint(
							skel->progs.rawtp_test,
							"bpf_testmod_test_read");
	if (!ASSERT_OK_PTR(skel->links.rawtp_test, "bpf_program__attach_raw_tracepoint"))
		goto cleanup;

	/* get bpf program address */
	info.jited_ksyms = ptr_to_u64(&bpf_prog_ksym);
	info.nr_jited_ksyms = 1;
	err = bpf_prog_get_info_by_fd(bpf_program__fd(skel->progs.rawtp_test),
				      &info, &info_len);
	if (!ASSERT_OK(err, "bpf_prog_get_info_by_fd"))
		goto cleanup;

	trigger_module_test_read(1);

	load_kallsyms();

	check_stacktrace_ips(bpf_map__fd(skel->maps.stackmap), skel->bss->stack_key, 2,
			     bpf_prog_ksym,
			     ksym_get_addr("bpf_trace_run2"));

cleanup:
	stacktrace_ips__destroy(skel);
}

static void test_stacktrace_ips_kprobe(bool retprobe)
{
	LIBBPF_OPTS(bpf_kprobe_opts, opts,
		.retprobe = retprobe
	);
	LIBBPF_OPTS(bpf_test_run_opts, topts);
	struct stacktrace_ips *skel;

	skel = stacktrace_ips__open_and_load();
	if (!ASSERT_OK_PTR(skel, "stacktrace_ips__open_and_load"))
		return;

	if (!skel->kconfig->CONFIG_UNWINDER_ORC) {
		test__skip();
		goto cleanup;
	}

	skel->links.kprobe_test = bpf_program__attach_kprobe_opts(
						skel->progs.kprobe_test,
						"bpf_testmod_stacktrace_test", &opts);
	if (!ASSERT_OK_PTR(skel->links.kprobe_test, "bpf_program__attach_kprobe_opts"))
		goto cleanup;

	trigger_module_test_read(1);

	load_kallsyms();

	if (retprobe) {
		check_stacktrace_ips(bpf_map__fd(skel->maps.stackmap), skel->bss->stack_key, 4,
				     ksym_get_addr("bpf_testmod_stacktrace_test_3"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_2"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_1"),
				     ksym_get_addr("bpf_testmod_test_read"));
	} else {
		check_stacktrace_ips(bpf_map__fd(skel->maps.stackmap), skel->bss->stack_key, 5,
				     ksym_get_addr("bpf_testmod_stacktrace_test"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_3"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_2"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_1"),
				     ksym_get_addr("bpf_testmod_test_read"));
	}

cleanup:
	stacktrace_ips__destroy(skel);
}

static void test_stacktrace_ips_trampoline(bool retprobe)
{
	LIBBPF_OPTS(bpf_test_run_opts, topts);
	struct stacktrace_ips *skel;

	skel = stacktrace_ips__open_and_load();
	if (!ASSERT_OK_PTR(skel, "stacktrace_ips__open_and_load"))
		return;

	if (!skel->kconfig->CONFIG_UNWINDER_ORC) {
		test__skip();
		goto cleanup;
	}

	if (retprobe) {
		skel->links.fexit_test = bpf_program__attach_trace(skel->progs.fexit_test);
		if (!ASSERT_OK_PTR(skel->links.fexit_test, "bpf_program__attach_trace"))
			goto cleanup;
	} else {
		skel->links.fentry_test = bpf_program__attach_trace(skel->progs.fentry_test);
		if (!ASSERT_OK_PTR(skel->links.fentry_test, "bpf_program__attach_trace"))
			goto cleanup;
	}

	trigger_module_test_read(1);

	load_kallsyms();

	if (retprobe) {
		check_stacktrace_ips(bpf_map__fd(skel->maps.stackmap), skel->bss->stack_key, 4,
				     ksym_get_addr("bpf_testmod_stacktrace_test_3"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_2"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_1"),
				     ksym_get_addr("bpf_testmod_test_read"));
	} else {
		check_stacktrace_ips(bpf_map__fd(skel->maps.stackmap), skel->bss->stack_key, 5,
				     ksym_get_addr("bpf_testmod_stacktrace_test"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_3"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_2"),
				     ksym_get_addr("bpf_testmod_stacktrace_test_1"),
				     ksym_get_addr("bpf_testmod_test_read"));
	}

cleanup:
	stacktrace_ips__destroy(skel);
}

static void __test_stacktrace_ips(void)
{
	if (test__start_subtest("kprobe_multi"))
		test_stacktrace_ips_kprobe_multi(false);
	if (test__start_subtest("kretprobe_multi"))
		test_stacktrace_ips_kprobe_multi(true);
	if (test__start_subtest("raw_tp"))
		test_stacktrace_ips_raw_tp();
	if (test__start_subtest("kprobe"))
		test_stacktrace_ips_kprobe(false);
	if (test__start_subtest("kretprobe"))
		test_stacktrace_ips_kprobe(true);
	if (test__start_subtest("fentry"))
		test_stacktrace_ips_trampoline(false);
	if (test__start_subtest("fexit"))
		test_stacktrace_ips_trampoline(true);
}
#else
static void __test_stacktrace_ips(void)
{
	test__skip();
}
#endif

void test_stacktrace_ips(void)
{
	__test_stacktrace_ips();
}