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

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

#define CPUID_MWAIT (1u << 3)

enum monitor_mwait_testcases {
	MWAIT_QUIRK_DISABLED = BIT(0),
	MISC_ENABLES_QUIRK_DISABLED = BIT(1),
	MWAIT_DISABLED = BIT(2),
	CPUID_DISABLED = BIT(3),
	TEST_MAX = CPUID_DISABLED * 2 - 1,
};

/*
 * If both MWAIT and its quirk are disabled, MONITOR/MWAIT should #UD, in all
 * other scenarios KVM should emulate them as nops.
 */
#define GUEST_ASSERT_MONITOR_MWAIT(insn, testcase, vector)		\
do {									\
	bool fault_wanted = ((testcase) & MWAIT_QUIRK_DISABLED) &&	\
			    ((testcase) & MWAIT_DISABLED);		\
									\
	if (fault_wanted)						\
		__GUEST_ASSERT((vector) == UD_VECTOR,			\
			       "Expected #UD on " insn " for testcase '0x%x', got %s", \
			       testcase, ex_str(vector));		\
	else								\
		__GUEST_ASSERT(!(vector),				\
			       "Expected success on " insn " for testcase '0x%x', got %s", \
			       testcase, ex_str(vector));		\
} while (0)

static void guest_monitor_wait(void *arg)
{
	int testcase = (int) (long) arg;
	u8 vector;

	u64 val = rdmsr(MSR_IA32_MISC_ENABLE) & ~MSR_IA32_MISC_ENABLE_MWAIT;
	if (!(testcase & MWAIT_DISABLED))
		val |= MSR_IA32_MISC_ENABLE_MWAIT;
	wrmsr(MSR_IA32_MISC_ENABLE, val);

	__GUEST_ASSERT(this_cpu_has(X86_FEATURE_MWAIT) == !(testcase & MWAIT_DISABLED),
		       "Expected CPUID.MWAIT %s\n",
		       (testcase & MWAIT_DISABLED) ? "cleared" : "set");

	/*
	 * Arbitrarily MONITOR this function, SVM performs fault checks before
	 * intercept checks, so the inputs for MONITOR and MWAIT must be valid.
	 */
	vector = kvm_asm_safe("monitor", "a"(guest_monitor_wait), "c"(0), "d"(0));
	GUEST_ASSERT_MONITOR_MWAIT("MONITOR", testcase, vector);

	vector = kvm_asm_safe("mwait", "a"(guest_monitor_wait), "c"(0), "d"(0));
	GUEST_ASSERT_MONITOR_MWAIT("MWAIT", testcase, vector);

	GUEST_DONE();
}

int main(int argc, char *argv[])
{
	uint64_t disabled_quirks;
	struct kvm_vcpu *vcpu;
	struct kvm_vm *vm;
	struct ucall uc;
	int testcase;
	char test[80];

	TEST_REQUIRE(this_cpu_has(X86_FEATURE_MWAIT));
	TEST_REQUIRE(kvm_has_cap(KVM_CAP_DISABLE_QUIRKS2));

	ksft_print_header();
	ksft_set_plan(12);
	for (testcase = 0; testcase <= TEST_MAX; testcase++) {
		vm = vm_create_with_one_vcpu(&vcpu, guest_monitor_wait);
		vcpu_args_set(vcpu, 1, (void *)(long)testcase);

		disabled_quirks = 0;
		if (testcase & MWAIT_QUIRK_DISABLED) {
			disabled_quirks |= KVM_X86_QUIRK_MWAIT_NEVER_UD_FAULTS;
			strcpy(test, "MWAIT can fault");
		} else {
			strcpy(test, "MWAIT never faults");
		}
		if (testcase & MISC_ENABLES_QUIRK_DISABLED) {
			disabled_quirks |= KVM_X86_QUIRK_MISC_ENABLE_NO_MWAIT;
			strcat(test, ", MISC_ENABLE updates CPUID");
		} else {
			strcat(test, ", no CPUID updates");
		}

		vm_enable_cap(vm, KVM_CAP_DISABLE_QUIRKS2, disabled_quirks);

		if (!(testcase & MISC_ENABLES_QUIRK_DISABLED) &&
		    (!!(testcase & CPUID_DISABLED) ^ !!(testcase & MWAIT_DISABLED)))
			continue;

		if (testcase & CPUID_DISABLED) {
			strcat(test, ", CPUID clear");
			vcpu_clear_cpuid_feature(vcpu, X86_FEATURE_MWAIT);
		} else {
			strcat(test, ", CPUID set");
			vcpu_set_cpuid_feature(vcpu, X86_FEATURE_MWAIT);
		}

		if (testcase & MWAIT_DISABLED)
			strcat(test, ", MWAIT disabled");

		vcpu_run(vcpu);
		TEST_ASSERT_KVM_EXIT_REASON(vcpu, KVM_EXIT_IO);

		switch (get_ucall(vcpu, &uc)) {
		case UCALL_ABORT:
			/* Detected in vcpu_run */
			break;
		case UCALL_DONE:
			ksft_test_result_pass("%s\n", test);
			break;
		default:
			TEST_FAIL("Unknown ucall %lu", uc.cmd);
			break;
		}
		kvm_vm_free(vm);
	}
	ksft_finished();

	return 0;
}
