// SPDX-License-Identifier: GPL-2.0-or-later
/*
 * test_fprobe.c - simple sanity test for fprobe
 */

#include <linux/kernel.h>
#include <linux/fprobe.h>
#include <linux/random.h>
#include <kunit/test.h>

#define div_factor 3

static struct kunit *current_test;

static u32 rand1, entry_only_val, entry_val, exit_val;
static u32 entry_only_count, entry_count, exit_count;

/* Use indirect calls to avoid inlining the target functions */
static u32 (*target)(u32 value);
static u32 (*target2)(u32 value);
static unsigned long target_ip;
static unsigned long target2_ip;
static int entry_return_value;

static noinline u32 fprobe_selftest_target(u32 value)
{
        return (value / div_factor);
}

static noinline u32 fprobe_selftest_target2(u32 value)
{
        return (value / div_factor) + 1;
}

static notrace int fp_entry_handler(struct fprobe *fp, unsigned long ip,
                                    unsigned long ret_ip,
                                    struct ftrace_regs *fregs, void *data)
{
        KUNIT_EXPECT_FALSE(current_test, preemptible());
        /* This can be called on the fprobe_selftest_target and the fprobe_selftest_target2 */
        if (ip != target_ip)
                KUNIT_EXPECT_EQ(current_test, ip, target2_ip);
        entry_val = (rand1 / div_factor);
        if (fp->entry_data_size) {
                KUNIT_EXPECT_NOT_NULL(current_test, data);
                if (data)
                        *(u32 *)data = entry_val;
        } else
                KUNIT_EXPECT_NULL(current_test, data);

        return entry_return_value;
}

static notrace void fp_exit_handler(struct fprobe *fp, unsigned long ip,
                                    unsigned long ret_ip,
                                    struct ftrace_regs *fregs, void *data)
{
        unsigned long ret = ftrace_regs_get_return_value(fregs);

        KUNIT_EXPECT_FALSE(current_test, preemptible());
        if (ip != target_ip) {
                KUNIT_EXPECT_EQ(current_test, ip, target2_ip);
                KUNIT_EXPECT_EQ(current_test, ret, (rand1 / div_factor) + 1);
        } else
                KUNIT_EXPECT_EQ(current_test, ret, (rand1 / div_factor));
        KUNIT_EXPECT_EQ(current_test, entry_val, (rand1 / div_factor));
        exit_val = entry_val + div_factor;
        if (fp->entry_data_size) {
                KUNIT_EXPECT_NOT_NULL(current_test, data);
                if (data)
                        KUNIT_EXPECT_EQ(current_test, *(u32 *)data, entry_val);
        } else
                KUNIT_EXPECT_NULL(current_test, data);
}

/* Test entry only (no rethook) */
static void test_fprobe_entry(struct kunit *test)
{
        struct fprobe fp_entry = {
                .entry_handler = fp_entry_handler,
        };

        current_test = test;

        /* Before register, unregister should be failed. */
        KUNIT_EXPECT_NE(test, 0, unregister_fprobe(&fp_entry));
        KUNIT_EXPECT_EQ(test, 0, register_fprobe(&fp_entry, "fprobe_selftest_target*", NULL));

        entry_val = 0;
        exit_val = 0;
        target(rand1);
        KUNIT_EXPECT_NE(test, 0, entry_val);
        KUNIT_EXPECT_EQ(test, 0, exit_val);

        entry_val = 0;
        exit_val = 0;
        target2(rand1);
        KUNIT_EXPECT_NE(test, 0, entry_val);
        KUNIT_EXPECT_EQ(test, 0, exit_val);

        KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp_entry));
}

static void test_fprobe(struct kunit *test)
{
        struct fprobe fp = {
                .entry_handler = fp_entry_handler,
                .exit_handler = fp_exit_handler,
        };

        current_test = test;
        KUNIT_EXPECT_EQ(test, 0, register_fprobe(&fp, "fprobe_selftest_target*", NULL));

        entry_val = 0;
        exit_val = 0;
        target(rand1);
        KUNIT_EXPECT_NE(test, 0, entry_val);
        KUNIT_EXPECT_EQ(test, entry_val + div_factor, exit_val);

        entry_val = 0;
        exit_val = 0;
        target2(rand1);
        KUNIT_EXPECT_NE(test, 0, entry_val);
        KUNIT_EXPECT_EQ(test, entry_val + div_factor, exit_val);

        KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp));
}

static void test_fprobe_syms(struct kunit *test)
{
        static const char *syms[] = {"fprobe_selftest_target", "fprobe_selftest_target2"};
        struct fprobe fp = {
                .entry_handler = fp_entry_handler,
                .exit_handler = fp_exit_handler,
        };

        current_test = test;
        KUNIT_EXPECT_EQ(test, 0, register_fprobe_syms(&fp, syms, 2));

        entry_val = 0;
        exit_val = 0;
        target(rand1);
        KUNIT_EXPECT_NE(test, 0, entry_val);
        KUNIT_EXPECT_EQ(test, entry_val + div_factor, exit_val);

        entry_val = 0;
        exit_val = 0;
        target2(rand1);
        KUNIT_EXPECT_NE(test, 0, entry_val);
        KUNIT_EXPECT_EQ(test, entry_val + div_factor, exit_val);

        KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp));
}

/* Test private entry_data */
static void test_fprobe_data(struct kunit *test)
{
        struct fprobe fp = {
                .entry_handler = fp_entry_handler,
                .exit_handler = fp_exit_handler,
                .entry_data_size = sizeof(u32),
        };

        current_test = test;
        KUNIT_EXPECT_EQ(test, 0, register_fprobe(&fp, "fprobe_selftest_target", NULL));

        target(rand1);

        KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp));
}

static void test_fprobe_skip(struct kunit *test)
{
        struct fprobe fp = {
                .entry_handler = fp_entry_handler,
                .exit_handler = fp_exit_handler,
        };

        current_test = test;
        KUNIT_EXPECT_EQ(test, 0, register_fprobe(&fp, "fprobe_selftest_target", NULL));

        entry_return_value = 1;
        entry_val = 0;
        exit_val = 0;
        target(rand1);
        KUNIT_EXPECT_NE(test, 0, entry_val);
        KUNIT_EXPECT_EQ(test, 0, exit_val);
        KUNIT_EXPECT_EQ(test, 0, fp.nmissed);
        entry_return_value = 0;

        KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp));
}

/* Handler for fprobe entry only case */
static notrace int entry_only_handler(struct fprobe *fp, unsigned long ip,
                                      unsigned long ret_ip,
                                      struct ftrace_regs *fregs, void *data)
{
        KUNIT_EXPECT_FALSE(current_test, preemptible());
        KUNIT_EXPECT_EQ(current_test, ip, target_ip);

        entry_only_count++;
        entry_only_val = (rand1 / div_factor);

        return 0;
}

static notrace int fprobe_entry_multi_handler(struct fprobe *fp, unsigned long ip,
                                              unsigned long ret_ip,
                                              struct ftrace_regs *fregs,
                                              void *data)
{
        KUNIT_EXPECT_FALSE(current_test, preemptible());
        KUNIT_EXPECT_EQ(current_test, ip, target_ip);

        entry_count++;
        entry_val = (rand1 / div_factor);

        return 0;
}

static notrace void fprobe_exit_multi_handler(struct fprobe *fp, unsigned long ip,
                                              unsigned long ret_ip,
                                              struct ftrace_regs *fregs,
                                              void *data)
{
        unsigned long ret = ftrace_regs_get_return_value(fregs);

        KUNIT_EXPECT_FALSE(current_test, preemptible());
        KUNIT_EXPECT_EQ(current_test, ip, target_ip);
        KUNIT_EXPECT_EQ(current_test, ret, (rand1 / div_factor));

        exit_count++;
        exit_val = ret;
}

static void check_fprobe_multi(struct kunit *test)
{
        entry_only_count = entry_count = exit_count = 0;
        entry_only_val = entry_val = exit_val = 0;

        target(rand1);

        /* Verify all handlers were called */
        KUNIT_EXPECT_EQ(test, 1, entry_only_count);
        KUNIT_EXPECT_EQ(test, 1, entry_count);
        KUNIT_EXPECT_EQ(test, 1, exit_count);

        /* Verify values are correct */
        KUNIT_EXPECT_EQ(test, (rand1 / div_factor), entry_only_val);
        KUNIT_EXPECT_EQ(test, (rand1 / div_factor), entry_val);
        KUNIT_EXPECT_EQ(test, (rand1 / div_factor), exit_val);
}

/* Test multiple fprobes hooking the same target function */
static void test_fprobe_multi(struct kunit *test)
{
        struct fprobe fp1 = {
                .entry_handler = fprobe_entry_multi_handler,
                .exit_handler = fprobe_exit_multi_handler,
        };
        struct fprobe fp2 = {
                .entry_handler = entry_only_handler,
        };

        current_test = test;

        /* Test Case 1: Register in order 1 -> 2 */
        KUNIT_EXPECT_EQ(test, 0, register_fprobe(&fp1, "fprobe_selftest_target", NULL));
        KUNIT_EXPECT_EQ(test, 0, register_fprobe(&fp2, "fprobe_selftest_target", NULL));

        check_fprobe_multi(test);

        /* Unregister all */
        KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp1));
        KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp2));

        /* Test Case 2: Register in order 2 -> 1 */
        KUNIT_EXPECT_EQ(test, 0, register_fprobe(&fp2, "fprobe_selftest_target", NULL));
        KUNIT_EXPECT_EQ(test, 0, register_fprobe(&fp1, "fprobe_selftest_target", NULL));

        check_fprobe_multi(test);

        /* Unregister all */
        KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp1));
        KUNIT_EXPECT_EQ(test, 0, unregister_fprobe(&fp2));
}

static unsigned long get_ftrace_location(void *func)
{
        unsigned long size, addr = (unsigned long)func;

        if (!kallsyms_lookup_size_offset(addr, &size, NULL) || !size)
                return 0;

        return ftrace_location_range(addr, addr + size - 1);
}

static int fprobe_test_init(struct kunit *test)
{
        rand1 = get_random_u32_above(div_factor);
        target = fprobe_selftest_target;
        target2 = fprobe_selftest_target2;
        target_ip = get_ftrace_location(target);
        target2_ip = get_ftrace_location(target2);

        return 0;
}

static struct kunit_case fprobe_testcases[] = {
        KUNIT_CASE(test_fprobe_entry),
        KUNIT_CASE(test_fprobe),
        KUNIT_CASE(test_fprobe_syms),
        KUNIT_CASE(test_fprobe_data),
        KUNIT_CASE(test_fprobe_skip),
        KUNIT_CASE(test_fprobe_multi),
        {}
};

static struct kunit_suite fprobe_test_suite = {
        .name = "fprobe_test",
        .init = fprobe_test_init,
        .test_cases = fprobe_testcases,
};

kunit_test_suites(&fprobe_test_suite);