// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2026, Google LLC.
 *
 * Test that KVM correctly virtualizes the PAT MSR and VMCB g_pat field
 * for nested SVM guests:
 *
 * o With nested NPT disabled:
 *     - L1 and L2 share the same PAT
 *     - The vmcb12.g_pat is ignored
 * o With nested NPT enabled:
 *     - Invalid g_pat in vmcb12 should cause VMEXIT_INVALID
 *     - L2 should see vmcb12.g_pat via RDMSR, not L1's PAT
 *     - L2's writes to PAT should be saved to vmcb12 on exit
 *     - L1's PAT should be restored after #VMEXIT from L2
 *     - State save/restore should preserve both L1's and L2's PAT values
 */
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

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

#define L2_GUEST_STACK_SIZE 256

#define PAT_DEFAULT		0x0007040600070406ULL
#define L1_PAT_VALUE		0x0007040600070404ULL  /* Change PA0 to WT */
#define L2_VMCB12_PAT		0x0606060606060606ULL  /* All WB */
#define L2_PAT_MODIFIED		0x0606060606060604ULL  /* Change PA0 to WT */
#define INVALID_PAT_VALUE	0x0808080808080808ULL  /* 8 is reserved */

bool npt_enabled;
int nr_iterations;

static void l2_guest_code(void)
{
	u64 expected_pat = npt_enabled ? L2_VMCB12_PAT : L1_PAT_VALUE;
	int i;

	for (i = 0; i < nr_iterations; i++) {
		GUEST_ASSERT_EQ(rdmsr(MSR_IA32_CR_PAT), expected_pat);
		GUEST_SYNC(1);
		GUEST_ASSERT_EQ(rdmsr(MSR_IA32_CR_PAT), expected_pat);

		wrmsr(MSR_IA32_CR_PAT, L2_PAT_MODIFIED);
		expected_pat = L2_PAT_MODIFIED;

		GUEST_ASSERT_EQ(rdmsr(MSR_IA32_CR_PAT), L2_PAT_MODIFIED);
		GUEST_SYNC(2);
		GUEST_ASSERT_EQ(rdmsr(MSR_IA32_CR_PAT), L2_PAT_MODIFIED);

		vmmcall();
	}
}

static void l1_guest_code(struct svm_test_data *svm)
{
	unsigned long l2_guest_stack[L2_GUEST_STACK_SIZE];
	struct vmcb *vmcb = svm->vmcb;
	int i;

	wrmsr(MSR_IA32_CR_PAT, L1_PAT_VALUE);
	GUEST_ASSERT_EQ(rdmsr(MSR_IA32_CR_PAT), L1_PAT_VALUE);

	generic_svm_setup(svm, l2_guest_code, &l2_guest_stack[L2_GUEST_STACK_SIZE]);

	vmcb->save.g_pat = L2_VMCB12_PAT;
	vmcb->control.intercept &= ~(1ULL << INTERCEPT_MSR_PROT);

	for (i = 0; i < nr_iterations; i++) {
		run_guest(vmcb, svm->vmcb_gpa);

		GUEST_ASSERT_EQ(vmcb->control.exit_code, SVM_EXIT_VMMCALL);

		/*
		 * If NPT is enabled by L1, L2 has a unique PAT and L1's PAT is
		 * unchanged. Otherwise, PAT is shared between L1 and L2.
		 */
		if (npt_enabled) {
			GUEST_ASSERT_EQ(vmcb->save.g_pat, L2_PAT_MODIFIED);
			GUEST_ASSERT_EQ(rdmsr(MSR_IA32_CR_PAT), L1_PAT_VALUE);
		} else {
			GUEST_ASSERT_EQ(rdmsr(MSR_IA32_CR_PAT), L2_PAT_MODIFIED);
		}
		vmcb->save.rip += 3; /* skip over VMMCALL */
	}

	GUEST_DONE();
}

static void l1_guest_code_invalid_gpat(struct svm_test_data *svm)
{
	unsigned long l2_guest_stack[L2_GUEST_STACK_SIZE];
	struct vmcb *vmcb = svm->vmcb;

	/* VMRUN should fail without running L2 */
	generic_svm_setup(svm, NULL, &l2_guest_stack[L2_GUEST_STACK_SIZE]);

	vmcb->save.g_pat = INVALID_PAT_VALUE;
	run_guest(vmcb, svm->vmcb_gpa);

	GUEST_ASSERT_EQ(vmcb->control.exit_code, SVM_EXIT_ERR);
	GUEST_DONE();
}

static void run_test(void *guest_code, bool do_save_restore, int nr_iters)
{
	struct kvm_x86_state *state;
	struct kvm_vcpu *vcpu;
	struct kvm_vm *vm;
	struct ucall uc;
	gva_t svm_gva;

	vm = vm_create_with_one_vcpu(&vcpu, guest_code);
	vm_enable_cap(vm, KVM_CAP_DISABLE_QUIRKS2,
		      KVM_X86_QUIRK_NESTED_SVM_SHARED_PAT);

	if (npt_enabled)
		vm_enable_npt(vm);

	vcpu_alloc_svm(vm, &svm_gva);

	if (npt_enabled)
		tdp_identity_map_default_memslots(vm);

	vcpu_args_set(vcpu, 1, svm_gva);

	nr_iterations = nr_iters;
	sync_global_to_guest(vm, npt_enabled);
	sync_global_to_guest(vm, nr_iterations);

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

		switch (get_ucall(vcpu, &uc)) {
		case UCALL_ABORT:
			REPORT_GUEST_ASSERT(uc);
			/* NOT REACHED */
		case UCALL_SYNC:
			if (do_save_restore) {
				state = vcpu_save_state(vcpu);
				kvm_vm_release(vm);
				vcpu = vm_recreate_with_one_vcpu(vm);
				vm_enable_cap(vm, KVM_CAP_DISABLE_QUIRKS2,
					      KVM_X86_QUIRK_NESTED_SVM_SHARED_PAT);
				vcpu_load_state(vcpu, state);
				kvm_x86_state_cleanup(state);
			}
			break;
		case UCALL_DONE:
			kvm_vm_free(vm);
			return;
		default:
			TEST_FAIL("Unknown ucall %lu", uc.cmd);
		}
	}
}

#define gpat_test(test_name, guest_code, npt_setting)			\
do {									\
	npt_setting;							\
									\
	if (npt_enabled && !kvm_cpu_has(X86_FEATURE_NPT)) {		\
		pr_info("Skipping: " test_name " (no NPT support)\n");	\
		break;							\
	}								\
									\
	pr_info("Testing: " test_name "\n");				\
	run_test(guest_code, false, 1);					\
									\
	if (guest_code == l1_guest_code) {				\
		pr_info("Testing: " test_name " Save/Restore\n");	\
		run_test(guest_code, true, 1);				\
									\
		pr_info("Testing: " test_name " Multiple VMRUNs\n");	\
		run_test(guest_code, false, 10);			\
	}								\
} while (0)

int main(int argc, char *argv[])
{
	TEST_REQUIRE(kvm_cpu_has(X86_FEATURE_SVM));
	TEST_REQUIRE(kvm_has_cap(KVM_CAP_NESTED_STATE));
	TEST_REQUIRE(kvm_check_cap(KVM_CAP_DISABLE_QUIRKS2) &
		     KVM_X86_QUIRK_NESTED_SVM_SHARED_PAT);

	gpat_test("Invalid gPAT", l1_guest_code_invalid_gpat, npt_enabled = true);
	gpat_test("Nested NPT enabled", l1_guest_code, npt_enabled = true);
	gpat_test("Nested NPT disabled", l1_guest_code, npt_enabled = false);
	return 0;
}