// SPDX-License-Identifier: GPL-2.0-only
#include <stdio.h>
#include <sys/mman.h>
#include <unistd.h>

#include <uapi/linux/types.h>
#include <linux/pci_regs.h>
#include <linux/sizes.h>
#include <linux/vfio.h>

#include <libvfio.h>

#include "../kselftest_harness.h"

static const char *device_bdf;

static struct vfio_pci_bar *largest_mapped_bar(struct vfio_pci_device *device)
{
        u32 flags = VFIO_REGION_INFO_FLAG_READ | VFIO_REGION_INFO_FLAG_WRITE;
        struct vfio_pci_bar *largest = NULL;
        u64 bar_size = 0;

        for (int i = 0; i < PCI_STD_NUM_BARS; i++) {
                struct vfio_pci_bar *bar = &device->bars[i];

                if (!bar->vaddr)
                        continue;

                /*
                 * iommu_map() maps with READ|WRITE, so require the same
                 * abilities for the underlying VFIO region.
                 */
                if ((bar->info.flags & flags) != flags)
                        continue;

                if (bar->info.size > bar_size) {
                        bar_size = bar->info.size;
                        largest = bar;
                }
        }

        return largest;
}

FIXTURE(vfio_dma_mapping_mmio_test) {
        struct iommu *iommu;
        struct vfio_pci_device *device;
        struct iova_allocator *iova_allocator;
        struct vfio_pci_bar *bar;
};

FIXTURE_VARIANT(vfio_dma_mapping_mmio_test) {
        const char *iommu_mode;
};

#define FIXTURE_VARIANT_ADD_IOMMU_MODE(_iommu_mode)                            \
FIXTURE_VARIANT_ADD(vfio_dma_mapping_mmio_test, _iommu_mode) {                 \
        .iommu_mode = #_iommu_mode,                                            \
}

FIXTURE_VARIANT_ADD_ALL_IOMMU_MODES();

#undef FIXTURE_VARIANT_ADD_IOMMU_MODE

FIXTURE_SETUP(vfio_dma_mapping_mmio_test)
{
        self->iommu = iommu_init(variant->iommu_mode);
        self->device = vfio_pci_device_init(device_bdf, self->iommu);
        self->iova_allocator = iova_allocator_init(self->iommu);
        self->bar = largest_mapped_bar(self->device);

        if (!self->bar)
                SKIP(return, "No mappable BAR found on device %s", device_bdf);
}

FIXTURE_TEARDOWN(vfio_dma_mapping_mmio_test)
{
        iova_allocator_cleanup(self->iova_allocator);
        vfio_pci_device_cleanup(self->device);
        iommu_cleanup(self->iommu);
}

static void do_mmio_map_test(struct iommu *iommu,
                             struct iova_allocator *iova_allocator,
                             void *vaddr, size_t size)
{
        struct dma_region region = {
                .vaddr = vaddr,
                .size = size,
                .iova = iova_allocator_alloc(iova_allocator, size),
        };

        /*
         * NOTE: Check for iommufd compat success once it lands. Native iommufd
         * will never support this.
         */
        if (!strcmp(iommu->mode->name, MODE_VFIO_TYPE1V2_IOMMU) ||
            !strcmp(iommu->mode->name, MODE_VFIO_TYPE1_IOMMU)) {
                iommu_map(iommu, &region);
                iommu_unmap(iommu, &region);
        } else {
                VFIO_ASSERT_NE(__iommu_map(iommu, &region), 0);
                VFIO_ASSERT_NE(__iommu_unmap(iommu, &region, NULL), 0);
        }
}

TEST_F(vfio_dma_mapping_mmio_test, map_full_bar)
{
        do_mmio_map_test(self->iommu, self->iova_allocator,
                         self->bar->vaddr, self->bar->info.size);
}

TEST_F(vfio_dma_mapping_mmio_test, map_partial_bar)
{
        if (self->bar->info.size < 2 * getpagesize())
                SKIP(return, "BAR too small (size=0x%llx)", self->bar->info.size);

        do_mmio_map_test(self->iommu, self->iova_allocator,
                         self->bar->vaddr, getpagesize());
}

/* Test IOMMU mapping of BAR mmap with intentionally poor vaddr alignment. */
TEST_F(vfio_dma_mapping_mmio_test, map_bar_misaligned)
{
        /* Limit size to bound test time for large BARs */
        size_t size = min_t(size_t, self->bar->info.size, SZ_1G);
        void *vaddr;

        vaddr = mmap_reserve(size, SZ_1G, getpagesize());
        vaddr = mmap(vaddr, size, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_FIXED,
                     self->device->fd, self->bar->info.offset);
        VFIO_ASSERT_NE(vaddr, MAP_FAILED);

        do_mmio_map_test(self->iommu, self->iova_allocator, vaddr, size);

        VFIO_ASSERT_EQ(munmap(vaddr, size), 0);
}

int main(int argc, char *argv[])
{
        device_bdf = vfio_selftests_get_bdf(&argc, argv);
        return test_harness_run(argc, argv);
}