// SPDX-License-Identifier: GPL-2.0-or-later
/*
 * Test for s390x KVM_S390_MEM_OP
 *
 * Copyright (C) 2019, Red Hat, Inc.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <pthread.h>

#include <linux/bits.h>

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

enum mop_target {
	LOGICAL,
	SIDA,
	ABSOLUTE,
	INVALID,
};

enum mop_access_mode {
	READ,
	WRITE,
	CMPXCHG,
};

struct mop_desc {
	uintptr_t gaddr;
	uintptr_t gaddr_v;
	uint64_t set_flags;
	unsigned int f_check : 1;
	unsigned int f_inject : 1;
	unsigned int f_key : 1;
	unsigned int _gaddr_v : 1;
	unsigned int _set_flags : 1;
	unsigned int _sida_offset : 1;
	unsigned int _ar : 1;
	uint32_t size;
	enum mop_target target;
	enum mop_access_mode mode;
	void *buf;
	uint32_t sida_offset;
	void *old;
	uint8_t old_value[16];
	bool *cmpxchg_success;
	uint8_t ar;
	uint8_t key;
};

const uint8_t NO_KEY = 0xff;

static struct kvm_s390_mem_op ksmo_from_desc(struct mop_desc *desc)
{
	struct kvm_s390_mem_op ksmo = {
		.gaddr = (uintptr_t)desc->gaddr,
		.size = desc->size,
		.buf = ((uintptr_t)desc->buf),
		.reserved = "ignored_ignored_ignored_ignored"
	};

	switch (desc->target) {
	case LOGICAL:
		if (desc->mode == READ)
			ksmo.op = KVM_S390_MEMOP_LOGICAL_READ;
		if (desc->mode == WRITE)
			ksmo.op = KVM_S390_MEMOP_LOGICAL_WRITE;
		break;
	case SIDA:
		if (desc->mode == READ)
			ksmo.op = KVM_S390_MEMOP_SIDA_READ;
		if (desc->mode == WRITE)
			ksmo.op = KVM_S390_MEMOP_SIDA_WRITE;
		break;
	case ABSOLUTE:
		if (desc->mode == READ)
			ksmo.op = KVM_S390_MEMOP_ABSOLUTE_READ;
		if (desc->mode == WRITE)
			ksmo.op = KVM_S390_MEMOP_ABSOLUTE_WRITE;
		if (desc->mode == CMPXCHG) {
			ksmo.op = KVM_S390_MEMOP_ABSOLUTE_CMPXCHG;
			ksmo.old_addr = (uint64_t)desc->old;
			memcpy(desc->old_value, desc->old, desc->size);
		}
		break;
	case INVALID:
		ksmo.op = -1;
	}
	if (desc->f_check)
		ksmo.flags |= KVM_S390_MEMOP_F_CHECK_ONLY;
	if (desc->f_inject)
		ksmo.flags |= KVM_S390_MEMOP_F_INJECT_EXCEPTION;
	if (desc->_set_flags)
		ksmo.flags = desc->set_flags;
	if (desc->f_key && desc->key != NO_KEY) {
		ksmo.flags |= KVM_S390_MEMOP_F_SKEY_PROTECTION;
		ksmo.key = desc->key;
	}
	if (desc->_ar)
		ksmo.ar = desc->ar;
	else
		ksmo.ar = 0;
	if (desc->_sida_offset)
		ksmo.sida_offset = desc->sida_offset;

	return ksmo;
}

struct test_info {
	struct kvm_vm *vm;
	struct kvm_vcpu *vcpu;
};

#define PRINT_MEMOP false
static void print_memop(struct kvm_vcpu *vcpu, const struct kvm_s390_mem_op *ksmo)
{
	if (!PRINT_MEMOP)
		return;

	if (!vcpu)
		printf("vm memop(");
	else
		printf("vcpu memop(");
	switch (ksmo->op) {
	case KVM_S390_MEMOP_LOGICAL_READ:
		printf("LOGICAL, READ, ");
		break;
	case KVM_S390_MEMOP_LOGICAL_WRITE:
		printf("LOGICAL, WRITE, ");
		break;
	case KVM_S390_MEMOP_SIDA_READ:
		printf("SIDA, READ, ");
		break;
	case KVM_S390_MEMOP_SIDA_WRITE:
		printf("SIDA, WRITE, ");
		break;
	case KVM_S390_MEMOP_ABSOLUTE_READ:
		printf("ABSOLUTE, READ, ");
		break;
	case KVM_S390_MEMOP_ABSOLUTE_WRITE:
		printf("ABSOLUTE, WRITE, ");
		break;
	case KVM_S390_MEMOP_ABSOLUTE_CMPXCHG:
		printf("ABSOLUTE, CMPXCHG, ");
		break;
	}
	printf("gaddr=%llu, size=%u, buf=%llu, ar=%u, key=%u, old_addr=%llx",
	       ksmo->gaddr, ksmo->size, ksmo->buf, ksmo->ar, ksmo->key,
	       ksmo->old_addr);
	if (ksmo->flags & KVM_S390_MEMOP_F_CHECK_ONLY)
		printf(", CHECK_ONLY");
	if (ksmo->flags & KVM_S390_MEMOP_F_INJECT_EXCEPTION)
		printf(", INJECT_EXCEPTION");
	if (ksmo->flags & KVM_S390_MEMOP_F_SKEY_PROTECTION)
		printf(", SKEY_PROTECTION");
	puts(")");
}

static int err_memop_ioctl(struct test_info info, struct kvm_s390_mem_op *ksmo,
			   struct mop_desc *desc)
{
	struct kvm_vcpu *vcpu = info.vcpu;

	if (!vcpu)
		return __vm_ioctl(info.vm, KVM_S390_MEM_OP, ksmo);
	else
		return __vcpu_ioctl(vcpu, KVM_S390_MEM_OP, ksmo);
}

static void memop_ioctl(struct test_info info, struct kvm_s390_mem_op *ksmo,
			struct mop_desc *desc)
{
	int r;

	r = err_memop_ioctl(info, ksmo, desc);
	if (ksmo->op == KVM_S390_MEMOP_ABSOLUTE_CMPXCHG) {
		if (desc->cmpxchg_success) {
			int diff = memcmp(desc->old_value, desc->old, desc->size);
			*desc->cmpxchg_success = !diff;
		}
	}
	TEST_ASSERT(!r, __KVM_IOCTL_ERROR("KVM_S390_MEM_OP", r));
}

#define MEMOP(err, info_p, mop_target_p, access_mode_p, buf_p, size_p, ...)	\
({										\
	struct test_info __info = (info_p);					\
	struct mop_desc __desc = {						\
		.target = (mop_target_p),					\
		.mode = (access_mode_p),					\
		.buf = (buf_p),							\
		.size = (size_p),						\
		__VA_ARGS__							\
	};									\
	struct kvm_s390_mem_op __ksmo;						\
										\
	if (__desc._gaddr_v) {							\
		if (__desc.target == ABSOLUTE)					\
			__desc.gaddr = addr_gva2gpa(__info.vm, __desc.gaddr_v);	\
		else								\
			__desc.gaddr = __desc.gaddr_v;				\
	}									\
	__ksmo = ksmo_from_desc(&__desc);					\
	print_memop(__info.vcpu, &__ksmo);					\
	err##memop_ioctl(__info, &__ksmo, &__desc);				\
})

#define MOP(...) MEMOP(, __VA_ARGS__)
#define ERR_MOP(...) MEMOP(err_, __VA_ARGS__)

#define GADDR(a) .gaddr = ((uintptr_t)a)
#define GADDR_V(v) ._gaddr_v = 1, .gaddr_v = ((uintptr_t)v)
#define CHECK_ONLY .f_check = 1
#define SET_FLAGS(f) ._set_flags = 1, .set_flags = (f)
#define SIDA_OFFSET(o) ._sida_offset = 1, .sida_offset = (o)
#define AR(a) ._ar = 1, .ar = (a)
#define KEY(a) .f_key = 1, .key = (a)
#define INJECT .f_inject = 1
#define CMPXCHG_OLD(o) .old = (o)
#define CMPXCHG_SUCCESS(s) .cmpxchg_success = (s)

#define CHECK_N_DO(f, ...) ({ f(__VA_ARGS__, CHECK_ONLY); f(__VA_ARGS__); })

#define PAGE_SHIFT 12
#define PAGE_SIZE (1ULL << PAGE_SHIFT)
#define PAGE_MASK (~(PAGE_SIZE - 1))
#define CR0_FETCH_PROTECTION_OVERRIDE	(1UL << (63 - 38))
#define CR0_STORAGE_PROTECTION_OVERRIDE	(1UL << (63 - 39))

static uint8_t __aligned(PAGE_SIZE) mem1[65536];
static uint8_t __aligned(PAGE_SIZE) mem2[65536];

struct test_default {
	struct kvm_vm *kvm_vm;
	struct test_info vm;
	struct test_info vcpu;
	struct kvm_run *run;
	int size;
};

static struct test_default test_default_init(void *guest_code)
{
	struct kvm_vcpu *vcpu;
	struct test_default t;

	t.size = min((size_t)kvm_check_cap(KVM_CAP_S390_MEM_OP), sizeof(mem1));
	t.kvm_vm = vm_create_with_one_vcpu(&vcpu, guest_code);
	t.vm = (struct test_info) { t.kvm_vm, NULL };
	t.vcpu = (struct test_info) { t.kvm_vm, vcpu };
	t.run = vcpu->run;
	return t;
}

enum stage {
	/* Synced state set by host, e.g. DAT */
	STAGE_INITED,
	/* Guest did nothing */
	STAGE_IDLED,
	/* Guest set storage keys (specifics up to test case) */
	STAGE_SKEYS_SET,
	/* Guest copied memory (locations up to test case) */
	STAGE_COPIED,
	/* End of guest code reached */
	STAGE_DONE,
};

#define HOST_SYNC(info_p, stage)					\
({									\
	struct test_info __info = (info_p);				\
	struct kvm_vcpu *__vcpu = __info.vcpu;				\
	struct ucall uc;						\
	int __stage = (stage);						\
									\
	vcpu_run(__vcpu);						\
	get_ucall(__vcpu, &uc);						\
	if (uc.cmd == UCALL_ABORT) {					\
		REPORT_GUEST_ASSERT(uc);				\
	}								\
	TEST_ASSERT_EQ(uc.cmd, UCALL_SYNC);				\
	TEST_ASSERT_EQ(uc.args[1], __stage);				\
})									\

static void prepare_mem12(void)
{
	int i;

	for (i = 0; i < sizeof(mem1); i++)
		mem1[i] = rand();
	memset(mem2, 0xaa, sizeof(mem2));
}

#define ASSERT_MEM_EQ(p1, p2, size) \
	TEST_ASSERT(!memcmp(p1, p2, size), "Memory contents do not match!")

static void default_write_read(struct test_info copy_cpu, struct test_info mop_cpu,
			       enum mop_target mop_target, uint32_t size, uint8_t key)
{
	prepare_mem12();
	CHECK_N_DO(MOP, mop_cpu, mop_target, WRITE, mem1, size,
		   GADDR_V(mem1), KEY(key));
	HOST_SYNC(copy_cpu, STAGE_COPIED);
	CHECK_N_DO(MOP, mop_cpu, mop_target, READ, mem2, size,
		   GADDR_V(mem2), KEY(key));
	ASSERT_MEM_EQ(mem1, mem2, size);
}

static void default_read(struct test_info copy_cpu, struct test_info mop_cpu,
			 enum mop_target mop_target, uint32_t size, uint8_t key)
{
	prepare_mem12();
	CHECK_N_DO(MOP, mop_cpu, mop_target, WRITE, mem1, size, GADDR_V(mem1));
	HOST_SYNC(copy_cpu, STAGE_COPIED);
	CHECK_N_DO(MOP, mop_cpu, mop_target, READ, mem2, size,
		   GADDR_V(mem2), KEY(key));
	ASSERT_MEM_EQ(mem1, mem2, size);
}

static void default_cmpxchg(struct test_default *test, uint8_t key)
{
	for (int size = 1; size <= 16; size *= 2) {
		for (int offset = 0; offset < 16; offset += size) {
			uint8_t __aligned(16) new[16] = {};
			uint8_t __aligned(16) old[16];
			bool succ;

			prepare_mem12();
			default_write_read(test->vcpu, test->vcpu, LOGICAL, 16, NO_KEY);

			memcpy(&old, mem1, 16);
			MOP(test->vm, ABSOLUTE, CMPXCHG, new + offset,
			    size, GADDR_V(mem1 + offset),
			    CMPXCHG_OLD(old + offset),
			    CMPXCHG_SUCCESS(&succ), KEY(key));
			HOST_SYNC(test->vcpu, STAGE_COPIED);
			MOP(test->vm, ABSOLUTE, READ, mem2, 16, GADDR_V(mem2));
			TEST_ASSERT(succ, "exchange of values should succeed");
			memcpy(mem1 + offset, new + offset, size);
			ASSERT_MEM_EQ(mem1, mem2, 16);

			memcpy(&old, mem1, 16);
			new[offset]++;
			old[offset]++;
			MOP(test->vm, ABSOLUTE, CMPXCHG, new + offset,
			    size, GADDR_V(mem1 + offset),
			    CMPXCHG_OLD(old + offset),
			    CMPXCHG_SUCCESS(&succ), KEY(key));
			HOST_SYNC(test->vcpu, STAGE_COPIED);
			MOP(test->vm, ABSOLUTE, READ, mem2, 16, GADDR_V(mem2));
			TEST_ASSERT(!succ, "exchange of values should not succeed");
			ASSERT_MEM_EQ(mem1, mem2, 16);
			ASSERT_MEM_EQ(&old, mem1, 16);
		}
	}
}

static void guest_copy(void)
{
	GUEST_SYNC(STAGE_INITED);
	memcpy(&mem2, &mem1, sizeof(mem2));
	GUEST_SYNC(STAGE_COPIED);
}

static void test_copy(void)
{
	struct test_default t = test_default_init(guest_copy);

	HOST_SYNC(t.vcpu, STAGE_INITED);

	default_write_read(t.vcpu, t.vcpu, LOGICAL, t.size, NO_KEY);

	kvm_vm_free(t.kvm_vm);
}

static void test_copy_access_register(void)
{
	struct test_default t = test_default_init(guest_copy);

	HOST_SYNC(t.vcpu, STAGE_INITED);

	prepare_mem12();
	t.run->psw_mask &= ~(3UL << (63 - 17));
	t.run->psw_mask |= 1UL << (63 - 17);  /* Enable AR mode */

	/*
	 * Primary address space gets used if an access register
	 * contains zero. The host makes use of AR[1] so is a good
	 * candidate to ensure the guest AR (of zero) is used.
	 */
	CHECK_N_DO(MOP, t.vcpu, LOGICAL, WRITE, mem1, t.size,
		   GADDR_V(mem1), AR(1));
	HOST_SYNC(t.vcpu, STAGE_COPIED);

	CHECK_N_DO(MOP, t.vcpu, LOGICAL, READ, mem2, t.size,
		   GADDR_V(mem2), AR(1));
	ASSERT_MEM_EQ(mem1, mem2, t.size);

	kvm_vm_free(t.kvm_vm);
}

static void set_storage_key_range(void *addr, size_t len, uint8_t key)
{
	uintptr_t _addr, abs, i;
	int not_mapped = 0;

	_addr = (uintptr_t)addr;
	for (i = _addr & PAGE_MASK; i < _addr + len; i += PAGE_SIZE) {
		abs = i;
		asm volatile (
			       "lra	%[abs], 0(0,%[abs])\n"
			"	jz	0f\n"
			"	llill	%[not_mapped],1\n"
			"	j	1f\n"
			"0:	sske	%[key], %[abs]\n"
			"1:"
			: [abs] "+&a" (abs), [not_mapped] "+r" (not_mapped)
			: [key] "r" (key)
			: "cc"
		);
		GUEST_ASSERT_EQ(not_mapped, 0);
	}
}

static void guest_copy_key(void)
{
	set_storage_key_range(mem1, sizeof(mem1), 0x90);
	set_storage_key_range(mem2, sizeof(mem2), 0x90);
	GUEST_SYNC(STAGE_SKEYS_SET);

	for (;;) {
		memcpy(&mem2, &mem1, sizeof(mem2));
		GUEST_SYNC(STAGE_COPIED);
	}
}

static void test_copy_key(void)
{
	struct test_default t = test_default_init(guest_copy_key);

	HOST_SYNC(t.vcpu, STAGE_SKEYS_SET);

	/* vm, no key */
	default_write_read(t.vcpu, t.vm, ABSOLUTE, t.size, NO_KEY);

	/* vm/vcpu, machting key or key 0 */
	default_write_read(t.vcpu, t.vcpu, LOGICAL, t.size, 0);
	default_write_read(t.vcpu, t.vcpu, LOGICAL, t.size, 9);
	default_write_read(t.vcpu, t.vm, ABSOLUTE, t.size, 0);
	default_write_read(t.vcpu, t.vm, ABSOLUTE, t.size, 9);
	/*
	 * There used to be different code paths for key handling depending on
	 * if the region crossed a page boundary.
	 * There currently are not, but the more tests the merrier.
	 */
	default_write_read(t.vcpu, t.vcpu, LOGICAL, 1, 0);
	default_write_read(t.vcpu, t.vcpu, LOGICAL, 1, 9);
	default_write_read(t.vcpu, t.vm, ABSOLUTE, 1, 0);
	default_write_read(t.vcpu, t.vm, ABSOLUTE, 1, 9);

	/* vm/vcpu, mismatching keys on read, but no fetch protection */
	default_read(t.vcpu, t.vcpu, LOGICAL, t.size, 2);
	default_read(t.vcpu, t.vm, ABSOLUTE, t.size, 2);

	kvm_vm_free(t.kvm_vm);
}

static void test_cmpxchg_key(void)
{
	struct test_default t = test_default_init(guest_copy_key);

	HOST_SYNC(t.vcpu, STAGE_SKEYS_SET);

	default_cmpxchg(&t, NO_KEY);
	default_cmpxchg(&t, 0);
	default_cmpxchg(&t, 9);

	kvm_vm_free(t.kvm_vm);
}

static __uint128_t cut_to_size(int size, __uint128_t val)
{
	switch (size) {
	case 1:
		return (uint8_t)val;
	case 2:
		return (uint16_t)val;
	case 4:
		return (uint32_t)val;
	case 8:
		return (uint64_t)val;
	case 16:
		return val;
	}
	GUEST_FAIL("Invalid size = %u", size);
	return 0;
}

static bool popcount_eq(__uint128_t a, __uint128_t b)
{
	unsigned int count_a, count_b;

	count_a = __builtin_popcountl((uint64_t)(a >> 64)) +
		  __builtin_popcountl((uint64_t)a);
	count_b = __builtin_popcountl((uint64_t)(b >> 64)) +
		  __builtin_popcountl((uint64_t)b);
	return count_a == count_b;
}

static __uint128_t rotate(int size, __uint128_t val, int amount)
{
	unsigned int bits = size * 8;

	amount = (amount + bits) % bits;
	val = cut_to_size(size, val);
	if (!amount)
		return val;
	return (val << (bits - amount)) | (val >> amount);
}

const unsigned int max_block = 16;

static void choose_block(bool guest, int i, int *size, int *offset)
{
	unsigned int rand;

	rand = i;
	if (guest) {
		rand = rand * 19 + 11;
		*size = 1 << ((rand % 3) + 2);
		rand = rand * 19 + 11;
		*offset = (rand % max_block) & ~(*size - 1);
	} else {
		rand = rand * 17 + 5;
		*size = 1 << (rand % 5);
		rand = rand * 17 + 5;
		*offset = (rand % max_block) & ~(*size - 1);
	}
}

static __uint128_t permutate_bits(bool guest, int i, int size, __uint128_t old)
{
	unsigned int rand;
	int amount;
	bool swap;

	rand = i;
	rand = rand * 3 + 1;
	if (guest)
		rand = rand * 3 + 1;
	swap = rand % 2 == 0;
	if (swap) {
		int i, j;
		__uint128_t new;
		uint8_t byte0, byte1;

		rand = rand * 3 + 1;
		i = rand % size;
		rand = rand * 3 + 1;
		j = rand % size;
		if (i == j)
			return old;
		new = rotate(16, old, i * 8);
		byte0 = new & 0xff;
		new &= ~0xff;
		new = rotate(16, new, -i * 8);
		new = rotate(16, new, j * 8);
		byte1 = new & 0xff;
		new = (new & ~0xff) | byte0;
		new = rotate(16, new, -j * 8);
		new = rotate(16, new, i * 8);
		new = new | byte1;
		new = rotate(16, new, -i * 8);
		return new;
	}
	rand = rand * 3 + 1;
	amount = rand % (size * 8);
	return rotate(size, old, amount);
}

static bool _cmpxchg(int size, void *target, __uint128_t *old_addr, __uint128_t new)
{
	bool ret;

	switch (size) {
	case 4: {
			uint32_t old = *old_addr;

			asm volatile ("cs %[old],%[new],%[address]"
			    : [old] "+d" (old),
			      [address] "+Q" (*(uint32_t *)(target))
			    : [new] "d" ((uint32_t)new)
			    : "cc"
			);
			ret = old == (uint32_t)*old_addr;
			*old_addr = old;
			return ret;
		}
	case 8: {
			uint64_t old = *old_addr;

			asm volatile ("csg %[old],%[new],%[address]"
			    : [old] "+d" (old),
			      [address] "+Q" (*(uint64_t *)(target))
			    : [new] "d" ((uint64_t)new)
			    : "cc"
			);
			ret = old == (uint64_t)*old_addr;
			*old_addr = old;
			return ret;
		}
	case 16: {
			__uint128_t old = *old_addr;

			asm volatile ("cdsg %[old],%[new],%[address]"
			    : [old] "+d" (old),
			      [address] "+Q" (*(__uint128_t *)(target))
			    : [new] "d" (new)
			    : "cc"
			);
			ret = old == *old_addr;
			*old_addr = old;
			return ret;
		}
	}
	GUEST_FAIL("Invalid size = %u", size);
	return 0;
}

const unsigned int cmpxchg_iter_outer = 100, cmpxchg_iter_inner = 10000;

static void guest_cmpxchg_key(void)
{
	int size, offset;
	__uint128_t old, new;

	set_storage_key_range(mem1, max_block, 0x10);
	set_storage_key_range(mem2, max_block, 0x10);
	GUEST_SYNC(STAGE_SKEYS_SET);

	for (int i = 0; i < cmpxchg_iter_outer; i++) {
		do {
			old = 1;
		} while (!_cmpxchg(16, mem1, &old, 0));
		for (int j = 0; j < cmpxchg_iter_inner; j++) {
			choose_block(true, i + j, &size, &offset);
			do {
				new = permutate_bits(true, i + j, size, old);
			} while (!_cmpxchg(size, mem2 + offset, &old, new));
		}
	}

	GUEST_SYNC(STAGE_DONE);
}

static void *run_guest(void *data)
{
	struct test_info *info = data;

	HOST_SYNC(*info, STAGE_DONE);
	return NULL;
}

static char *quad_to_char(__uint128_t *quad, int size)
{
	return ((char *)quad) + (sizeof(*quad) - size);
}

static void test_cmpxchg_key_concurrent(void)
{
	struct test_default t = test_default_init(guest_cmpxchg_key);
	int size, offset;
	__uint128_t old, new;
	bool success;
	pthread_t thread;

	HOST_SYNC(t.vcpu, STAGE_SKEYS_SET);
	prepare_mem12();
	MOP(t.vcpu, LOGICAL, WRITE, mem1, max_block, GADDR_V(mem2));
	pthread_create(&thread, NULL, run_guest, &t.vcpu);

	for (int i = 0; i < cmpxchg_iter_outer; i++) {
		do {
			old = 0;
			new = 1;
			MOP(t.vm, ABSOLUTE, CMPXCHG, &new,
			    sizeof(new), GADDR_V(mem1),
			    CMPXCHG_OLD(&old),
			    CMPXCHG_SUCCESS(&success), KEY(1));
		} while (!success);
		for (int j = 0; j < cmpxchg_iter_inner; j++) {
			choose_block(false, i + j, &size, &offset);
			do {
				new = permutate_bits(false, i + j, size, old);
				MOP(t.vm, ABSOLUTE, CMPXCHG, quad_to_char(&new, size),
				    size, GADDR_V(mem2 + offset),
				    CMPXCHG_OLD(quad_to_char(&old, size)),
				    CMPXCHG_SUCCESS(&success), KEY(1));
			} while (!success);
		}
	}

	pthread_join(thread, NULL);

	MOP(t.vcpu, LOGICAL, READ, mem2, max_block, GADDR_V(mem2));
	TEST_ASSERT(popcount_eq(*(__uint128_t *)mem1, *(__uint128_t *)mem2),
		    "Must retain number of set bits");

	kvm_vm_free(t.kvm_vm);
}

static void guest_copy_key_fetch_prot(void)
{
	/*
	 * For some reason combining the first sync with override enablement
	 * results in an exception when calling HOST_SYNC.
	 */
	GUEST_SYNC(STAGE_INITED);
	/* Storage protection override applies to both store and fetch. */
	set_storage_key_range(mem1, sizeof(mem1), 0x98);
	set_storage_key_range(mem2, sizeof(mem2), 0x98);
	GUEST_SYNC(STAGE_SKEYS_SET);

	for (;;) {
		memcpy(&mem2, &mem1, sizeof(mem2));
		GUEST_SYNC(STAGE_COPIED);
	}
}

static void test_copy_key_storage_prot_override(void)
{
	struct test_default t = test_default_init(guest_copy_key_fetch_prot);

	HOST_SYNC(t.vcpu, STAGE_INITED);
	t.run->s.regs.crs[0] |= CR0_STORAGE_PROTECTION_OVERRIDE;
	t.run->kvm_dirty_regs = KVM_SYNC_CRS;
	HOST_SYNC(t.vcpu, STAGE_SKEYS_SET);

	/* vcpu, mismatching keys, storage protection override in effect */
	default_write_read(t.vcpu, t.vcpu, LOGICAL, t.size, 2);

	kvm_vm_free(t.kvm_vm);
}

static void test_copy_key_fetch_prot(void)
{
	struct test_default t = test_default_init(guest_copy_key_fetch_prot);

	HOST_SYNC(t.vcpu, STAGE_INITED);
	HOST_SYNC(t.vcpu, STAGE_SKEYS_SET);

	/* vm/vcpu, matching key, fetch protection in effect */
	default_read(t.vcpu, t.vcpu, LOGICAL, t.size, 9);
	default_read(t.vcpu, t.vm, ABSOLUTE, t.size, 9);

	kvm_vm_free(t.kvm_vm);
}

#define ERR_PROT_MOP(...)							\
({										\
	int rv;									\
										\
	rv = ERR_MOP(__VA_ARGS__);						\
	TEST_ASSERT(rv == 4, "Should result in protection exception");		\
})

static void guest_error_key(void)
{
	GUEST_SYNC(STAGE_INITED);
	set_storage_key_range(mem1, PAGE_SIZE, 0x18);
	set_storage_key_range(mem1 + PAGE_SIZE, sizeof(mem1) - PAGE_SIZE, 0x98);
	GUEST_SYNC(STAGE_SKEYS_SET);
	GUEST_SYNC(STAGE_IDLED);
}

static void test_errors_key(void)
{
	struct test_default t = test_default_init(guest_error_key);

	HOST_SYNC(t.vcpu, STAGE_INITED);
	HOST_SYNC(t.vcpu, STAGE_SKEYS_SET);

	/* vm/vcpu, mismatching keys, fetch protection in effect */
	CHECK_N_DO(ERR_PROT_MOP, t.vcpu, LOGICAL, WRITE, mem1, t.size, GADDR_V(mem1), KEY(2));
	CHECK_N_DO(ERR_PROT_MOP, t.vcpu, LOGICAL, READ, mem2, t.size, GADDR_V(mem1), KEY(2));
	CHECK_N_DO(ERR_PROT_MOP, t.vm, ABSOLUTE, WRITE, mem1, t.size, GADDR_V(mem1), KEY(2));
	CHECK_N_DO(ERR_PROT_MOP, t.vm, ABSOLUTE, READ, mem2, t.size, GADDR_V(mem1), KEY(2));

	kvm_vm_free(t.kvm_vm);
}

static void test_errors_cmpxchg_key(void)
{
	struct test_default t = test_default_init(guest_copy_key_fetch_prot);
	int i;

	HOST_SYNC(t.vcpu, STAGE_INITED);
	HOST_SYNC(t.vcpu, STAGE_SKEYS_SET);

	for (i = 1; i <= 16; i *= 2) {
		__uint128_t old = 0;

		ERR_PROT_MOP(t.vm, ABSOLUTE, CMPXCHG, mem2, i, GADDR_V(mem2),
			     CMPXCHG_OLD(&old), KEY(2));
	}

	kvm_vm_free(t.kvm_vm);
}

static void test_termination(void)
{
	struct test_default t = test_default_init(guest_error_key);
	uint64_t prefix;
	uint64_t teid;
	uint64_t teid_mask = BIT(63 - 56) | BIT(63 - 60) | BIT(63 - 61);
	uint64_t psw[2];

	HOST_SYNC(t.vcpu, STAGE_INITED);
	HOST_SYNC(t.vcpu, STAGE_SKEYS_SET);

	/* vcpu, mismatching keys after first page */
	ERR_PROT_MOP(t.vcpu, LOGICAL, WRITE, mem1, t.size, GADDR_V(mem1), KEY(1), INJECT);
	/*
	 * The memop injected a program exception and the test needs to check the
	 * Translation-Exception Identification (TEID). It is necessary to run
	 * the guest in order to be able to read the TEID from guest memory.
	 * Set the guest program new PSW, so the guest state is not clobbered.
	 */
	prefix = t.run->s.regs.prefix;
	psw[0] = t.run->psw_mask;
	psw[1] = t.run->psw_addr;
	MOP(t.vm, ABSOLUTE, WRITE, psw, sizeof(psw), GADDR(prefix + 464));
	HOST_SYNC(t.vcpu, STAGE_IDLED);
	MOP(t.vm, ABSOLUTE, READ, &teid, sizeof(teid), GADDR(prefix + 168));
	/* Bits 56, 60, 61 form a code, 0 being the only one allowing for termination */
	TEST_ASSERT_EQ(teid & teid_mask, 0);

	kvm_vm_free(t.kvm_vm);
}

static void test_errors_key_storage_prot_override(void)
{
	struct test_default t = test_default_init(guest_copy_key_fetch_prot);

	HOST_SYNC(t.vcpu, STAGE_INITED);
	t.run->s.regs.crs[0] |= CR0_STORAGE_PROTECTION_OVERRIDE;
	t.run->kvm_dirty_regs = KVM_SYNC_CRS;
	HOST_SYNC(t.vcpu, STAGE_SKEYS_SET);

	/* vm, mismatching keys, storage protection override not applicable to vm */
	CHECK_N_DO(ERR_PROT_MOP, t.vm, ABSOLUTE, WRITE, mem1, t.size, GADDR_V(mem1), KEY(2));
	CHECK_N_DO(ERR_PROT_MOP, t.vm, ABSOLUTE, READ, mem2, t.size, GADDR_V(mem2), KEY(2));

	kvm_vm_free(t.kvm_vm);
}

const uint64_t last_page_addr = -PAGE_SIZE;

static void guest_copy_key_fetch_prot_override(void)
{
	int i;
	char *page_0 = 0;

	GUEST_SYNC(STAGE_INITED);
	set_storage_key_range(0, PAGE_SIZE, 0x18);
	set_storage_key_range((void *)last_page_addr, PAGE_SIZE, 0x0);
	asm volatile ("sske %[key],%[addr]\n" :: [addr] "r"(0L), [key] "r"(0x18) : "cc");
	GUEST_SYNC(STAGE_SKEYS_SET);

	for (;;) {
		for (i = 0; i < PAGE_SIZE; i++)
			page_0[i] = mem1[i];
		GUEST_SYNC(STAGE_COPIED);
	}
}

static void test_copy_key_fetch_prot_override(void)
{
	struct test_default t = test_default_init(guest_copy_key_fetch_prot_override);
	vm_vaddr_t guest_0_page, guest_last_page;

	guest_0_page = vm_vaddr_alloc(t.kvm_vm, PAGE_SIZE, 0);
	guest_last_page = vm_vaddr_alloc(t.kvm_vm, PAGE_SIZE, last_page_addr);
	if (guest_0_page != 0 || guest_last_page != last_page_addr) {
		print_skip("did not allocate guest pages at required positions");
		goto out;
	}

	HOST_SYNC(t.vcpu, STAGE_INITED);
	t.run->s.regs.crs[0] |= CR0_FETCH_PROTECTION_OVERRIDE;
	t.run->kvm_dirty_regs = KVM_SYNC_CRS;
	HOST_SYNC(t.vcpu, STAGE_SKEYS_SET);

	/* vcpu, mismatching keys on fetch, fetch protection override applies */
	prepare_mem12();
	MOP(t.vcpu, LOGICAL, WRITE, mem1, PAGE_SIZE, GADDR_V(mem1));
	HOST_SYNC(t.vcpu, STAGE_COPIED);
	CHECK_N_DO(MOP, t.vcpu, LOGICAL, READ, mem2, 2048, GADDR_V(guest_0_page), KEY(2));
	ASSERT_MEM_EQ(mem1, mem2, 2048);

	/*
	 * vcpu, mismatching keys on fetch, fetch protection override applies,
	 * wraparound
	 */
	prepare_mem12();
	MOP(t.vcpu, LOGICAL, WRITE, mem1, 2 * PAGE_SIZE, GADDR_V(guest_last_page));
	HOST_SYNC(t.vcpu, STAGE_COPIED);
	CHECK_N_DO(MOP, t.vcpu, LOGICAL, READ, mem2, PAGE_SIZE + 2048,
		   GADDR_V(guest_last_page), KEY(2));
	ASSERT_MEM_EQ(mem1, mem2, 2048);

out:
	kvm_vm_free(t.kvm_vm);
}

static void test_errors_key_fetch_prot_override_not_enabled(void)
{
	struct test_default t = test_default_init(guest_copy_key_fetch_prot_override);
	vm_vaddr_t guest_0_page, guest_last_page;

	guest_0_page = vm_vaddr_alloc(t.kvm_vm, PAGE_SIZE, 0);
	guest_last_page = vm_vaddr_alloc(t.kvm_vm, PAGE_SIZE, last_page_addr);
	if (guest_0_page != 0 || guest_last_page != last_page_addr) {
		print_skip("did not allocate guest pages at required positions");
		goto out;
	}
	HOST_SYNC(t.vcpu, STAGE_INITED);
	HOST_SYNC(t.vcpu, STAGE_SKEYS_SET);

	/* vcpu, mismatching keys on fetch, fetch protection override not enabled */
	CHECK_N_DO(ERR_PROT_MOP, t.vcpu, LOGICAL, READ, mem2, 2048, GADDR_V(0), KEY(2));

out:
	kvm_vm_free(t.kvm_vm);
}

static void test_errors_key_fetch_prot_override_enabled(void)
{
	struct test_default t = test_default_init(guest_copy_key_fetch_prot_override);
	vm_vaddr_t guest_0_page, guest_last_page;

	guest_0_page = vm_vaddr_alloc(t.kvm_vm, PAGE_SIZE, 0);
	guest_last_page = vm_vaddr_alloc(t.kvm_vm, PAGE_SIZE, last_page_addr);
	if (guest_0_page != 0 || guest_last_page != last_page_addr) {
		print_skip("did not allocate guest pages at required positions");
		goto out;
	}
	HOST_SYNC(t.vcpu, STAGE_INITED);
	t.run->s.regs.crs[0] |= CR0_FETCH_PROTECTION_OVERRIDE;
	t.run->kvm_dirty_regs = KVM_SYNC_CRS;
	HOST_SYNC(t.vcpu, STAGE_SKEYS_SET);

	/*
	 * vcpu, mismatching keys on fetch,
	 * fetch protection override does not apply because memory range exceeded
	 */
	CHECK_N_DO(ERR_PROT_MOP, t.vcpu, LOGICAL, READ, mem2, 2048 + 1, GADDR_V(0), KEY(2));
	CHECK_N_DO(ERR_PROT_MOP, t.vcpu, LOGICAL, READ, mem2, PAGE_SIZE + 2048 + 1,
		   GADDR_V(guest_last_page), KEY(2));
	/* vm, fetch protected override does not apply */
	CHECK_N_DO(ERR_PROT_MOP, t.vm, ABSOLUTE, READ, mem2, 2048, GADDR(0), KEY(2));
	CHECK_N_DO(ERR_PROT_MOP, t.vm, ABSOLUTE, READ, mem2, 2048, GADDR_V(guest_0_page), KEY(2));

out:
	kvm_vm_free(t.kvm_vm);
}

static void guest_idle(void)
{
	GUEST_SYNC(STAGE_INITED); /* for consistency's sake */
	for (;;)
		GUEST_SYNC(STAGE_IDLED);
}

static void _test_errors_common(struct test_info info, enum mop_target target, int size)
{
	int rv;

	/* Bad size: */
	rv = ERR_MOP(info, target, WRITE, mem1, -1, GADDR_V(mem1));
	TEST_ASSERT(rv == -1 && errno == E2BIG, "ioctl allows insane sizes");

	/* Zero size: */
	rv = ERR_MOP(info, target, WRITE, mem1, 0, GADDR_V(mem1));
	TEST_ASSERT(rv == -1 && (errno == EINVAL || errno == ENOMEM),
		    "ioctl allows 0 as size");

	/* Bad flags: */
	rv = ERR_MOP(info, target, WRITE, mem1, size, GADDR_V(mem1), SET_FLAGS(-1));
	TEST_ASSERT(rv == -1 && errno == EINVAL, "ioctl allows all flags");

	/* Bad guest address: */
	rv = ERR_MOP(info, target, WRITE, mem1, size, GADDR((void *)~0xfffUL), CHECK_ONLY);
	TEST_ASSERT(rv > 0, "ioctl does not report bad guest memory address with CHECK_ONLY");
	rv = ERR_MOP(info, target, WRITE, mem1, size, GADDR((void *)~0xfffUL));
	TEST_ASSERT(rv > 0, "ioctl does not report bad guest memory address on write");

	/* Bad host address: */
	rv = ERR_MOP(info, target, WRITE, 0, size, GADDR_V(mem1));
	TEST_ASSERT(rv == -1 && errno == EFAULT,
		    "ioctl does not report bad host memory address");

	/* Bad key: */
	rv = ERR_MOP(info, target, WRITE, mem1, size, GADDR_V(mem1), KEY(17));
	TEST_ASSERT(rv == -1 && errno == EINVAL, "ioctl allows invalid key");
}

static void test_errors(void)
{
	struct test_default t = test_default_init(guest_idle);
	int rv;

	HOST_SYNC(t.vcpu, STAGE_INITED);

	_test_errors_common(t.vcpu, LOGICAL, t.size);
	_test_errors_common(t.vm, ABSOLUTE, t.size);

	/* Bad operation: */
	rv = ERR_MOP(t.vcpu, INVALID, WRITE, mem1, t.size, GADDR_V(mem1));
	TEST_ASSERT(rv == -1 && errno == EINVAL, "ioctl allows bad operations");
	/* virtual addresses are not translated when passing INVALID */
	rv = ERR_MOP(t.vm, INVALID, WRITE, mem1, PAGE_SIZE, GADDR(0));
	TEST_ASSERT(rv == -1 && errno == EINVAL, "ioctl allows bad operations");

	/* Bad access register: */
	t.run->psw_mask &= ~(3UL << (63 - 17));
	t.run->psw_mask |= 1UL << (63 - 17);  /* Enable AR mode */
	HOST_SYNC(t.vcpu, STAGE_IDLED); /* To sync new state to SIE block */
	rv = ERR_MOP(t.vcpu, LOGICAL, WRITE, mem1, t.size, GADDR_V(mem1), AR(17));
	TEST_ASSERT(rv == -1 && errno == EINVAL, "ioctl allows ARs > 15");
	t.run->psw_mask &= ~(3UL << (63 - 17));   /* Disable AR mode */
	HOST_SYNC(t.vcpu, STAGE_IDLED); /* Run to sync new state */

	/* Check that the SIDA calls are rejected for non-protected guests */
	rv = ERR_MOP(t.vcpu, SIDA, READ, mem1, 8, GADDR(0), SIDA_OFFSET(0x1c0));
	TEST_ASSERT(rv == -1 && errno == EINVAL,
		    "ioctl does not reject SIDA_READ in non-protected mode");
	rv = ERR_MOP(t.vcpu, SIDA, WRITE, mem1, 8, GADDR(0), SIDA_OFFSET(0x1c0));
	TEST_ASSERT(rv == -1 && errno == EINVAL,
		    "ioctl does not reject SIDA_WRITE in non-protected mode");

	kvm_vm_free(t.kvm_vm);
}

static void test_errors_cmpxchg(void)
{
	struct test_default t = test_default_init(guest_idle);
	__uint128_t old;
	int rv, i, power = 1;

	HOST_SYNC(t.vcpu, STAGE_INITED);

	for (i = 0; i < 32; i++) {
		if (i == power) {
			power *= 2;
			continue;
		}
		rv = ERR_MOP(t.vm, ABSOLUTE, CMPXCHG, mem1, i, GADDR_V(mem1),
			     CMPXCHG_OLD(&old));
		TEST_ASSERT(rv == -1 && errno == EINVAL,
			    "ioctl allows bad size for cmpxchg");
	}
	for (i = 1; i <= 16; i *= 2) {
		rv = ERR_MOP(t.vm, ABSOLUTE, CMPXCHG, mem1, i, GADDR((void *)~0xfffUL),
			     CMPXCHG_OLD(&old));
		TEST_ASSERT(rv > 0, "ioctl allows bad guest address for cmpxchg");
	}
	for (i = 2; i <= 16; i *= 2) {
		rv = ERR_MOP(t.vm, ABSOLUTE, CMPXCHG, mem1, i, GADDR_V(mem1 + 1),
			     CMPXCHG_OLD(&old));
		TEST_ASSERT(rv == -1 && errno == EINVAL,
			    "ioctl allows bad alignment for cmpxchg");
	}

	kvm_vm_free(t.kvm_vm);
}

int main(int argc, char *argv[])
{
	int extension_cap, idx;

	TEST_REQUIRE(kvm_has_cap(KVM_CAP_S390_MEM_OP));
	extension_cap = kvm_check_cap(KVM_CAP_S390_MEM_OP_EXTENSION);

	struct testdef {
		const char *name;
		void (*test)(void);
		bool requirements_met;
	} testlist[] = {
		{
			.name = "simple copy",
			.test = test_copy,
			.requirements_met = true,
		},
		{
			.name = "generic error checks",
			.test = test_errors,
			.requirements_met = true,
		},
		{
			.name = "copy with storage keys",
			.test = test_copy_key,
			.requirements_met = extension_cap > 0,
		},
		{
			.name = "cmpxchg with storage keys",
			.test = test_cmpxchg_key,
			.requirements_met = extension_cap & 0x2,
		},
		{
			.name = "concurrently cmpxchg with storage keys",
			.test = test_cmpxchg_key_concurrent,
			.requirements_met = extension_cap & 0x2,
		},
		{
			.name = "copy with key storage protection override",
			.test = test_copy_key_storage_prot_override,
			.requirements_met = extension_cap > 0,
		},
		{
			.name = "copy with key fetch protection",
			.test = test_copy_key_fetch_prot,
			.requirements_met = extension_cap > 0,
		},
		{
			.name = "copy with key fetch protection override",
			.test = test_copy_key_fetch_prot_override,
			.requirements_met = extension_cap > 0,
		},
		{
			.name = "copy with access register mode",
			.test = test_copy_access_register,
			.requirements_met = true,
		},
		{
			.name = "error checks with key",
			.test = test_errors_key,
			.requirements_met = extension_cap > 0,
		},
		{
			.name = "error checks for cmpxchg with key",
			.test = test_errors_cmpxchg_key,
			.requirements_met = extension_cap & 0x2,
		},
		{
			.name = "error checks for cmpxchg",
			.test = test_errors_cmpxchg,
			.requirements_met = extension_cap & 0x2,
		},
		{
			.name = "termination",
			.test = test_termination,
			.requirements_met = extension_cap > 0,
		},
		{
			.name = "error checks with key storage protection override",
			.test = test_errors_key_storage_prot_override,
			.requirements_met = extension_cap > 0,
		},
		{
			.name = "error checks without key fetch prot override",
			.test = test_errors_key_fetch_prot_override_not_enabled,
			.requirements_met = extension_cap > 0,
		},
		{
			.name = "error checks with key fetch prot override",
			.test = test_errors_key_fetch_prot_override_enabled,
			.requirements_met = extension_cap > 0,
		},
	};

	ksft_print_header();
	ksft_set_plan(ARRAY_SIZE(testlist));

	for (idx = 0; idx < ARRAY_SIZE(testlist); idx++) {
		if (testlist[idx].requirements_met) {
			testlist[idx].test();
			ksft_test_result_pass("%s\n", testlist[idx].name);
		} else {
			ksft_test_result_skip("%s - requirements not met (kernel has extension cap %#x)\n",
					      testlist[idx].name, extension_cap);
		}
	}

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