// SPDX-License-Identifier: GPL-2.0-only
/*
 * CPU-agnostic ARM page table allocator.
 *
 * Copyright (C) 2014 ARM Limited
 *
 * Author: Will Deacon <will.deacon@arm.com>
 */

#define pr_fmt(fmt)     "arm-lpae io-pgtable: " fmt

#include <kunit/device.h>
#include <kunit/test.h>
#include <linux/io-pgtable.h>
#include <linux/kernel.h>

#include "io-pgtable-arm.h"

static struct io_pgtable_cfg *cfg_cookie;

static void dummy_tlb_flush_all(void *cookie)
{
        WARN_ON(cookie != cfg_cookie);
}

static void dummy_tlb_flush(unsigned long iova, size_t size,
                            size_t granule, void *cookie)
{
        WARN_ON(cookie != cfg_cookie);
        WARN_ON(!(size & cfg_cookie->pgsize_bitmap));
}

static void dummy_tlb_add_page(struct iommu_iotlb_gather *gather,
                               unsigned long iova, size_t granule,
                               void *cookie)
{
        dummy_tlb_flush(iova, granule, granule, cookie);
}

static const struct iommu_flush_ops dummy_tlb_ops = {
        .tlb_flush_all  = dummy_tlb_flush_all,
        .tlb_flush_walk = dummy_tlb_flush,
        .tlb_add_page   = dummy_tlb_add_page,
};

#define __FAIL(test, i) ({                                                      \
                KUNIT_FAIL(test, "test failed for fmt idx %d\n", (i));          \
                -EFAULT;                                                        \
})

static int arm_lpae_run_tests(struct kunit *test, struct io_pgtable_cfg *cfg)
{
        static const enum io_pgtable_fmt fmts[] = {
                ARM_64_LPAE_S1,
                ARM_64_LPAE_S2,
        };

        int i, j;
        unsigned long iova;
        size_t size, mapped;
        struct io_pgtable_ops *ops;

        for (i = 0; i < ARRAY_SIZE(fmts); ++i) {
                cfg_cookie = cfg;
                ops = alloc_io_pgtable_ops(fmts[i], cfg, cfg);
                if (!ops) {
                        kunit_err(test, "failed to allocate io pgtable ops\n");
                        return -ENOMEM;
                }

                /*
                 * Initial sanity checks.
                 * Empty page tables shouldn't provide any translations.
                 */
                if (ops->iova_to_phys(ops, 42))
                        return __FAIL(test, i);

                if (ops->iova_to_phys(ops, SZ_1G + 42))
                        return __FAIL(test, i);

                if (ops->iova_to_phys(ops, SZ_2G + 42))
                        return __FAIL(test, i);

                /*
                 * Distinct mappings of different granule sizes.
                 */
                iova = 0;
                for_each_set_bit(j, &cfg->pgsize_bitmap, BITS_PER_LONG) {
                        size = 1UL << j;

                        if (ops->map_pages(ops, iova, iova, size, 1,
                                           IOMMU_READ | IOMMU_WRITE |
                                           IOMMU_NOEXEC | IOMMU_CACHE,
                                           GFP_KERNEL, &mapped))
                                return __FAIL(test, i);

                        /* Overlapping mappings */
                        if (!ops->map_pages(ops, iova, iova + size, size, 1,
                                            IOMMU_READ | IOMMU_NOEXEC,
                                            GFP_KERNEL, &mapped))
                                return __FAIL(test, i);

                        if (ops->iova_to_phys(ops, iova + 42) != (iova + 42))
                                return __FAIL(test, i);

                        iova += SZ_1G;
                }

                /* Full unmap */
                iova = 0;
                for_each_set_bit(j, &cfg->pgsize_bitmap, BITS_PER_LONG) {
                        size = 1UL << j;

                        if (ops->unmap_pages(ops, iova, size, 1, NULL) != size)
                                return __FAIL(test, i);

                        if (ops->iova_to_phys(ops, iova + 42))
                                return __FAIL(test, i);

                        /* Remap full block */
                        if (ops->map_pages(ops, iova, iova, size, 1,
                                           IOMMU_WRITE, GFP_KERNEL, &mapped))
                                return __FAIL(test, i);

                        if (ops->iova_to_phys(ops, iova + 42) != (iova + 42))
                                return __FAIL(test, i);

                        iova += SZ_1G;
                }

                /*
                 * Map/unmap the last largest supported page of the IAS, this can
                 * trigger corner cases in the concatednated page tables.
                 */
                mapped = 0;
                size = 1UL << __fls(cfg->pgsize_bitmap);
                iova = (1UL << cfg->ias) - size;
                if (ops->map_pages(ops, iova, iova, size, 1,
                                   IOMMU_READ | IOMMU_WRITE |
                                   IOMMU_NOEXEC | IOMMU_CACHE,
                                   GFP_KERNEL, &mapped))
                        return __FAIL(test, i);
                if (mapped != size)
                        return __FAIL(test, i);
                if (ops->unmap_pages(ops, iova, size, 1, NULL) != size)
                        return __FAIL(test, i);

                free_io_pgtable_ops(ops);
        }

        return 0;
}

static void arm_lpae_do_selftests(struct kunit *test)
{
        static const unsigned long pgsize[] = {
                SZ_4K | SZ_2M | SZ_1G,
                SZ_16K | SZ_32M,
                SZ_64K | SZ_512M,
        };

        static const unsigned int address_size[] = {
                32, 36, 40, 42, 44, 48,
        };

        int i, j, k, pass = 0, fail = 0;
        struct device *dev;
        struct io_pgtable_cfg cfg = {
                .tlb = &dummy_tlb_ops,
                .coherent_walk = true,
                .quirks = IO_PGTABLE_QUIRK_NO_WARN,
        };

        dev = kunit_device_register(test, "io-pgtable-test");
        KUNIT_EXPECT_NOT_ERR_OR_NULL(test, dev);
        if (IS_ERR_OR_NULL(dev))
                return;

        cfg.iommu_dev = dev;

        for (i = 0; i < ARRAY_SIZE(pgsize); ++i) {
                for (j = 0; j < ARRAY_SIZE(address_size); ++j) {
                        /* Don't use ias > oas as it is not valid for stage-2. */
                        for (k = 0; k <= j; ++k) {
                                cfg.pgsize_bitmap = pgsize[i];
                                cfg.ias = address_size[k];
                                cfg.oas = address_size[j];
                                kunit_info(test, "pgsize_bitmap 0x%08lx, IAS %u OAS %u\n",
                                           pgsize[i], cfg.ias, cfg.oas);
                                if (arm_lpae_run_tests(test, &cfg))
                                        fail++;
                                else
                                        pass++;
                        }
                }
        }

        kunit_info(test, "completed with %d PASS %d FAIL\n", pass, fail);
}

static struct kunit_case io_pgtable_arm_test_cases[] = {
        KUNIT_CASE(arm_lpae_do_selftests),
        {},
};

static struct kunit_suite io_pgtable_arm_test = {
        .name = "io-pgtable-arm-test",
        .test_cases = io_pgtable_arm_test_cases,
};

kunit_test_suite(io_pgtable_arm_test);

MODULE_DESCRIPTION("io-pgtable-arm library kunit tests");
MODULE_LICENSE("GPL");