// SPDX-License-Identifier: GPL-2.0
/* Copyright (c) 2025 Meta Platforms, Inc. and affiliates. */

#include <test_progs.h>
#include <unistd.h>
#include "test_sleepable_tracepoints.skel.h"
#include "test_sleepable_tracepoints_fail.skel.h"

static void run_test(struct test_sleepable_tracepoints *skel)
{
	char buf[PATH_MAX] = "/";

	skel->bss->target_pid = getpid();
	skel->bss->prog_triggered = 0;
	skel->bss->err = 0;
	skel->bss->copied_byte = 0;

	syscall(__NR_getcwd, buf, sizeof(buf));

	ASSERT_EQ(skel->bss->prog_triggered, 1, "prog_triggered");
	ASSERT_EQ(skel->bss->err, 0, "err");
	ASSERT_EQ(skel->bss->copied_byte, '/', "copied_byte");
}

static void run_auto_attach_test(struct bpf_program *prog,
				 struct test_sleepable_tracepoints *skel)
{
	struct bpf_link *link;

	link = bpf_program__attach(prog);
	if (!ASSERT_OK_PTR(link, "prog_attach"))
		return;

	run_test(skel);
	bpf_link__destroy(link);
}

static void test_attach_only(struct bpf_program *prog)
{
	struct bpf_link *link;

	link = bpf_program__attach(prog);
	if (ASSERT_OK_PTR(link, "attach"))
		bpf_link__destroy(link);
}

static void test_attach_reject(struct bpf_program *prog)
{
	struct bpf_link *link;

	link = bpf_program__attach(prog);
	if (!ASSERT_ERR_PTR(link, "attach_should_fail"))
		bpf_link__destroy(link);
}

static void test_raw_tp_bare(struct test_sleepable_tracepoints *skel)
{
	struct bpf_link *link;

	link = bpf_program__attach_raw_tracepoint(skel->progs.handle_raw_tp_bare,
						  "sys_enter");
	if (ASSERT_OK_PTR(link, "attach"))
		bpf_link__destroy(link);
}

static void test_tp_bare(struct test_sleepable_tracepoints *skel)
{
	struct bpf_link *link;

	link = bpf_program__attach_tracepoint(skel->progs.handle_tp_bare,
					      "syscalls", "sys_enter_getcwd");
	if (ASSERT_OK_PTR(link, "attach"))
		bpf_link__destroy(link);
}

static void test_test_run(struct test_sleepable_tracepoints *skel)
{
	__u64 args[2] = {0x1234ULL, 0x5678ULL};
	LIBBPF_OPTS(bpf_test_run_opts, topts,
		.ctx_in = args,
		.ctx_size_in = sizeof(args),
	);
	int fd, err;

	fd = bpf_program__fd(skel->progs.handle_test_run);
	err = bpf_prog_test_run_opts(fd, &topts);
	ASSERT_OK(err, "test_run");
	ASSERT_EQ(topts.retval, args[0] + args[1], "test_run_retval");
}

static void test_test_run_on_cpu_reject(struct test_sleepable_tracepoints *skel)
{
	__u64 args[2] = {};
	LIBBPF_OPTS(bpf_test_run_opts, topts,
		.ctx_in = args,
		.ctx_size_in = sizeof(args),
		.flags = BPF_F_TEST_RUN_ON_CPU,
	);
	int fd, err;

	fd = bpf_program__fd(skel->progs.handle_test_run);
	err = bpf_prog_test_run_opts(fd, &topts);
	ASSERT_ERR(err, "test_run_on_cpu_reject");
}

void test_sleepable_tracepoints(void)
{
	struct test_sleepable_tracepoints *skel;

	skel = test_sleepable_tracepoints__open_and_load();
	if (!ASSERT_OK_PTR(skel, "open_and_load"))
		return;

	if (test__start_subtest("tp_btf"))
		run_auto_attach_test(skel->progs.handle_sys_enter_tp_btf, skel);
	if (test__start_subtest("raw_tp"))
		run_auto_attach_test(skel->progs.handle_sys_enter_raw_tp, skel);
	if (test__start_subtest("tracepoint"))
		run_auto_attach_test(skel->progs.handle_sys_enter_tp, skel);
	if (test__start_subtest("sys_exit"))
		run_auto_attach_test(skel->progs.handle_sys_exit_tp, skel);
	if (test__start_subtest("tracepoint_alias"))
		test_attach_only(skel->progs.handle_sys_enter_tp_alias);
	if (test__start_subtest("raw_tracepoint_alias"))
		test_attach_only(skel->progs.handle_sys_enter_raw_tp_alias);
	if (test__start_subtest("raw_tp_bare"))
		test_raw_tp_bare(skel);
	if (test__start_subtest("tp_bare"))
		test_tp_bare(skel);
	if (test__start_subtest("test_run"))
		test_test_run(skel);
	if (test__start_subtest("test_run_on_cpu_reject"))
		test_test_run_on_cpu_reject(skel);
	if (test__start_subtest("raw_tp_non_faultable"))
		test_attach_reject(skel->progs.handle_raw_tp_non_faultable);
	if (test__start_subtest("tp_non_syscall"))
		test_attach_reject(skel->progs.handle_tp_non_syscall);
	if (test__start_subtest("tp_btf_non_faultable_reject"))
		RUN_TESTS(test_sleepable_tracepoints_fail);

	test_sleepable_tracepoints__destroy(skel);
}