// SPDX-License-Identifier: GPL-2.0

/*
 * Copyright (c) 2025, Google LLC.
 * Pasha Tatashin <pasha.tatashin@soleen.com>
 */

#define pr_fmt(fmt) KBUILD_MODNAME " test: " fmt

#include <linux/cleanup.h>
#include <linux/errno.h>
#include <linux/init.h>
#include <linux/liveupdate.h>
#include <linux/module.h>
#include "../../kernel/liveupdate/luo_internal.h"

static const struct liveupdate_flb_ops test_flb_ops;
#define DEFINE_TEST_FLB(i) {                                            \
        .ops = &test_flb_ops,                                           \
        .compatible = LIVEUPDATE_TEST_FLB_COMPATIBLE(i),                \
}

/* Number of Test FLBs to register with every file handler */
#define TEST_NFLBS 3
static struct liveupdate_flb test_flbs[TEST_NFLBS] = {
        DEFINE_TEST_FLB(0),
        DEFINE_TEST_FLB(1),
        DEFINE_TEST_FLB(2),
};

#define TEST_FLB_MAGIC_BASE 0xFEEDF00DCAFEBEE0ULL

static int test_flb_preserve(struct liveupdate_flb_op_args *argp)
{
        ptrdiff_t index = argp->flb - test_flbs;

        pr_info("%s: preserve was triggered\n", argp->flb->compatible);
        argp->data = TEST_FLB_MAGIC_BASE + index;

        return 0;
}

static void test_flb_unpreserve(struct liveupdate_flb_op_args *argp)
{
        pr_info("%s: unpreserve was triggered\n", argp->flb->compatible);
}

static int test_flb_retrieve(struct liveupdate_flb_op_args *argp)
{
        ptrdiff_t index = argp->flb - test_flbs;
        u64 expected_data = TEST_FLB_MAGIC_BASE + index;

        if (argp->data == expected_data) {
                pr_info("%s: found flb data from the previous boot\n",
                        argp->flb->compatible);
                argp->obj = (void *)argp->data;
        } else {
                pr_err("%s: ERROR - incorrect data handle: %llx, expected %llx\n",
                       argp->flb->compatible, argp->data, expected_data);
                return -EINVAL;
        }

        return 0;
}

static void test_flb_finish(struct liveupdate_flb_op_args *argp)
{
        ptrdiff_t index = argp->flb - test_flbs;
        void *expected_obj = (void *)(TEST_FLB_MAGIC_BASE + index);

        if (argp->obj == expected_obj) {
                pr_info("%s: finish was triggered\n", argp->flb->compatible);
        } else {
                pr_err("%s: ERROR - finish called with invalid object\n",
                       argp->flb->compatible);
        }
}

static const struct liveupdate_flb_ops test_flb_ops = {
        .preserve       = test_flb_preserve,
        .unpreserve     = test_flb_unpreserve,
        .retrieve       = test_flb_retrieve,
        .finish         = test_flb_finish,
        .owner          = THIS_MODULE,
};

static void liveupdate_test_init(void)
{
        static DEFINE_MUTEX(init_lock);
        static bool initialized;
        int i;

        guard(mutex)(&init_lock);

        if (initialized)
                return;

        for (i = 0; i < TEST_NFLBS; i++) {
                struct liveupdate_flb *flb = &test_flbs[i];
                void *obj;
                int err;

                err = liveupdate_flb_get_incoming(flb, &obj);
                if (err && err != -ENODATA && err != -ENOENT) {
                        pr_err("liveupdate_flb_get_incoming for %s failed: %pe\n",
                               flb->compatible, ERR_PTR(err));
                }
        }
        initialized = true;
}

void liveupdate_test_register(struct liveupdate_file_handler *fh)
{
        int err, i;

        liveupdate_test_init();

        for (i = 0; i < TEST_NFLBS; i++) {
                struct liveupdate_flb *flb = &test_flbs[i];

                err = liveupdate_register_flb(fh, flb);
                if (err) {
                        pr_err("Failed to register %s %pe\n",
                               flb->compatible, ERR_PTR(err));
                }
        }

        err = liveupdate_register_flb(fh, &test_flbs[0]);
        if (!err || err != -EEXIST) {
                pr_err("Failed: %s should be already registered, but got err: %pe\n",
                       test_flbs[0].compatible, ERR_PTR(err));
        }

        pr_info("Registered %d FLBs with file handler: [%s]\n",
                TEST_NFLBS, fh->compatible);
}

void liveupdate_test_unregister(struct liveupdate_file_handler *fh)
{
        int err, i;

        for (i = 0; i < TEST_NFLBS; i++) {
                struct liveupdate_flb *flb = &test_flbs[i];

                err = liveupdate_unregister_flb(fh, flb);
                if (err) {
                        pr_err("Failed to unregister %s %pe\n",
                               flb->compatible, ERR_PTR(err));
                }
        }

        pr_info("Unregistered %d FLBs from file handler: [%s]\n",
                TEST_NFLBS, fh->compatible);
}

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Pasha Tatashin <pasha.tatashin@soleen.com>");
MODULE_DESCRIPTION("In-kernel test for LUO mechanism");