// SPDX-License-Identifier: GPL-2.0-only
/*
 * Test for s390x KVM_S390_KEYOP
 *
 * Copyright IBM Corp. 2026
 *
 * Authors:
 *  Claudio Imbrenda <imbrenda@linux.ibm.com>
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>

#include <linux/bits.h>

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

#define BUF_PAGES 128UL
#define GUEST_PAGES 256UL

#define BUF_START_GFN   (GUEST_PAGES - BUF_PAGES)
#define BUF_START_ADDR  (BUF_START_GFN << PAGE_SHIFT)

#define KEY_BITS_ACC    0xf0
#define KEY_BIT_F       0x08
#define KEY_BIT_R       0x04
#define KEY_BIT_C       0x02

#define KEY_BITS_RC     (KEY_BIT_R | KEY_BIT_C)
#define KEY_BITS_ALL    (KEY_BITS_ACC | KEY_BIT_F | KEY_BITS_RC)

static unsigned char tmp[BUF_PAGES];
static unsigned char old[BUF_PAGES];
static unsigned char expected[BUF_PAGES];

static int _get_skeys(struct kvm_vcpu *vcpu, unsigned char skeys[])
{
        struct kvm_s390_skeys skeys_ioctl = {
                .start_gfn = BUF_START_GFN,
                .count = BUF_PAGES,
                .skeydata_addr = (unsigned long)skeys,
        };

        return __vm_ioctl(vcpu->vm, KVM_S390_GET_SKEYS, &skeys_ioctl);
}

static void get_skeys(struct kvm_vcpu *vcpu, unsigned char skeys[])
{
        int r = _get_skeys(vcpu, skeys);

        TEST_ASSERT(!r, "Failed to get storage keys, r=%d", r);
}

static void set_skeys(struct kvm_vcpu *vcpu, unsigned char skeys[])
{
        struct kvm_s390_skeys skeys_ioctl = {
                .start_gfn = BUF_START_GFN,
                .count = BUF_PAGES,
                .skeydata_addr = (unsigned long)skeys,
        };
        int r;

        r = __vm_ioctl(vcpu->vm, KVM_S390_SET_SKEYS, &skeys_ioctl);
        TEST_ASSERT(!r, "Failed to set storage keys, r=%d", r);
}

static int do_keyop(struct kvm_vcpu *vcpu, int op, unsigned long page_idx, unsigned char skey)
{
        struct kvm_s390_keyop keyop = {
                .guest_addr = BUF_START_ADDR + page_idx * PAGE_SIZE,
                .key = skey,
                .operation = op,
        };
        int r;

        r = __vm_ioctl(vcpu->vm, KVM_S390_KEYOP, &keyop);
        TEST_ASSERT(!r, "Failed to perform keyop, r=%d", r);
        TEST_ASSERT((keyop.key & 1) == 0,
                    "Last bit of key is 1, should be 0! page %lu, new key=%#x, old key=%#x",
                    page_idx, skey, keyop.key);

        return keyop.key;
}

static void fault_in_buffer(struct kvm_vcpu *vcpu, int where, int cur_loc)
{
        unsigned long i;
        int r;

        if (where != cur_loc)
                return;

        for (i = 0; i < BUF_PAGES; i++) {
                r = ioctl(vcpu->fd, KVM_S390_VCPU_FAULT, BUF_START_ADDR + i * PAGE_SIZE);
                TEST_ASSERT(!r, "Faulting in buffer page %lu, r=%d", i, r);
        }
}

static inline void set_pattern(unsigned char skeys[])
{
        int i;

        for (i = 0; i < BUF_PAGES; i++)
                skeys[i] = i << 1;
}

static void dump_sk(const unsigned char skeys[], const char *descr)
{
        int i, j;

        fprintf(stderr, "# %s:\n", descr);
        for (i = 0; i < BUF_PAGES; i += 32) {
                fprintf(stderr, "# %3d: ", i);
                for (j = 0; j < 32; j++)
                        fprintf(stderr, "%02x ", skeys[i + j]);
                fprintf(stderr, "\n");
        }
}

static inline void compare(const unsigned char what[], const unsigned char expected[],
                           const char *descr, int fault_in_loc)
{
        int i;

        for (i = 0; i < BUF_PAGES; i++) {
                if (expected[i] != what[i]) {
                        dump_sk(expected, "Expected");
                        dump_sk(what, "Got");
                }
                TEST_ASSERT(expected[i] == what[i],
                            "%s! fault-in location %d, page %d, expected %#x, got %#x",
                            descr, fault_in_loc, i, expected[i], what[i]);
        }
}

static inline void clear_all(void)
{
        memset(tmp, 0, BUF_PAGES);
        memset(old, 0, BUF_PAGES);
        memset(expected, 0, BUF_PAGES);
}

static void test_init(struct kvm_vcpu *vcpu, int fault_in)
{
        /* Set all storage keys to zero */
        fault_in_buffer(vcpu, fault_in, 1);
        set_skeys(vcpu, expected);

        fault_in_buffer(vcpu, fault_in, 2);
        get_skeys(vcpu, tmp);
        compare(tmp, expected, "Setting keys not zero", fault_in);

        /* Set storage keys to a sequential pattern */
        fault_in_buffer(vcpu, fault_in, 3);
        set_pattern(expected);
        set_skeys(vcpu, expected);

        fault_in_buffer(vcpu, fault_in, 4);
        get_skeys(vcpu, tmp);
        compare(tmp, expected, "Setting storage keys failed", fault_in);
}

static void test_rrbe(struct kvm_vcpu *vcpu, int fault_in)
{
        unsigned char k;
        int i;

        /* Set storage keys to a sequential pattern */
        fault_in_buffer(vcpu, fault_in, 1);
        set_pattern(expected);
        set_skeys(vcpu, expected);

        /* Call the RRBE KEYOP ioctl on each page and verify the result */
        fault_in_buffer(vcpu, fault_in, 2);
        for (i = 0; i < BUF_PAGES; i++) {
                k = do_keyop(vcpu, KVM_S390_KEYOP_RRBE, i, 0xff);
                TEST_ASSERT((expected[i] & KEY_BITS_RC) == k,
                            "Old R or C value mismatch! expected: %#x, got %#x",
                            expected[i] & KEY_BITS_RC, k);
                if (i == BUF_PAGES / 2)
                        fault_in_buffer(vcpu, fault_in, 3);
        }

        for (i = 0; i < BUF_PAGES; i++)
                expected[i] &= ~KEY_BIT_R;

        /* Verify that only the R bit has been cleared */
        fault_in_buffer(vcpu, fault_in, 4);
        get_skeys(vcpu, tmp);
        compare(tmp, expected, "New value mismatch", fault_in);
}

static void test_iske(struct kvm_vcpu *vcpu, int fault_in)
{
        int i;

        /* Set storage keys to a sequential pattern */
        fault_in_buffer(vcpu, fault_in, 1);
        set_pattern(expected);
        set_skeys(vcpu, expected);

        /* Call the ISKE KEYOP ioctl on each page and verify the result */
        fault_in_buffer(vcpu, fault_in, 2);
        for (i = 0; i < BUF_PAGES; i++) {
                tmp[i] = do_keyop(vcpu, KVM_S390_KEYOP_ISKE, i, 0xff);
                if (i == BUF_PAGES / 2)
                        fault_in_buffer(vcpu, fault_in, 3);
        }
        compare(tmp, expected, "Old value mismatch", fault_in);

        /* Check storage keys have not changed */
        fault_in_buffer(vcpu, fault_in, 4);
        get_skeys(vcpu, tmp);
        compare(tmp, expected, "Storage keys values changed", fault_in);
}

static void test_sske(struct kvm_vcpu *vcpu, int fault_in)
{
        int i;

        /* Set storage keys to a sequential pattern */
        fault_in_buffer(vcpu, fault_in, 1);
        set_pattern(tmp);
        set_skeys(vcpu, tmp);

        /* Call the SSKE KEYOP ioctl on each page and verify the result */
        fault_in_buffer(vcpu, fault_in, 2);
        for (i = 0; i < BUF_PAGES; i++) {
                expected[i] = ~tmp[i] & KEY_BITS_ALL;
                /* Set the new storage keys to be the bit-inversion of the previous ones */
                old[i] = do_keyop(vcpu, KVM_S390_KEYOP_SSKE, i, expected[i] | 1);
                if (i == BUF_PAGES / 2)
                        fault_in_buffer(vcpu, fault_in, 3);
        }
        compare(old, tmp, "Old value mismatch", fault_in);

        /* Verify that the storage keys have been set correctly */
        fault_in_buffer(vcpu, fault_in, 4);
        get_skeys(vcpu, tmp);
        compare(tmp, expected, "New value mismatch", fault_in);
}

static struct testdef {
        const char *name;
        void (*test)(struct kvm_vcpu *vcpu, int fault_in_location);
        int n_fault_in_locations;
} testplan[] = {
        { "Initialization", test_init, 5 },
        { "RRBE", test_rrbe, 5 },
        { "ISKE", test_iske, 5 },
        { "SSKE", test_sske, 5 },
};

static void run_test(void (*the_test)(struct kvm_vcpu *, int), int fault_in_location)
{
        struct kvm_vcpu *vcpu;
        struct kvm_vm *vm;
        int r;

        vm = vm_create_barebones();
        vm_userspace_mem_region_add(vm, VM_MEM_SRC_ANONYMOUS, 0, 0, GUEST_PAGES, 0);
        vcpu = __vm_vcpu_add(vm, 0);

        r = _get_skeys(vcpu, tmp);
        TEST_ASSERT(r == KVM_S390_GET_SKEYS_NONE,
                    "Storage keys are not disabled initially, r=%d", r);

        clear_all();

        the_test(vcpu, fault_in_location);

        kvm_vm_free(vm);
}

int main(int argc, char *argv[])
{
        int i, f;

        TEST_REQUIRE(kvm_has_cap(KVM_CAP_S390_KEYOP));
        TEST_REQUIRE(kvm_has_cap(KVM_CAP_S390_UCONTROL));

        ksft_print_header();
        for (i = 0, f = 0; i < ARRAY_SIZE(testplan); i++)
                f += testplan[i].n_fault_in_locations;
        ksft_set_plan(f);

        for (i = 0; i < ARRAY_SIZE(testplan); i++) {
                for (f = 0; f < testplan[i].n_fault_in_locations; f++) {
                        run_test(testplan[i].test, f);
                        ksft_test_result_pass("%s (fault-in location %d)\n", testplan[i].name, f);
                }
        }

        ksft_finished();        /* Print results and exit() accordingly */
}