// SPDX-License-Identifier: GPL-2.0

#include <test_progs.h>
#include <bpf/btf.h>
#include <search.h>
#include "bpf/libbpf_internal.h"
#include "tracing_multi.skel.h"
#include "tracing_multi_module.skel.h"
#include "tracing_multi_intersect.skel.h"
#include "tracing_multi_session.skel.h"
#include "tracing_multi_fail.skel.h"
#include "tracing_multi_verifier.skel.h"
#include "tracing_multi_bench.skel.h"
#include "tracing_multi_rollback.skel.h"
#include "trace_helpers.h"

static __u64 bpf_fentry_test_cookies[] = {
	8,  /* bpf_fentry_test1 */
	9,  /* bpf_fentry_test2 */
	7,  /* bpf_fentry_test3 */
	5,  /* bpf_fentry_test4 */
	4,  /* bpf_fentry_test5 */
	2,  /* bpf_fentry_test6 */
	3,  /* bpf_fentry_test7 */
	1,  /* bpf_fentry_test8 */
	10, /* bpf_fentry_test9 */
	6,  /* bpf_fentry_test10 */
};

static const char * const bpf_fentry_test[] = {
	"bpf_fentry_test1",
	"bpf_fentry_test2",
	"bpf_fentry_test3",
	"bpf_fentry_test4",
	"bpf_fentry_test5",
	"bpf_fentry_test6",
	"bpf_fentry_test7",
	"bpf_fentry_test8",
	"bpf_fentry_test9",
	"bpf_fentry_test10",
};

static const char * const bpf_testmod_fentry_test[] = {
	"bpf_testmod_fentry_test1",
	"bpf_testmod_fentry_test2",
	"bpf_testmod_fentry_test3",
	"bpf_testmod_fentry_test7",
	"bpf_testmod_fentry_test11",
};

#define FUNCS_CNT (ARRAY_SIZE(bpf_fentry_test))

static int get_random_funcs(const char **funcs)
{
	int i, cnt = 0;

	for (i = 0; i < FUNCS_CNT; i++) {
		if (rand() % 2)
			funcs[cnt++] = bpf_fentry_test[i];
	}
	/* we always need at least one.. */
	if (!cnt)
		funcs[cnt++] = bpf_fentry_test[rand() % FUNCS_CNT];
	return cnt;
}

static int compare(const void *ppa, const void *ppb)
{
	const char *pa = *(const char **) ppa;
	const char *pb = *(const char **) ppb;

	return strcmp(pa, pb);
}

static void tdestroy_free_nop(void *ptr)
{
}

static __u32 *get_ids(const char * const funcs[], int funcs_cnt, const char *mod)
{
	struct btf *btf, *vmlinux_btf = NULL;
	__u32 nr, type_id, cnt = 0;
	void *root = NULL;
	__u32 *ids = NULL;
	int i, err = 0;

	btf = btf__load_vmlinux_btf();
	if (!ASSERT_OK_PTR(btf, "btf__load_vmlinux_btf"))
		return NULL;

	if (mod) {
		vmlinux_btf = btf;
		btf = btf__load_module_btf(mod, vmlinux_btf);
		if (!ASSERT_OK_PTR(btf, "btf__load_module_btf")) {
			btf__free(vmlinux_btf);
			return NULL;
		}
	}

	ids = calloc(funcs_cnt, sizeof(ids[0]));
	if (!ids)
		goto out;

	/*
	 * We sort function names by name and search them
	 * below for each function.
	 */
	for (i = 0; i < funcs_cnt; i++) {
		if (!tsearch(&funcs[i], &root, compare)) {
			ASSERT_FAIL("tsearch failed");
			err = -1;
			goto error;
		}
	}

	nr = btf__type_cnt(btf);
	for (type_id = 1; type_id < nr && cnt < funcs_cnt; type_id++) {
		const struct btf_type *type;
		const char *str, ***val;
		unsigned int idx;

		type = btf__type_by_id(btf, type_id);
		if (!type) {
			err = -1;
			break;
		}

		if (BTF_INFO_KIND(type->info) != BTF_KIND_FUNC)
			continue;

		str = btf__name_by_offset(btf, type->name_off);
		if (!str) {
			err = -1;
			break;
		}

		val = tfind(&str, &root, compare);
		if (!val)
			continue;

		/*
		 * We keep pointer for each function name so we can get the original
		 * array index and have the resulting ids array matching the original
		 * function array.
		 *
		 * Doing it this way allow us to easily test the cookies support,
		 * because each cookie is attached to particular function/id.
		 */
		idx = *val - funcs;
		ids[idx] = type_id;
		cnt++;
	}

error:
	if (err) {
		free(ids);
		ids = NULL;
	}

out:
	tdestroy(root, tdestroy_free_nop);
	btf__free(vmlinux_btf);
	btf__free(btf);
	return ids;
}

static void tracing_multi_test_run(struct tracing_multi *skel)
{
	LIBBPF_OPTS(bpf_test_run_opts, topts);
	int err, prog_fd;

	prog_fd = bpf_program__fd(skel->progs.test_fentry);
	err = bpf_prog_test_run_opts(prog_fd, &topts);
	ASSERT_OK(err, "test_run");

	/* extra +1 count for sleepable programs */
	ASSERT_EQ(skel->bss->test_result_fentry, FUNCS_CNT + 1, "test_result_fentry");
	ASSERT_EQ(skel->bss->test_result_fexit, FUNCS_CNT + 1, "test_result_fexit");
}

static void test_skel_api(void)
{
	struct tracing_multi *skel;
	int err;

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

	skel->bss->pid = getpid();

	err = tracing_multi__attach(skel);
	if (!ASSERT_OK(err, "tracing_multi__attach"))
		goto cleanup;

	tracing_multi_test_run(skel);

cleanup:
	tracing_multi__destroy(skel);
}

static void test_link_api_pattern(void)
{
	struct tracing_multi *skel;

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

	skel->bss->pid = getpid();

	skel->links.test_fentry = bpf_program__attach_tracing_multi(skel->progs.test_fentry,
					"bpf_fentry_test*", NULL);
	if (!ASSERT_OK_PTR(skel->links.test_fentry, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	skel->links.test_fexit = bpf_program__attach_tracing_multi(skel->progs.test_fexit,
					"bpf_fentry_test*", NULL);
	if (!ASSERT_OK_PTR(skel->links.test_fexit, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	skel->links.test_fentry_s = bpf_program__attach_tracing_multi(skel->progs.test_fentry_s,
					"bpf_fentry_test1", NULL);
	if (!ASSERT_OK_PTR(skel->links.test_fentry_s, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	skel->links.test_fexit_s = bpf_program__attach_tracing_multi(skel->progs.test_fexit_s,
					"bpf_fentry_test1", NULL);
	if (!ASSERT_OK_PTR(skel->links.test_fexit_s, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	tracing_multi_test_run(skel);

cleanup:
	tracing_multi__destroy(skel);
}

static void test_link_api_ids(bool test_cookies)
{
	LIBBPF_OPTS(bpf_tracing_multi_opts, opts);
	struct tracing_multi *skel;
	size_t cnt = FUNCS_CNT;
	__u32 *ids;

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

	skel->bss->pid = getpid();
	skel->bss->test_cookies = test_cookies;

	ids = get_ids(bpf_fentry_test, cnt, NULL);
	if (!ASSERT_OK_PTR(ids, "get_ids"))
		goto cleanup;

	opts.ids = ids;
	opts.cnt = cnt;

	if (test_cookies)
		opts.cookies = bpf_fentry_test_cookies;

	skel->links.test_fentry = bpf_program__attach_tracing_multi(skel->progs.test_fentry,
						NULL, &opts);
	if (!ASSERT_OK_PTR(skel->links.test_fentry, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	skel->links.test_fexit = bpf_program__attach_tracing_multi(skel->progs.test_fexit,
						NULL, &opts);
	if (!ASSERT_OK_PTR(skel->links.test_fexit, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	/* Only bpf_fentry_test1 is allowed for sleepable programs. */
	opts.cnt = 1;
	skel->links.test_fentry_s = bpf_program__attach_tracing_multi(skel->progs.test_fentry_s,
						NULL, &opts);
	if (!ASSERT_OK_PTR(skel->links.test_fentry_s, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	skel->links.test_fexit_s = bpf_program__attach_tracing_multi(skel->progs.test_fexit_s,
						NULL, &opts);
	if (!ASSERT_OK_PTR(skel->links.test_fexit_s, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	tracing_multi_test_run(skel);

cleanup:
	tracing_multi__destroy(skel);
	free(ids);
}

static void test_module_skel_api(void)
{
	struct tracing_multi_module *skel = NULL;
	int err;

	skel = tracing_multi_module__open_and_load();
	if (!ASSERT_OK_PTR(skel, "tracing_multi__open_and_load"))
		return;

	skel->bss->pid = getpid();

	err = tracing_multi_module__attach(skel);
	if (!ASSERT_OK(err, "tracing_multi__attach"))
		goto cleanup;

	ASSERT_OK(trigger_module_test_read(1), "trigger_read");
	ASSERT_EQ(skel->bss->test_result_fentry, 5, "test_result_fentry");
	ASSERT_EQ(skel->bss->test_result_fexit, 5, "test_result_fexit");

cleanup:
	tracing_multi_module__destroy(skel);
}

static void test_module_link_api_pattern(void)
{
	struct tracing_multi_module *skel = NULL;

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

	skel->bss->pid = getpid();

	skel->links.test_fentry = bpf_program__attach_tracing_multi(skel->progs.test_fentry,
					"bpf_testmod:bpf_testmod_fentry_test*", NULL);
	if (!ASSERT_OK_PTR(skel->links.test_fentry, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	skel->links.test_fexit = bpf_program__attach_tracing_multi(skel->progs.test_fexit,
					"bpf_testmod:bpf_testmod_fentry_test*", NULL);
	if (!ASSERT_OK_PTR(skel->links.test_fexit, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	ASSERT_OK(trigger_module_test_read(1), "trigger_read");
	ASSERT_EQ(skel->bss->test_result_fentry, 5, "test_result_fentry");
	ASSERT_EQ(skel->bss->test_result_fexit, 5, "test_result_fexit");

cleanup:
	tracing_multi_module__destroy(skel);
}

static void test_module_link_api_ids(void)
{
	size_t cnt = ARRAY_SIZE(bpf_testmod_fentry_test);
	LIBBPF_OPTS(bpf_tracing_multi_opts, opts);
	struct tracing_multi_module *skel = NULL;
	__u32 *ids;

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

	skel->bss->pid = getpid();

	ids = get_ids(bpf_testmod_fentry_test, cnt, "bpf_testmod");
	if (!ASSERT_OK_PTR(ids, "get_ids"))
		goto cleanup;

	opts.ids = ids;
	opts.cnt = cnt;

	skel->links.test_fentry = bpf_program__attach_tracing_multi(skel->progs.test_fentry,
						NULL, &opts);
	if (!ASSERT_OK_PTR(skel->links.test_fentry, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	skel->links.test_fexit = bpf_program__attach_tracing_multi(skel->progs.test_fexit,
						NULL, &opts);
	if (!ASSERT_OK_PTR(skel->links.test_fexit, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	ASSERT_OK(trigger_module_test_read(1), "trigger_read");
	ASSERT_EQ(skel->bss->test_result_fentry, 5, "test_result_fentry");
	ASSERT_EQ(skel->bss->test_result_fexit, 5, "test_result_fexit");

cleanup:
	tracing_multi_module__destroy(skel);
	free(ids);
}

static bool is_set(__u32 mask, __u32 bit)
{
	return (1 << bit) & mask;
}

static void __test_intersect(__u32 mask, const struct bpf_program *progs[4], __u64 *test_results[4])
{
	LIBBPF_OPTS(bpf_tracing_multi_opts, opts);
	LIBBPF_OPTS(bpf_test_run_opts, topts);
	struct bpf_link *links[4] = { NULL };
	const char *funcs[FUNCS_CNT];
	__u64 expected[4];
	__u32 *ids, i;
	int err, cnt;

	/*
	 * We have 4 programs in progs and the mask bits pick which
	 * of them gets attached to randomly chosen functions.
	 */
	for (i = 0; i < 4; i++) {
		if (!is_set(mask, i))
			continue;

		cnt = get_random_funcs(funcs);
		ids = get_ids(funcs, cnt, NULL);
		if (!ASSERT_OK_PTR(ids, "get_ids"))
			goto cleanup;

		opts.ids = ids;
		opts.cnt = cnt;
		links[i] = bpf_program__attach_tracing_multi(progs[i], NULL, &opts);
		free(ids);

		if (!ASSERT_OK_PTR(links[i], "bpf_program__attach_tracing_multi"))
			goto cleanup;

		expected[i] = *test_results[i] + cnt;
	}

	err = bpf_prog_test_run_opts(bpf_program__fd(progs[0]), &topts);
	ASSERT_OK(err, "test_run");

	for (i = 0; i < 4; i++) {
		if (!is_set(mask, i))
			continue;
		ASSERT_EQ(*test_results[i], expected[i], "test_results");
	}

cleanup:
	for (i = 0; i < 4; i++)
		bpf_link__destroy(links[i]);
}

static void test_intersect(void)
{
	struct tracing_multi_intersect *skel;
	const struct bpf_program *progs[4];
	__u64 *test_results[4];
	__u32 i;

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

	skel->bss->pid = getpid();

	progs[0] = skel->progs.fentry_1;
	progs[1] = skel->progs.fexit_1;
	progs[2] = skel->progs.fentry_2;
	progs[3] = skel->progs.fexit_2;

	test_results[0] = &skel->bss->test_result_fentry_1;
	test_results[1] = &skel->bss->test_result_fexit_1;
	test_results[2] = &skel->bss->test_result_fentry_2;
	test_results[3] = &skel->bss->test_result_fexit_2;

	for (i = 1; i < 16; i++)
		__test_intersect(i, progs, test_results);

	tracing_multi_intersect__destroy(skel);
}

static void test_session(void)
{
	LIBBPF_OPTS(bpf_test_run_opts, topts);
	struct tracing_multi_session *skel;
	int err, prog_fd;

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

	skel->bss->pid = getpid();

	err = tracing_multi_session__attach(skel);
	if (!ASSERT_OK(err, "tracing_multi_session__attach"))
		goto cleanup;

	/* execute kernel session */
	prog_fd = bpf_program__fd(skel->progs.test_session_1);
	err = bpf_prog_test_run_opts(prog_fd, &topts);
	ASSERT_OK(err, "test_run");

	/* 10 for test_session_1, 1 for test_fsession_s */
	ASSERT_EQ(skel->bss->test_result_fentry, 11, "test_result_fentry");
	/* extra count (+1 for each fexit execution) for test_result_fexit cookie check/inc */
	ASSERT_EQ(skel->bss->test_result_fexit, 22, "test_result_fexit");

	skel->bss->test_result_fentry = 0;
	skel->bss->test_result_fexit = 0;

	/* execute bpf_testmo.ko session */
	ASSERT_OK(trigger_module_test_read(1), "trigger_read");

	/* 5 for test_session_2 */
	ASSERT_EQ(skel->bss->test_result_fentry, 5, "test_result_fentry");
	/* extra count (+1 for each fexit execution) for test_result_fexit cookie */
	ASSERT_EQ(skel->bss->test_result_fexit, 10, "test_result_fexit");


cleanup:
	tracing_multi_session__destroy(skel);
}

static void test_attach_api_fails(void)
{
	LIBBPF_OPTS(bpf_tracing_multi_opts, opts);
	static const char * const func[] = {
		"bpf_fentry_test2",
	};
	struct tracing_multi_fail *skel = NULL;
	__u32 ids[2] = {}, *ids2 = NULL;
	__u64 cookies[2];

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

	/* fail#1 (libbpf) pattern and opts NULL */
	skel->links.test_fentry = bpf_program__attach_tracing_multi(skel->progs.test_fentry,
						NULL, NULL);
	if (!ASSERT_EQ(libbpf_get_error(skel->links.test_fentry), -EINVAL, "fail_1"))
		goto cleanup;

	/* fail#2 (libbpf) pattern and ids */
	LIBBPF_OPTS_RESET(opts,
		.ids = ids,
		.cnt = 2,
	);

	skel->links.test_fentry = bpf_program__attach_tracing_multi(skel->progs.test_fentry,
						"bpf_fentry_test*", &opts);
	if (!ASSERT_EQ(libbpf_get_error(skel->links.test_fentry), -EINVAL, "fail_2"))
		goto cleanup;

	/* fail#3 (libbpf) pattern and cookies */
	LIBBPF_OPTS_RESET(opts,
		.ids = NULL,
		.cnt = 2,
		.cookies = cookies,
	);

	skel->links.test_fentry = bpf_program__attach_tracing_multi(skel->progs.test_fentry,
						"bpf_fentry_test*", &opts);
	if (!ASSERT_EQ(libbpf_get_error(skel->links.test_fentry), -EINVAL, "fail_3"))
		goto cleanup;

	/* fail#4 (libbpf) bogus pattern */
	skel->links.test_fentry = bpf_program__attach_tracing_multi(skel->progs.test_fentry,
						"bpf_not_really_a_function*", NULL);
	if (!ASSERT_EQ(libbpf_get_error(skel->links.test_fentry), -EINVAL, "fail_4"))
		goto cleanup;

	/* fail#5 (kernel) abnormal cnt */
	LIBBPF_OPTS_RESET(opts,
		.ids = ids,
		.cnt = INT_MAX,
	);

	skel->links.test_fentry = bpf_program__attach_tracing_multi(skel->progs.test_fentry,
						NULL, &opts);
	if (!ASSERT_EQ(libbpf_get_error(skel->links.test_fentry), -E2BIG, "fail_5"))
		goto cleanup;

	/* fail#6 (kernel) attach sleepable program to not-allowed function */
	ids2 = get_ids(func, 1, NULL);
	if (!ASSERT_OK_PTR(ids2, "get_ids"))
		goto cleanup;

	LIBBPF_OPTS_RESET(opts,
		.ids = ids2,
		.cnt = 1,
	);

	skel->links.test_fentry_s = bpf_program__attach_tracing_multi(skel->progs.test_fentry_s,
						NULL, &opts);
	if (!ASSERT_EQ(libbpf_get_error(skel->links.test_fentry_s), -EINVAL, "fail_6"))
		goto cleanup;

	/* fail#7 (kernel) attach with duplicate id */
	ids[0] = ids2[0];
	ids[1] = ids2[0];

	LIBBPF_OPTS_RESET(opts,
		.ids = ids,
		.cnt = 2,
	);

	skel->links.test_fentry = bpf_program__attach_tracing_multi(skel->progs.test_fentry,
						NULL, &opts);
	ASSERT_EQ(libbpf_get_error(skel->links.test_fentry), -EINVAL, "fail_7");

cleanup:
	tracing_multi_fail__destroy(skel);
	free(ids2);
}

void serial_test_tracing_multi_bench_attach(void)
{
	LIBBPF_OPTS(bpf_tracing_multi_opts, opts);
	struct tracing_multi_bench *skel = NULL;
	long attach_start_ns, attach_end_ns;
	long detach_start_ns, detach_end_ns;
	double attach_delta, detach_delta;
	struct bpf_link *link = NULL;
	size_t i, cap = 0, cnt = 0;
	struct ksyms *ksyms = NULL;
	void *root = NULL;
	void *dups = NULL;
	__u32 *ids = NULL;
	__u32 nr, type_id;
	struct btf *btf;
	int err;

#ifndef __x86_64__
	test__skip();
	return;
#endif

	btf = btf__load_vmlinux_btf();
	if (!ASSERT_OK_PTR(btf, "btf__load_vmlinux_btf"))
		return;

	skel = tracing_multi_bench__open_and_load();
	if (!ASSERT_OK_PTR(skel, "tracing_multi_bench__open_and_load"))
		goto cleanup;

	if (!ASSERT_OK(bpf_get_ksyms(&ksyms, true), "get_syms"))
		goto cleanup;

	/* Get all ftrace 'safe' symbols.. */
	for (i = 0; i < ksyms->filtered_cnt; i++) {
		if (!tsearch(&ksyms->filtered_syms[i], &root, compare)) {
			ASSERT_FAIL("tsearch failed");
			goto cleanup;
		}
	}

	/*
	 * Collect names that are not unique in kallsyms. The kernel resolves a
	 * tracing-multi BTF id to an address with kallsyms_lookup_name(), which
	 * returns the first symbol of that name. For a duplicate name that may
	 * be a different (non-ftrace-able) instance than the ftrace-able one in
	 * available_filter_functions, so attaching to it by BTF id fails with
	 * -ENOENT (e.g. t_start/t_next/t_stop). ksyms->syms is sorted by name,
	 * so equal names are adjacent.
	 */
	for (i = 1; i < ksyms->sym_cnt; i++) {
		if (strcmp(ksyms->syms[i].name, ksyms->syms[i - 1].name))
			continue;
		if (!tsearch(&ksyms->syms[i].name, &dups, compare)) {
			ASSERT_FAIL("tsearch failed");
			goto cleanup;
		}
	}

	/* ..and filter them through BTF and btf_type_is_traceable_func. */
	nr = btf__type_cnt(btf);
	for (type_id = 1; type_id < nr; type_id++) {
		const struct btf_type *type;
		const char *str;

		type = btf__type_by_id(btf, type_id);
		if (!type)
			break;

		if (BTF_INFO_KIND(type->info) != BTF_KIND_FUNC)
			continue;

		str = btf__name_by_offset(btf, type->name_off);
		if (!str)
			break;

		if (!tfind(&str, &root, compare))
			continue;

		/* Skip names that are not unique in kallsyms, see above. */
		if (tfind(&str, &dups, compare))
			continue;

		if (!btf_type_is_traceable_func(btf, type))
			continue;

		err = libbpf_ensure_mem((void **) &ids, &cap, sizeof(*ids), cnt + 1);
		if (err)
			goto cleanup;

		ids[cnt++] = type_id;
	}

	opts.ids = ids;
	opts.cnt = cnt;

	attach_start_ns = get_time_ns();
	link = bpf_program__attach_tracing_multi(skel->progs.bench, NULL, &opts);
	attach_end_ns = get_time_ns();

	if (!ASSERT_OK_PTR(link, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	detach_start_ns = get_time_ns();
	bpf_link__destroy(link);
	detach_end_ns = get_time_ns();

	attach_delta = (attach_end_ns - attach_start_ns) / 1000000000.0;
	detach_delta = (detach_end_ns - detach_start_ns) / 1000000000.0;

	printf("%s: found %lu functions\n", __func__, cnt);
	printf("%s: attached in %7.3lfs\n", __func__, attach_delta);
	printf("%s: detached in %7.3lfs\n", __func__, detach_delta);

cleanup:
	tracing_multi_bench__destroy(skel);
	tdestroy(root, tdestroy_free_nop);
	tdestroy(dups, tdestroy_free_nop);
	free_kallsyms_local(ksyms);
	free(ids);
	btf__free(btf);
}

static void tracing_multi_rollback_run(struct tracing_multi_rollback *skel)
{
	LIBBPF_OPTS(bpf_test_run_opts, topts);
	int err, prog_fd;

	prog_fd = bpf_program__fd(skel->progs.test_fentry);
	err = bpf_prog_test_run_opts(prog_fd, &topts);
	ASSERT_OK(err, "test_run");

	/* make sure the rollback code did not leave any program attached */
	ASSERT_EQ(skel->bss->test_result_fentry, 0, "test_result_fentry");
	ASSERT_EQ(skel->bss->test_result_fexit, 0, "test_result_fexit");
}

static void test_rollback_put(void)
{
	LIBBPF_OPTS(bpf_tracing_multi_opts, opts);
	struct tracing_multi_rollback *skel = NULL;
	size_t cnt = FUNCS_CNT;
	__u32 *ids = NULL;
	int err;

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

	bpf_program__set_autoload(skel->progs.test_fentry, true);
	bpf_program__set_autoload(skel->progs.test_fexit, true);

	err = tracing_multi_rollback__load(skel);
	if (!ASSERT_OK(err, "tracing_multi_rollback__load"))
		goto cleanup;

	ids = get_ids(bpf_fentry_test, cnt, NULL);
	if (!ASSERT_OK_PTR(ids, "get_ids"))
		goto cleanup;

	/*
	 * Mangle last id to trigger rollback, which needs to do put
	 * on get-ed trampolines.
	 */
	ids[9] = 0;

	opts.ids = ids;
	opts.cnt = cnt;

	skel->bss->pid = getpid();

	skel->links.test_fentry = bpf_program__attach_tracing_multi(skel->progs.test_fentry,
						NULL, &opts);
	if (!ASSERT_ERR_PTR(skel->links.test_fentry, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	skel->links.test_fexit = bpf_program__attach_tracing_multi(skel->progs.test_fexit,
						NULL, &opts);
	if (!ASSERT_ERR_PTR(skel->links.test_fexit, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	/* We don't really attach any program, but let's make sure. */
	tracing_multi_rollback_run(skel);

cleanup:
	tracing_multi_rollback__destroy(skel);
	free(ids);
}

static void fillers_cleanup(struct tracing_multi_rollback **skels, int cnt)
{
	int i;

	for (i = 0; i < cnt; i++)
		tracing_multi_rollback__destroy(skels[i]);

	free(skels);
}

static struct tracing_multi_rollback *extra_load_and_link(void)
{
	struct tracing_multi_rollback *skel;
	int err;

	skel = tracing_multi_rollback__open();
	if (!ASSERT_OK_PTR(skel, "tracing_multi_rollback__open"))
		goto cleanup;

	bpf_program__set_autoload(skel->progs.extra, true);

	err = tracing_multi_rollback__load(skel);
	if (!ASSERT_OK(err, "tracing_multi_rollback__load"))
		goto cleanup;

	skel->links.extra = bpf_program__attach_trace(skel->progs.extra);
	if (!ASSERT_OK_PTR(skel->links.extra, "bpf_program__attach_trace"))
		goto cleanup;

	return skel;

cleanup:
	tracing_multi_rollback__destroy(skel);
	return NULL;
}

static struct tracing_multi_rollback **fillers_load_and_link(int max)
{
	struct tracing_multi_rollback **skels, *skel;
	int i, err;

	skels = calloc(max + 1, sizeof(*skels));
	if (!ASSERT_OK_PTR(skels, "calloc"))
		return NULL;

	for (i = 0; i < max; i++) {
		skel = skels[i] = tracing_multi_rollback__open();
		if (!ASSERT_OK_PTR(skels[i], "tracing_multi_rollback__open"))
			goto cleanup;

		bpf_program__set_autoload(skel->progs.filler, true);

		err = tracing_multi_rollback__load(skel);
		if (!ASSERT_OK(err, "tracing_multi_rollback__load"))
			goto cleanup;

		skel->links.filler = bpf_program__attach_trace(skel->progs.filler);
		if (!ASSERT_OK_PTR(skels[i]->links.filler, "bpf_program__attach_trace"))
			goto cleanup;
	}

	return skels;

cleanup:
	fillers_cleanup(skels, i + 1);
	return NULL;
}

static void test_rollback_unlink(void)
{
	struct tracing_multi_rollback *skel = NULL, *extra;
	LIBBPF_OPTS(bpf_tracing_multi_opts, opts);
	struct tracing_multi_rollback **fillers;
	size_t cnt = FUNCS_CNT;
	__u32 *ids = NULL;
	int err, max;

	max = get_bpf_max_tramp_links();
	if (!ASSERT_GE(max, 1, "bpf_max_tramp_links"))
		return;

	/* Attach maximum allowed programs to bpf_fentry_test10 */
	fillers = fillers_load_and_link(max);
	if (!ASSERT_OK_PTR(fillers, "fillers_load_and_link"))
		return;

	extra = extra_load_and_link();
	if (!ASSERT_OK_PTR(extra, "extra_load_and_link"))
		goto cleanup;

	skel = tracing_multi_rollback__open();
	if (!ASSERT_OK_PTR(skel, "tracing_multi_rollback__open"))
		goto cleanup;

	bpf_program__set_autoload(skel->progs.test_fentry, true);
	bpf_program__set_autoload(skel->progs.test_fexit, true);

	/*
	 * Attach tracing_multi link on bpf_fentry_test1-10, which will
	 * fail on bpf_fentry_test10 function, because it already has
	 * maximum allowed programs attached.
	 *
	 * The rollback needs to unlink already link-ed trampolines and
	 * put all of them.
	 */
	err = tracing_multi_rollback__load(skel);
	if (!ASSERT_OK(err, "tracing_multi_rollback__load"))
		goto cleanup;

	ids = get_ids(bpf_fentry_test, cnt, NULL);
	if (!ASSERT_OK_PTR(ids, "get_ids"))
		goto cleanup;

	opts.ids = ids;
	opts.cnt = cnt;

	skel->bss->pid = getpid();

	skel->links.test_fentry = bpf_program__attach_tracing_multi(skel->progs.test_fentry,
						NULL, &opts);
	if (!ASSERT_ERR_PTR(skel->links.test_fentry, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	skel->links.test_fexit = bpf_program__attach_tracing_multi(skel->progs.test_fexit,
						NULL, &opts);
	if (!ASSERT_ERR_PTR(skel->links.test_fexit, "bpf_program__attach_tracing_multi"))
		goto cleanup;

	tracing_multi_rollback_run(skel);

cleanup:
	fillers_cleanup(fillers, max);
	tracing_multi_rollback__destroy(extra);
	tracing_multi_rollback__destroy(skel);
	free(ids);
}

void serial_test_tracing_multi_attach_rollback(void)
{
	if (test__start_subtest("put"))
		test_rollback_put();
	if (test__start_subtest("unlink"))
		test_rollback_unlink();
}

void test_tracing_multi_test(void)
{
#ifndef __x86_64__
	test__skip();
	return;
#endif

	if (test__start_subtest("skel_api"))
		test_skel_api();
	if (test__start_subtest("link_api_pattern"))
		test_link_api_pattern();
	if (test__start_subtest("link_api_ids"))
		test_link_api_ids(false);
	if (test__start_subtest("module_skel_api"))
		test_module_skel_api();
	if (test__start_subtest("module_link_api_pattern"))
		test_module_link_api_pattern();
	if (test__start_subtest("module_link_api_ids"))
		test_module_link_api_ids();
	if (test__start_subtest("intersect"))
		test_intersect();
	if (test__start_subtest("cookies"))
		test_link_api_ids(true);
	if (test__start_subtest("session"))
		test_session();
	if (test__start_subtest("attach_api_fails"))
		test_attach_api_fails();
	RUN_TESTS(tracing_multi_verifier);
}