// SPDX-License-Identifier: GPL-2.0-only
#include <fcntl.h>
#include <string.h>
#include <sys/ioctl.h>

#include "test_util.h"
#include "kvm_util.h"
#include "processor.h"
#include "sev.h"

#define BUFFER_SIZE     (PAGE_SIZE * 2)

static u8 *data;
static u8 src[BUFFER_SIZE] __aligned(PAGE_SIZE);
static u8 dst[BUFFER_SIZE] __aligned(PAGE_SIZE);

static void validate_dst(int i, int nr_bytes, u8 pattern)
{
        for ( ; i < nr_bytes; i++)
                TEST_ASSERT(dst[i] == pattern,
                            "Expected 0x%x at byte %u, got 0x%x",
                            pattern, i, dst[i]);
}

static void validate_buffers(void)
{
        int i;

        for (i = 0; i < BUFFER_SIZE; i++)
                TEST_ASSERT(src[i] == dst[i],
                            "Expected src[%u] (0x%x) == dst[%u] (0x%x)",
                            i, src[i], i, dst[i]);
}

static void ____test_sev_dbg(struct kvm_vm *vm, int i, int j, int nr_bytes)
{
        u8 pattern = guest_random_u32(&guest_rng);

        if (i + nr_bytes > BUFFER_SIZE || j + nr_bytes > BUFFER_SIZE)
                return;

        memset(&src[i], pattern, nr_bytes);
        sev_encrypt_memory(vm, &data[j], &src[i], nr_bytes);
        sev_decrypt_memory(vm, &dst[i], &data[j], nr_bytes);
        validate_buffers();
        validate_dst(i, nr_bytes, pattern);
}

static void __test_sev_dbg(struct kvm_vm *vm, int nr_bytes)
{
        /*
         * In a perfect world, all sizes at all combinations within the buffers
         * would be tested.  In reality, even this much testing is quite slow.
         * Target sizes and offsets around the chunk (16 bytes) and page (4096
         * bytes) sizes.
         */
        int x[] = { 1, 8, 15, 16, 23 };
        int p = PAGE_SIZE - 24;
        int i, j;

        ____test_sev_dbg(vm, 0, 0, nr_bytes);

        for (i = 0; i < ARRAY_SIZE(x); i++) {
                for (j = 0; j < ARRAY_SIZE(x); j++) {
                        ____test_sev_dbg(vm, x[i], x[j], nr_bytes);
                        ____test_sev_dbg(vm, x[i], p + x[j], nr_bytes);
                        ____test_sev_dbg(vm, p + x[i], x[j], nr_bytes);
                        ____test_sev_dbg(vm, p + x[i], p + x[j], nr_bytes);
                }
        }
}

static void test_sev_dbg(u32 type, u64 policy)
{
        int sizes[] = { 1, 8, 15, 16, 17, 32, 33 };
        struct kvm_vcpu *vcpu;
        struct kvm_vm *vm;
        int i;

        if (!(kvm_check_cap(KVM_CAP_VM_TYPES) & BIT(type)))
                return;

        vm = vm_sev_create_with_one_vcpu(type, NULL, &vcpu);

        data = addr_gva2hva(vm, vm_alloc(vm, BUFFER_SIZE, KVM_UTIL_MIN_VADDR));
        memset(data, 0xaa, BUFFER_SIZE);

        vm_sev_launch(vm, policy, NULL);

        sev_decrypt_memory(vm, dst, data, BUFFER_SIZE);
        validate_dst(0, BUFFER_SIZE, 0xaa);

        memset(src, 0x55, BUFFER_SIZE);
        sev_encrypt_memory(vm, data, src, BUFFER_SIZE);
        sev_decrypt_memory(vm, dst, data, BUFFER_SIZE);
        validate_dst(0, BUFFER_SIZE, 0x55);

        __test_sev_dbg(vm, PAGE_SIZE);

        for (i = 0; i < ARRAY_SIZE(sizes); i++) {
                __test_sev_dbg(vm, sizes[i]);
                __test_sev_dbg(vm, PAGE_SIZE - sizes[i]);
                __test_sev_dbg(vm, PAGE_SIZE + sizes[i]);
                __test_sev_dbg(vm, BUFFER_SIZE - sizes[i]);
        }

        kvm_vm_free(vm);
}

int main(int argc, char *argv[])
{
        TEST_REQUIRE(kvm_cpu_has(X86_FEATURE_SEV));

        /* Note, KVM doesn't support {de,en}crypt commands for SNP. */
        test_sev_dbg(KVM_X86_SEV_VM, 0);
        test_sev_dbg(KVM_X86_SEV_ES_VM, SEV_POLICY_ES);
        return 0;
}