// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2026, Google LLC.
 */
#include "kvm_util.h"
#include "vmx.h"
#include "svm_util.h"
#include "kselftest.h"

/*
 * Allocate two VMCB pages for testing. Both pages have different GVAs (shared
 * by both L1 and L2) and L1 GPAs. A single L2 GPA is used such that:
 * - L2 GPA == L1 GPA for VMCB0.
 * - L2 GPA is mapped to L1 GPA for VMCB1 using NPT in L1.
 *
 * This allows testing whether the GPA used by VMSAVE/VMLOAD in L2 is
 * interpreted as a direct L1 GPA or translated using NPT as an L2 GPA, depends
 * on which VMCB is accessed.
 */
#define TEST_MEM_SLOT_INDEX		1
#define TEST_MEM_PAGES			2
#define TEST_MEM_BASE			0xc0000000

#define TEST_GUEST_ADDR(idx)		(TEST_MEM_BASE + (idx) * PAGE_SIZE)

#define TEST_VMCB_L1_GPA(idx)		TEST_GUEST_ADDR(idx)
#define TEST_VMCB_GVA(idx)		TEST_GUEST_ADDR(idx)

#define TEST_VMCB_L2_GPA		TEST_VMCB_L1_GPA(0)

#define L2_GUEST_STACK_SIZE		64

static void l2_guest_code_vmsave(void)
{
	asm volatile("vmsave %0" : : "a"(TEST_VMCB_L2_GPA) : "memory");
}

static void l2_guest_code_vmload(void)
{
	asm volatile("vmload %0" : : "a"(TEST_VMCB_L2_GPA) : "memory");
}

static void l2_guest_code_vmcb(int vmcb_idx)
{
	wrmsr(MSR_KERNEL_GS_BASE, 0xaaaa);
	l2_guest_code_vmsave();

	/* Verify the VMCB used by VMSAVE and update KERNEL_GS_BASE to 0xbbbb */
	GUEST_SYNC(vmcb_idx);

	l2_guest_code_vmload();
	GUEST_ASSERT_EQ(rdmsr(MSR_KERNEL_GS_BASE), 0xbbbb);

	/* Reset MSR_KERNEL_GS_BASE */
	wrmsr(MSR_KERNEL_GS_BASE, 0);
	l2_guest_code_vmsave();

	vmmcall();
}

static void l2_guest_code_vmcb0(void)
{
	l2_guest_code_vmcb(0);
}

static void l2_guest_code_vmcb1(void)
{
	l2_guest_code_vmcb(1);
}

static void l1_guest_code(struct svm_test_data *svm)
{
	unsigned long l2_guest_stack[L2_GUEST_STACK_SIZE];

	/* Each test case initializes the guest RIP below */
	generic_svm_setup(svm, NULL, &l2_guest_stack[L2_GUEST_STACK_SIZE]);

	/* Set VMSAVE/VMLOAD intercepts and make sure they work with.. */
	svm->vmcb->control.intercept |= (BIT_ULL(INTERCEPT_VMSAVE) |
					 BIT_ULL(INTERCEPT_VMLOAD));

	 /* ..VIRTUAL_VMLOAD_VMSAVE_ENABLE_MASK cleared.. */
	svm->vmcb->control.virt_ext &= ~VIRTUAL_VMLOAD_VMSAVE_ENABLE_MASK;

	svm->vmcb->save.rip = (u64)l2_guest_code_vmsave;
	run_guest(svm->vmcb, svm->vmcb_gpa);
	GUEST_ASSERT_EQ(svm->vmcb->control.exit_code, SVM_EXIT_VMSAVE);

	svm->vmcb->save.rip = (u64)l2_guest_code_vmload;
	run_guest(svm->vmcb, svm->vmcb_gpa);
	GUEST_ASSERT_EQ(svm->vmcb->control.exit_code, SVM_EXIT_VMLOAD);

	/* ..and VIRTUAL_VMLOAD_VMSAVE_ENABLE_MASK set */
	svm->vmcb->control.virt_ext |= VIRTUAL_VMLOAD_VMSAVE_ENABLE_MASK;

	svm->vmcb->save.rip = (u64)l2_guest_code_vmsave;
	run_guest(svm->vmcb, svm->vmcb_gpa);
	GUEST_ASSERT_EQ(svm->vmcb->control.exit_code, SVM_EXIT_VMSAVE);

	svm->vmcb->save.rip = (u64)l2_guest_code_vmload;
	run_guest(svm->vmcb, svm->vmcb_gpa);
	GUEST_ASSERT_EQ(svm->vmcb->control.exit_code, SVM_EXIT_VMLOAD);

	/* Now clear the intercepts to test VMSAVE/VMLOAD behavior */
	svm->vmcb->control.intercept &= ~(BIT_ULL(INTERCEPT_VMSAVE) |
					  BIT_ULL(INTERCEPT_VMLOAD));

	/*
	 * Without VIRTUAL_VMLOAD_VMSAVE_ENABLE_MASK, the GPA will be
	 * interpreted as an L1 GPA, so VMCB0 should be used.
	 */
	svm->vmcb->save.rip = (u64)l2_guest_code_vmcb0;
	svm->vmcb->control.virt_ext &= ~VIRTUAL_VMLOAD_VMSAVE_ENABLE_MASK;
	run_guest(svm->vmcb, svm->vmcb_gpa);
	GUEST_ASSERT_EQ(svm->vmcb->control.exit_code, SVM_EXIT_VMMCALL);

	/*
	 * With VIRTUAL_VMLOAD_VMSAVE_ENABLE_MASK, the GPA will be interpeted as
	 * an L2 GPA, and translated through the NPT to VMCB1.
	 */
	svm->vmcb->save.rip = (u64)l2_guest_code_vmcb1;
	svm->vmcb->control.virt_ext |= VIRTUAL_VMLOAD_VMSAVE_ENABLE_MASK;
	run_guest(svm->vmcb, svm->vmcb_gpa);
	GUEST_ASSERT_EQ(svm->vmcb->control.exit_code, SVM_EXIT_VMMCALL);

	GUEST_DONE();
}

int main(int argc, char *argv[])
{
	vm_vaddr_t nested_gva = 0;
	struct vmcb *test_vmcb[2];
	struct kvm_vcpu *vcpu;
	struct kvm_vm *vm;
	int i;

	TEST_REQUIRE(kvm_cpu_has(X86_FEATURE_SVM));
	TEST_REQUIRE(kvm_cpu_has(X86_FEATURE_NPT));
	TEST_REQUIRE(kvm_cpu_has(X86_FEATURE_V_VMSAVE_VMLOAD));

	vm = vm_create_with_one_vcpu(&vcpu, l1_guest_code);
	vm_enable_tdp(vm);

	vcpu_alloc_svm(vm, &nested_gva);
	vcpu_args_set(vcpu, 1, nested_gva);

	vm_userspace_mem_region_add(vm, VM_MEM_SRC_ANONYMOUS,
				    TEST_MEM_BASE, TEST_MEM_SLOT_INDEX,
				    TEST_MEM_PAGES, 0);

	for (i = 0; i <= 1; i++) {
		virt_map(vm, TEST_VMCB_GVA(i), TEST_VMCB_L1_GPA(i), 1);
		test_vmcb[i] = (struct vmcb *)addr_gva2hva(vm, TEST_VMCB_GVA(i));
	}

	tdp_identity_map_default_memslots(vm);

	/*
	 * L2 GPA == L1_GPA(0), but map it to L1_GPA(1), to allow testing
	 * whether the L2 GPA is interpreted as an L1 GPA or translated through
	 * the NPT.
	 */
	TEST_ASSERT_EQ(TEST_VMCB_L2_GPA, TEST_VMCB_L1_GPA(0));
	tdp_map(vm, TEST_VMCB_L2_GPA, TEST_VMCB_L1_GPA(1), PAGE_SIZE);

	for (;;) {
		struct ucall uc;

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

		switch (get_ucall(vcpu, &uc)) {
		case UCALL_ABORT:
			REPORT_GUEST_ASSERT(uc);
		case UCALL_SYNC:
			i = uc.args[1];
			TEST_ASSERT(i == 0 || i == 1, "Unexpected VMCB idx: %d", i);

			/*
			 * Check that only the expected VMCB has KERNEL_GS_BASE
			 * set to 0xaaaa, and update it to 0xbbbb.
			 */
			TEST_ASSERT_EQ(test_vmcb[i]->save.kernel_gs_base, 0xaaaa);
			TEST_ASSERT_EQ(test_vmcb[1-i]->save.kernel_gs_base, 0);
			test_vmcb[i]->save.kernel_gs_base = 0xbbbb;
			break;
		case UCALL_DONE:
			goto done;
		default:
			TEST_FAIL("Unknown ucall %lu", uc.cmd);
		}
	}

done:
	kvm_vm_free(vm);
	return 0;
}