/*
 * Copyright (C) 2021-2026 Colin Ian King
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 *
 */
#include "stress-ng.h"
#include "core-arch.h"
#include "core-builtin.h"
#include "core-capabilities.h"
#include "core-madvise.h"
#include "core-mmap.h"

#include <sys/ioctl.h>

#if defined(HAVE_LINUX_KVM_H)
#include <linux/kvm.h>
#endif

static const stress_help_t help[] = {
	{ NULL,	"kvm N",	"start N workers exercising /dev/kvm" },
	{ NULL, "kvm-ops N",	"stop after N kvm create/run/destroy operations" },
	{ NULL,	NULL,		NULL }
};

#if defined(STRESS_ARCH_X86) &&	\
    !defined(__i386__) &&	\
    !defined(__i386)
#define STRESS_KVM_X86
#define PHYS_ADDR	(0x00000000)
#endif

#if defined(STRESS_ARCH_ARM) &&	\
    defined(__ARM_ARCH_ISA_A64)
#define STRESS_KVM_ARM
#define PHYS_ADDR	(0x04000000)
#endif

#if defined(STRESS_ARCH_RISCV)
#define STRESS_KVM_RISCV
#define PHYS_ADDR	(0x04000000)
#endif

#if defined(__linux__)	&&			\
    defined(HAVE_LINUX_KVM_H) && 		\
    defined(KVM_CREATE_VM) &&			\
    defined(KVM_SET_USER_MEMORY_REGION) &&	\
    defined(KVM_CREATE_VCPU) && 		\
    defined(KVM_GET_SREGS) &&			\
    defined(KVM_SET_SREGS) &&			\
    defined(KVM_SET_REGS) &&			\
    defined(KVM_GET_VCPU_MMAP_SIZE) &&		\
    defined(KVM_RUN) &&				\
    defined(KVM_EXIT_IO) &&			\
    defined(KVM_EXIT_SHUTDOWN) &&		\
    (defined(STRESS_KVM_X86) ||			\
     defined(STRESS_KVM_ARM) ||			\
     defined(STRESS_KVM_RISCV))

static int stress_kvm_open(const char *name, const bool report)
{
	int kvm_fd;

	if ((kvm_fd = open("/dev/kvm", O_RDWR)) < 0) {
		switch (errno) {
		case ENOENT:
			if (report)
				pr_inf_skip("%s: '/dev/kvm' not available, skipping stressor\n",
					name);
			break;
		case EPERM:
		case EACCES:
			if (report && !stress_capabilities_check(SHIM_CAP_SYS_ADMIN))
				pr_inf_skip("%s stressor will be skipped, "
					"need to be running with CAP_SYS_ADMIN "
					"rights for this stressor\n", name);
			break;
		default:
			if (report)
				pr_fail("%s: open '/dev/kvm' failed, errno=%d (%s), skipping stressor\n",
					name, errno, strerror(errno));
			break;
		}
		return -1;
	}
	return kvm_fd;
}

static int stress_kvm_supported(const char *name)
{
	int kvm_fd;

	kvm_fd = stress_kvm_open(name, true);
	if (kvm_fd < 0)
		return -1;

	(void)close(kvm_fd);
	return 0;
}

#if defined(STRESS_KVM_X86)
/*
 *  Minimal x86 kernel, read/increment/write port $80 loop
 */
static const uint8_t kvm_kernel[] = {
	0x31, 0xc0,  /* xor    %eax,%eax */
	0x0f, 0xa2,  /* cpuid */
	0x0f, 0x31,  /* rdtsc */
	0x0f, 0x08,  /* inv */
	0xe5, 0x80,  /* in     $0x80,%eax */
	0x40,        /* inc    %eax       */
	0xe7, 0x80,  /* out    %eax,$0x80 */
	0xeb, 0xf1,  /* jmp    0 <_start> */
};
#endif

#if defined(STRESS_KVM_ARM)
static const uint8_t kvm_kernel[] = {
	0x01, 0x00, 0xb0, 0xd2,	/* mov     x1, #0x80000000 */
	0x00, 0x00, 0x80, 0x12, /* mov     w0, #0xffffffff */
	0x20, 0x00, 0x00, 0x39, /* strb    w0, [x1] */
	0x20, 0x00, 0x00, 0x39, /* strb    w0, [x1] */
	0xfe, 0xff, 0xff, 0x17, /* b       768 <_start+0x8> */
};
#endif

#if defined(STRESS_KVM_RISCV)
static const uint8_t kvm_kernel[] = {
	0x85, 0x47,		/* li   a5, 1 */
	0xfe, 0x07,		/* slli a5,a5,0x1f */
	0x7d, 0x57,		/* li   a5, 0xff */
	0x23, 0x80, 0xe7, 0x00,	/* sb   a4,0(a5) */
	0x23, 0x80, 0xe7, 0x00, /* sb   a4,0(a5) */
	0xe5, 0xbf,		/* j    66e <_start+0x6) */
};
#endif

/*
 *  stress_kvm
 *	stress /dev/kvm
 */
static int stress_kvm(stress_args_t *args)
{
	bool pr_version = false;

	stress_proc_state_set(args->name, STRESS_STATE_SYNC_WAIT);
	stress_sync_start_wait(args);
	stress_proc_state_set(args->name, STRESS_STATE_RUN);

	do {
#if defined(STRESS_KVM_X86)
		struct kvm_sregs sregs;
		uint8_t value = 0;
#endif
		struct kvm_regs regs;
		struct kvm_userspace_memory_region kvm_mem;
		struct kvm_run *run;
		int kvm_fd;
		int vm_fd;
		int vcpu_fd;
		int version;
		int ret;
		int i;
		void *vm_mem;
#if defined(STRESS_KVM_ARM) ||	\
    defined(STRESS_KVM_RISCV)
		void *mmio_mem;
		int mmio_exits = 0;
		struct kvm_one_reg reg;
		const uint64_t entry_addr = PHYS_ADDR;
#endif

#if defined(STRESS_KVM_ARM)
		struct kvm_vcpu_init vcpu;
		const uint64_t arm_pc_index = ((uint8_t *)&regs.regs.pc - (uint8_t *)&regs) / sizeof(uint32_t);
#endif
		const size_t vm_mem_size = args->page_size;
		ssize_t run_size;
		bool run_ok = false;

		kvm_fd = stress_kvm_open(args->name, stress_instance_zero(args));
		if (kvm_fd < 0)
			return EXIT_NOT_IMPLEMENTED;

#if defined(KVM_GET_API_VERSION)
		version = ioctl(kvm_fd, KVM_GET_API_VERSION, 0);
		if ((version != -1) && (!pr_version) && (stress_instance_zero(args))) {
			pr_dbg("%s: KVM kernel API version %d\n", args->name, version);
			pr_version = true;
		}
#endif

		vm_fd = ioctl(kvm_fd, KVM_CREATE_VM, 0);
		if (vm_fd < 0) {
			if (errno == EINTR)
				goto tidy_kvm_fd;
			if (errno == EBUSY) {
				pr_inf_skip("%s: KVM device busy, skipping stressor\n", args->name);
				(void)close(kvm_fd);
				return EXIT_NO_RESOURCE;
			}
			pr_fail("%s: ioctl KVM_CREATE_VM failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_kvm_fd;
		}
		vm_mem = stress_mmap_populate(NULL, vm_mem_size,
			PROT_READ | PROT_WRITE,
			MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE,
			-1, 0);
		if (vm_mem == MAP_FAILED)
			goto tidy_vm_fd;
		(void)stress_madvise_mergeable(vm_mem, vm_mem_size);

		(void)shim_memset(&kvm_mem, 0, sizeof(kvm_mem));
		kvm_mem.slot = 0;
		kvm_mem.guest_phys_addr = 0;
		kvm_mem.guest_phys_addr = PHYS_ADDR;
		kvm_mem.memory_size = vm_mem_size;
		kvm_mem.userspace_addr = (uintptr_t)vm_mem;
		kvm_mem.flags = 0;
		ret = ioctl(vm_fd, KVM_SET_USER_MEMORY_REGION, &kvm_mem);
		if (ret < 0) {
			if (errno == EINTR)
				goto tidy_vm_mmap;
			pr_fail("%s: ioctl KVM_SET_USER_MEMORY_REGION for text region failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_vm_mmap;
		}

#if defined(STRESS_KVM_ARM) ||	\
    defined(STRESS_KVM_RISCV)
		/*
		 *  ARM: 1 page r/o MMIO space at 0x80000000
		 */
		mmio_mem = stress_mmap_populate(NULL, 2 * args->page_size,
			PROT_READ,
			MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE,
			-1, 0);
		if (mmio_mem == MAP_FAILED)
			goto tidy_vm_fd;

		kvm_mem.slot = 1;
		kvm_mem.guest_phys_addr = 0x80000000;
		kvm_mem.memory_size = args->page_size * 2;
		kvm_mem.userspace_addr = (uintptr_t)mmio_mem;
		kvm_mem.flags = KVM_MEM_READONLY;
		ret = ioctl(vm_fd, KVM_SET_USER_MEMORY_REGION, &kvm_mem);
		if (ret < 0) {
			if (errno == EINTR)
				goto tidy_vm_mmap;
			pr_fail("%s: ioctl KVM_SET_USER_MEMORY_REGION for MMIO memory failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_mmio_mmap;
		}
#endif

		vcpu_fd = ioctl(vm_fd, KVM_CREATE_VCPU, 0);
		if (vcpu_fd < 0) {
			if (errno == EINTR)
				goto tidy_vm_mmap;
			pr_fail("%s: ioctl KVM_CREATE_VCPU failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_mmio_mmap;
		}

		(void)shim_memcpy(vm_mem, &kvm_kernel, sizeof(kvm_kernel));

#if defined(STRESS_KVM_X86)
		if (ioctl(vcpu_fd, KVM_GET_SREGS, &sregs) < 0) {
			if (errno == EINTR)
				goto tidy_vcpu_fd;
			pr_fail("%s: ioctl KVM_GET_SREGS failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_vcpu_fd;
		}

		sregs.cs.selector = 0;
		sregs.cs.base = 0;
		sregs.ds.selector = 0;
		sregs.ds.base = 0;
		sregs.es.selector = 0;
		sregs.es.base = 0;
		sregs.fs.selector = 0;
		sregs.fs.base = 0;
		sregs.gs.selector = 0;
		sregs.gs.base = 0;
		sregs.ss.selector = 0;
		sregs.ss.base = 0;

		if (ioctl(vcpu_fd, KVM_SET_SREGS, &sregs) < 0) {
			pr_fail("%s: ioctl KVM_SET_SREGS failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_vcpu_fd;
		}
#endif

#if defined(STRESS_KVM_X86)
		if (ioctl(vcpu_fd, KVM_GET_REGS, &regs) < 0) {
			pr_fail("%s: ioctl KVM_GET_REGS failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_vcpu_fd;
		}
#endif

		(void)shim_memset(&regs, 0, sizeof(regs));
#if defined(STRESS_KVM_X86)
		regs.rflags = 2;
		regs.rip = 0;
		if (ioctl(vcpu_fd, KVM_SET_REGS, &regs) < 0) {
			if (errno == EINTR)
				goto tidy_vcpu_fd;
			pr_fail("%s: ioctl KVM_SET_REGS failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_vcpu_fd;
		}
#endif
#if defined(STRESS_KVM_ARM)
		(void)shim_memset(&vcpu, 0, sizeof(vcpu));
		if (ioctl(vm_fd, KVM_ARM_PREFERRED_TARGET, &vcpu) < 0) {
			pr_fail("%s: ioctl KVM_ARM_PREFERRED_TARGET failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_vcpu_fd;
		}
		vcpu.features[0] |= 1 << KVM_ARM_VCPU_PSCI_0_2;
		if (ioctl(vcpu_fd, KVM_ARM_VCPU_INIT, &vcpu) < 0) {
			pr_fail("%s: ioctl KVM_ARM_VCPU_INIT failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_vcpu_fd;
		}
		reg.id = KVM_REG_ARM64 | KVM_REG_SIZE_U64 | KVM_REG_ARM_CORE | arm_pc_index;
		reg.addr = (uint64_t)&entry_addr;

		if (ioctl(vcpu_fd, KVM_SET_ONE_REG, &reg) < 0) {
			pr_fail("%s: ioctl KVM_SET_ONE_REG failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_vcpu_fd;
		}
#endif
#if defined(STRESS_KVM_RISCV)
		/* set program counter */
		reg.id = KVM_REG_RISCV | KVM_REG_SIZE_U64 | KVM_REG_RISCV_CORE | 0;
		reg.addr = (uint64_t)&entry_addr;

		if (ioctl(vcpu_fd, KVM_SET_ONE_REG, &reg) < 0) {
			pr_fail("%s: ioctl KVM_SET_ONE_REG failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_vcpu_fd;
		}
#endif

		run_size = (ssize_t)ioctl(kvm_fd, KVM_GET_VCPU_MMAP_SIZE, 0);
		if (run_size < 0) {
			if (errno == EINTR)
				goto tidy_vcpu_fd;
			pr_fail("%s: ioctl KVM_GET_VCPU_MMAP_SIZE failed, errno=%d (%s)\n",
				args->name, errno, strerror(errno));
			goto tidy_vcpu_fd;
		}

		run = (struct kvm_run *)stress_mmap_populate(NULL, (size_t)run_size,
			PROT_READ | PROT_WRITE, MAP_SHARED, vcpu_fd, 0);
		if (run == MAP_FAILED) {
			pr_fail("%s: mmap on vcpu_fd failed%s, errno=%d (%s)\n",
				args->name, stress_memory_free_get(),
				errno, strerror(errno));
			goto tidy_vcpu_fd;
		}

		for (i = 0; LIKELY((i < 1000) && stress_continue(args)); i++) {
#if defined(STRESS_KVM_X86)
			uint8_t *port;
#endif

			ret = ioctl(vcpu_fd, KVM_RUN, 0);
			if (ret < 0) {
				if (errno != EINTR) {
					pr_fail("%s: ioctl KVM_RUN failed, errno=%d (%s)\n",
						args->name, errno, strerror(errno));
				}
				goto tidy_run;
			}
			switch (run->exit_reason) {
#if defined(STRESS_KVM_X86)
			case KVM_EXIT_IO:
				port = (uint8_t *)run + run->io.data_offset;
				if (run->io.direction == 0) {
					/* Read */
					*port = value;
				} else {
					/* Write */
					value = *port;
				}
				if (*port == 0xff) {
					run_ok = true;
					goto tidy_run;
				}
				break;
#endif
#if defined(STRESS_KVM_ARM) ||	\
    defined(STRESS_KVM_RISCV)
			case KVM_EXIT_MMIO:
				mmio_exits++;
				if (mmio_exits > 255) {
					run_ok = true;
					goto tidy_run;
				}
				break;
#endif
			case KVM_EXIT_SHUTDOWN:
				goto tidy_run;
			default:
				break;
			}
#if defined(STRESS_KVM_X86) &&	\
    defined(KVM_GET_REGS)
			{
				struct kvm_regs kregs;

				VOID_RET(int, ioctl(vcpu_fd, KVM_GET_REGS, &kregs));
			}
#endif
#if defined(STRESS_KVM_X86) &&	\
    defined(KVM_GET_FPU)
			{
				struct kvm_fpu fpu;

				VOID_RET(int, ioctl(vcpu_fd, KVM_GET_FPU, &fpu));
			}
#endif
#if defined(KVM_GET_MP_STATE)
			{
				struct kvm_mp_state state;

				VOID_RET(int, ioctl(vcpu_fd, KVM_GET_MP_STATE, &state));
			}
#endif
#if defined(STRESS_KVM_X86) &&	\
    defined(KVM_GET_XSAVE)
			{
				struct kvm_xsave xsave;

				VOID_RET(int, ioctl(vcpu_fd, KVM_GET_XSAVE, &xsave));
			}
#endif
#if defined(STRESS_KVM_X86) &&	\
    defined(KVM_GET_TSC_KHZ)
			VOID_RET(int, ioctl(vcpu_fd, KVM_GET_TSC_KHZ, 0));
#endif
		}
tidy_run:
		(void)munmap((void *)run, (size_t)run_size);
tidy_vcpu_fd:
		(void)close(vcpu_fd);
tidy_mmio_mmap:
#if defined(STRESS_KVM_ARM) ||	\
    defined(STRESS_KVM_RISCV)
		kvm_mem.slot = 1;
		kvm_mem.guest_phys_addr = 0x80000000;
		kvm_mem.memory_size = 0;
		kvm_mem.userspace_addr = (uintptr_t)mmio_mem;
		kvm_mem.flags = 0;
		VOID_RET(int, ioctl(vm_fd, KVM_SET_USER_MEMORY_REGION, &kvm_mem));

		(void)munmap((void *)mmio_mem, args->page_size);
#endif
tidy_vm_mmap:
		kvm_mem.slot = 0;
		kvm_mem.guest_phys_addr = 0;
		kvm_mem.guest_phys_addr = PHYS_ADDR;
		kvm_mem.memory_size = vm_mem_size;
		kvm_mem.userspace_addr = (uintptr_t)vm_mem;
		kvm_mem.flags = 0;
		VOID_RET(int, ioctl(vm_fd, KVM_SET_USER_MEMORY_REGION, &kvm_mem));

		(void)munmap((void *)vm_mem, vm_mem_size);
tidy_vm_fd:
		(void)close(vm_fd);
tidy_kvm_fd:
		(void)close(kvm_fd);
		if (run_ok)
			stress_bogo_inc(args);
	} while (stress_continue(args));

	stress_proc_state_set(args->name, STRESS_STATE_DEINIT);

	return EXIT_SUCCESS;
}

static const stress_exercises_t exercises[] = {
	STRESS_EX_FEATURE("lock-contention"),
	STRESS_EX_FEATURE("virt"),

	STRESS_EX_SYSCALL("close"),
	STRESS_EX_SYSCALL("ioctl"),
	STRESS_EX_SYSCALL("mmap"),
	STRESS_EX_SYSCALL("open"),
	STRESS_EX_SYSCALL("munmap"),
	STRESS_EX_END,
};

const stressor_info_t stress_kvm_info = {
	.stressor = stress_kvm,
	.classifier = CLASS_DEV | CLASS_OS,
	.supported = stress_kvm_supported,
	.verify = VERIFY_ALWAYS,
	.help = help,
	.exercises = exercises,
};
#else
const stressor_info_t stress_kvm_info = {
	.stressor = stress_unimplemented,
	.classifier = CLASS_DEV | CLASS_OS,
	.verify = VERIFY_ALWAYS,
	.help = help,
	.unimplemented_reason = "built on non-x86-64 without linux/kvm.h"
};
#endif
