// SPDX-License-Identifier: GPL-2.0

#include <test_progs.h>
#include "struct_ops_assoc.skel.h"
#include "struct_ops_assoc_reuse.skel.h"
#include "struct_ops_assoc_in_timer.skel.h"

static void test_st_ops_assoc(void)
{
        struct struct_ops_assoc *skel = NULL;
        int err, pid;

        skel = struct_ops_assoc__open_and_load();
        if (!ASSERT_OK_PTR(skel, "struct_ops_assoc__open"))
                goto out;

        /* cannot explicitly associate struct_ops program */
        err = bpf_program__assoc_struct_ops(skel->progs.test_1_a,
                                            skel->maps.st_ops_map_a, NULL);
        ASSERT_ERR(err, "bpf_program__assoc_struct_ops(test_1_a, st_ops_map_a)");

        err = bpf_program__assoc_struct_ops(skel->progs.syscall_prog_a,
                                            skel->maps.st_ops_map_a, NULL);
        ASSERT_OK(err, "bpf_program__assoc_struct_ops(syscall_prog_a, st_ops_map_a)");

        err = bpf_program__assoc_struct_ops(skel->progs.sys_enter_prog_a,
                                            skel->maps.st_ops_map_a, NULL);
        ASSERT_OK(err, "bpf_program__assoc_struct_ops(sys_enter_prog_a, st_ops_map_a)");

        err = bpf_program__assoc_struct_ops(skel->progs.syscall_prog_b,
                                            skel->maps.st_ops_map_b, NULL);
        ASSERT_OK(err, "bpf_program__assoc_struct_ops(syscall_prog_b, st_ops_map_b)");

        err = bpf_program__assoc_struct_ops(skel->progs.sys_enter_prog_b,
                                            skel->maps.st_ops_map_b, NULL);
        ASSERT_OK(err, "bpf_program__assoc_struct_ops(sys_enter_prog_b, st_ops_map_b)");

        /* sys_enter_prog_a already associated with map_a */
        err = bpf_program__assoc_struct_ops(skel->progs.sys_enter_prog_a,
                                            skel->maps.st_ops_map_b, NULL);
        ASSERT_ERR(err, "bpf_program__assoc_struct_ops(sys_enter_prog_a, st_ops_map_b)");

        err = struct_ops_assoc__attach(skel);
        if (!ASSERT_OK(err, "struct_ops_assoc__attach"))
                goto out;

        /* run tracing prog that calls .test_1 and checks return */
        pid = getpid();
        skel->bss->test_pid = pid;
        sys_gettid();
        skel->bss->test_pid = 0;

        ASSERT_EQ(skel->bss->test_err_a, 0, "skel->bss->test_err_a");
        ASSERT_EQ(skel->bss->test_err_b, 0, "skel->bss->test_err_b");

        /* run syscall_prog that calls .test_1 and checks return */
        err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.syscall_prog_a), NULL);
        ASSERT_OK(err, "bpf_prog_test_run_opts");

        err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.syscall_prog_b), NULL);
        ASSERT_OK(err, "bpf_prog_test_run_opts");

        ASSERT_EQ(skel->bss->test_err_a, 0, "skel->bss->test_err_a");
        ASSERT_EQ(skel->bss->test_err_b, 0, "skel->bss->test_err_b");

out:
        struct_ops_assoc__destroy(skel);
}

static void test_st_ops_assoc_reuse(void)
{
        struct struct_ops_assoc_reuse *skel = NULL;
        int err;

        skel = struct_ops_assoc_reuse__open_and_load();
        if (!ASSERT_OK_PTR(skel, "struct_ops_assoc_reuse__open"))
                goto out;

        err = bpf_program__assoc_struct_ops(skel->progs.syscall_prog_a,
                                            skel->maps.st_ops_map_a, NULL);
        ASSERT_OK(err, "bpf_program__assoc_struct_ops(syscall_prog_a, st_ops_map_a)");

        err = bpf_program__assoc_struct_ops(skel->progs.syscall_prog_b,
                                            skel->maps.st_ops_map_b, NULL);
        ASSERT_OK(err, "bpf_program__assoc_struct_ops(syscall_prog_b, st_ops_map_b)");

        err = struct_ops_assoc_reuse__attach(skel);
        if (!ASSERT_OK(err, "struct_ops_assoc__attach"))
                goto out;

        /* run syscall_prog that calls .test_1 and checks return */
        err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.syscall_prog_a), NULL);
        ASSERT_OK(err, "bpf_prog_test_run_opts");

        err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.syscall_prog_b), NULL);
        ASSERT_OK(err, "bpf_prog_test_run_opts");

        ASSERT_EQ(skel->bss->test_err_a, 0, "skel->bss->test_err_a");
        ASSERT_EQ(skel->bss->test_err_b, 0, "skel->bss->test_err_b");

out:
        struct_ops_assoc_reuse__destroy(skel);
}

static void test_st_ops_assoc_in_timer(void)
{
        struct struct_ops_assoc_in_timer *skel = NULL;
        int err;

        skel = struct_ops_assoc_in_timer__open_and_load();
        if (!ASSERT_OK_PTR(skel, "struct_ops_assoc_in_timer__open"))
                goto out;

        err = bpf_program__assoc_struct_ops(skel->progs.syscall_prog,
                                            skel->maps.st_ops_map, NULL);
        ASSERT_OK(err, "bpf_program__assoc_struct_ops");

        err = struct_ops_assoc_in_timer__attach(skel);
        if (!ASSERT_OK(err, "struct_ops_assoc__attach"))
                goto out;

        /*
         * Run .test_1 by calling kfunc bpf_kfunc_multi_st_ops_test_1_prog_arg() and checks
         * the return value. .test_1 will also schedule timer_cb that runs .test_1 again
         * immediately.
         */
        err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.syscall_prog), NULL);
        ASSERT_OK(err, "bpf_prog_test_run_opts");

        /* Check the return of the kfunc after timer_cb runs */
        while (!READ_ONCE(skel->bss->timer_cb_run))
                sched_yield();
        ASSERT_EQ(skel->bss->timer_test_1_ret, 1234, "skel->bss->timer_test_1_ret");
        ASSERT_EQ(skel->bss->test_err, 0, "skel->bss->test_err_a");
out:
        struct_ops_assoc_in_timer__destroy(skel);
}

static void test_st_ops_assoc_in_timer_no_uref(void)
{
        struct struct_ops_assoc_in_timer *skel = NULL;
        struct bpf_link *link;
        int err;

        skel = struct_ops_assoc_in_timer__open_and_load();
        if (!ASSERT_OK_PTR(skel, "struct_ops_assoc_in_timer__open"))
                goto out;

        err = bpf_program__assoc_struct_ops(skel->progs.syscall_prog,
                                            skel->maps.st_ops_map, NULL);
        ASSERT_OK(err, "bpf_program__assoc_struct_ops");

        link = bpf_map__attach_struct_ops(skel->maps.st_ops_map);
        if (!ASSERT_OK_PTR(link, "bpf_map__attach_struct_ops"))
                goto out;

        /*
         * Run .test_1 by calling kfunc bpf_kfunc_multi_st_ops_test_1_prog_arg() and checks
         * the return value. .test_1 will also schedule timer_cb that runs .test_1 again.
         * timer_cb will run 500ms after syscall_prog runs, when the user space no longer
         * holds a reference to st_ops_map.
         */
        skel->bss->timer_ns = 500000000;
        err = bpf_prog_test_run_opts(bpf_program__fd(skel->progs.syscall_prog), NULL);
        ASSERT_OK(err, "bpf_prog_test_run_opts");

        /* Detach and close struct_ops map to cause it to be freed */
        bpf_link__destroy(link);
        close(bpf_program__fd(skel->progs.syscall_prog));
        close(bpf_map__fd(skel->maps.st_ops_map));

        /* Check the return of the kfunc after timer_cb runs */
        while (!READ_ONCE(skel->bss->timer_cb_run))
                sched_yield();
        ASSERT_EQ(skel->bss->timer_test_1_ret, -1, "skel->bss->timer_test_1_ret");
        ASSERT_EQ(skel->bss->test_err, 0, "skel->bss->test_err_a");
out:
        struct_ops_assoc_in_timer__destroy(skel);
}

void test_struct_ops_assoc(void)
{
        if (test__start_subtest("st_ops_assoc"))
                test_st_ops_assoc();
        if (test__start_subtest("st_ops_assoc_reuse"))
                test_st_ops_assoc_reuse();
        if (test__start_subtest("st_ops_assoc_in_timer"))
                test_st_ops_assoc_in_timer();
        if (test__start_subtest("st_ops_assoc_in_timer_no_uref"))
                test_st_ops_assoc_in_timer_no_uref();
}