// SPDX-License-Identifier: GPL-2.0-only
/*
 * psci_test - Tests relating to KVM's PSCI implementation.
 *
 * Copyright (c) 2021 Google LLC.
 *
 * This test includes:
 *  - A regression test for a race between KVM servicing the PSCI CPU_ON call
 *    and userspace reading the targeted vCPU's registers.
 *  - A test for KVM's handling of PSCI SYSTEM_SUSPEND and the associated
 *    KVM_SYSTEM_EVENT_SUSPEND UAPI.
 */

#include <linux/kernel.h>
#include <linux/psci.h>
#include <asm/cputype.h>

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

#define CPU_ON_ENTRY_ADDR 0xfeedf00dul
#define CPU_ON_CONTEXT_ID 0xdeadc0deul

static u64 psci_cpu_on(u64 target_cpu, u64 entry_addr, u64 context_id)
{
	struct arm_smccc_res res;

	do_smccc(PSCI_0_2_FN64_CPU_ON, target_cpu, entry_addr, context_id,
		  0, 0, 0, 0, &res);

	return res.a0;
}

static u64 psci_affinity_info(u64 target_affinity, u64 lowest_affinity_level)
{
	struct arm_smccc_res res;

	do_smccc(PSCI_0_2_FN64_AFFINITY_INFO, target_affinity, lowest_affinity_level,
		  0, 0, 0, 0, 0, &res);

	return res.a0;
}

static u64 psci_system_suspend(u64 entry_addr, u64 context_id)
{
	struct arm_smccc_res res;

	do_smccc(PSCI_1_0_FN64_SYSTEM_SUSPEND, entry_addr, context_id,
		  0, 0, 0, 0, 0, &res);

	return res.a0;
}

static u64 psci_system_off2(u64 type, u64 cookie)
{
	struct arm_smccc_res res;

	do_smccc(PSCI_1_3_FN64_SYSTEM_OFF2, type, cookie, 0, 0, 0, 0, 0, &res);

	return res.a0;
}

static u64 psci_features(u32 func_id)
{
	struct arm_smccc_res res;

	do_smccc(PSCI_1_0_FN_PSCI_FEATURES, func_id, 0, 0, 0, 0, 0, 0, &res);

	return res.a0;
}

static void vcpu_power_off(struct kvm_vcpu *vcpu)
{
	struct kvm_mp_state mp_state = {
		.mp_state = KVM_MP_STATE_STOPPED,
	};

	vcpu_mp_state_set(vcpu, &mp_state);
}

static struct kvm_vm *setup_vm(void *guest_code, struct kvm_vcpu **source,
			       struct kvm_vcpu **target)
{
	struct kvm_vcpu_init init;
	struct kvm_vm *vm;

	vm = vm_create(2);

	kvm_get_default_vcpu_target(vm, &init);
	init.features[0] |= (1 << KVM_ARM_VCPU_PSCI_0_2);

	*source = aarch64_vcpu_add(vm, 0, &init, guest_code);
	*target = aarch64_vcpu_add(vm, 1, &init, guest_code);

	kvm_arch_vm_finalize_vcpus(vm);
	return vm;
}

static void enter_guest(struct kvm_vcpu *vcpu)
{
	struct ucall uc;

	vcpu_run(vcpu);
	if (get_ucall(vcpu, &uc) == UCALL_ABORT)
		REPORT_GUEST_ASSERT(uc);
}

static void assert_vcpu_reset(struct kvm_vcpu *vcpu)
{
	u64 obs_pc, obs_x0;

	obs_pc = vcpu_get_reg(vcpu, ARM64_CORE_REG(regs.pc));
	obs_x0 = vcpu_get_reg(vcpu, ARM64_CORE_REG(regs.regs[0]));

	TEST_ASSERT(obs_pc == CPU_ON_ENTRY_ADDR,
		    "unexpected target cpu pc: %lx (expected: %lx)",
		    obs_pc, CPU_ON_ENTRY_ADDR);
	TEST_ASSERT(obs_x0 == CPU_ON_CONTEXT_ID,
		    "unexpected target context id: %lx (expected: %lx)",
		    obs_x0, CPU_ON_CONTEXT_ID);
}

static void guest_test_cpu_on(u64 target_cpu)
{
	u64 target_state;

	GUEST_ASSERT(!psci_cpu_on(target_cpu, CPU_ON_ENTRY_ADDR, CPU_ON_CONTEXT_ID));

	do {
		target_state = psci_affinity_info(target_cpu, 0);

		GUEST_ASSERT((target_state == PSCI_0_2_AFFINITY_LEVEL_ON) ||
			     (target_state == PSCI_0_2_AFFINITY_LEVEL_OFF));
	} while (target_state != PSCI_0_2_AFFINITY_LEVEL_ON);

	GUEST_DONE();
}

static void host_test_cpu_on(void)
{
	struct kvm_vcpu *source, *target;
	u64 target_mpidr;
	struct kvm_vm *vm;
	struct ucall uc;

	vm = setup_vm(guest_test_cpu_on, &source, &target);

	/*
	 * make sure the target is already off when executing the test.
	 */
	vcpu_power_off(target);

	target_mpidr = vcpu_get_reg(target, KVM_ARM64_SYS_REG(SYS_MPIDR_EL1));
	vcpu_args_set(source, 1, target_mpidr & MPIDR_HWID_BITMASK);
	enter_guest(source);

	if (get_ucall(source, &uc) != UCALL_DONE)
		TEST_FAIL("Unhandled ucall: %lu", uc.cmd);

	assert_vcpu_reset(target);
	kvm_vm_free(vm);
}

static void guest_test_system_suspend(void)
{
	u64 ret;

	/* assert that SYSTEM_SUSPEND is discoverable */
	GUEST_ASSERT(!psci_features(PSCI_1_0_FN_SYSTEM_SUSPEND));
	GUEST_ASSERT(!psci_features(PSCI_1_0_FN64_SYSTEM_SUSPEND));

	ret = psci_system_suspend(CPU_ON_ENTRY_ADDR, CPU_ON_CONTEXT_ID);
	GUEST_SYNC(ret);
}

static void host_test_system_suspend(void)
{
	struct kvm_vcpu *source, *target;
	struct kvm_run *run;
	struct kvm_vm *vm;

	vm = setup_vm(guest_test_system_suspend, &source, &target);
	vm_enable_cap(vm, KVM_CAP_ARM_SYSTEM_SUSPEND, 0);

	vcpu_power_off(target);
	run = source->run;

	enter_guest(source);

	TEST_ASSERT_KVM_EXIT_REASON(source, KVM_EXIT_SYSTEM_EVENT);
	TEST_ASSERT(run->system_event.type == KVM_SYSTEM_EVENT_SUSPEND,
		    "Unhandled system event: %u (expected: %u)",
		    run->system_event.type, KVM_SYSTEM_EVENT_SUSPEND);

	kvm_vm_free(vm);
}

static void guest_test_system_off2(void)
{
	u64 ret;

	/* assert that SYSTEM_OFF2 is discoverable */
	GUEST_ASSERT(psci_features(PSCI_1_3_FN_SYSTEM_OFF2) &
		     PSCI_1_3_OFF_TYPE_HIBERNATE_OFF);
	GUEST_ASSERT(psci_features(PSCI_1_3_FN64_SYSTEM_OFF2) &
		     PSCI_1_3_OFF_TYPE_HIBERNATE_OFF);

	/* With non-zero 'cookie' field, it should fail */
	ret = psci_system_off2(PSCI_1_3_OFF_TYPE_HIBERNATE_OFF, 1);
	GUEST_ASSERT(ret == PSCI_RET_INVALID_PARAMS);

	/*
	 * This would normally never return, so KVM sets the return value
	 * to PSCI_RET_INTERNAL_FAILURE. The test case *does* return, so
	 * that it can test both values for HIBERNATE_OFF.
	 */
	ret = psci_system_off2(PSCI_1_3_OFF_TYPE_HIBERNATE_OFF, 0);
	GUEST_ASSERT(ret == PSCI_RET_INTERNAL_FAILURE);

	/*
	 * Revision F.b of the PSCI v1.3 specification documents zero as an
	 * alias for HIBERNATE_OFF, since that's the value used in earlier
	 * revisions of the spec and some implementations in the field.
	 */
	ret = psci_system_off2(0, 1);
	GUEST_ASSERT(ret == PSCI_RET_INVALID_PARAMS);

	ret = psci_system_off2(0, 0);
	GUEST_ASSERT(ret == PSCI_RET_INTERNAL_FAILURE);

	GUEST_DONE();
}

static void host_test_system_off2(void)
{
	struct kvm_vcpu *source, *target;
	struct kvm_mp_state mps;
	u64 psci_version = 0;
	int nr_shutdowns = 0;
	struct kvm_run *run;
	struct ucall uc;

	setup_vm(guest_test_system_off2, &source, &target);

	psci_version = vcpu_get_reg(target, KVM_REG_ARM_PSCI_VERSION);

	TEST_ASSERT(psci_version >= PSCI_VERSION(1, 3),
		    "Unexpected PSCI version %lu.%lu",
		    PSCI_VERSION_MAJOR(psci_version),
		    PSCI_VERSION_MINOR(psci_version));

	vcpu_power_off(target);
	run = source->run;

	enter_guest(source);
	while (run->exit_reason == KVM_EXIT_SYSTEM_EVENT) {
		TEST_ASSERT(run->system_event.type == KVM_SYSTEM_EVENT_SHUTDOWN,
			    "Unhandled system event: %u (expected: %u)",
			    run->system_event.type, KVM_SYSTEM_EVENT_SHUTDOWN);
		TEST_ASSERT(run->system_event.ndata >= 1,
			    "Unexpected amount of system event data: %u (expected, >= 1)",
			    run->system_event.ndata);
		TEST_ASSERT(run->system_event.data[0] & KVM_SYSTEM_EVENT_SHUTDOWN_FLAG_PSCI_OFF2,
			    "PSCI_OFF2 flag not set. Flags %llu (expected %llu)",
			    run->system_event.data[0], KVM_SYSTEM_EVENT_SHUTDOWN_FLAG_PSCI_OFF2);

		nr_shutdowns++;

		/* Restart the vCPU */
	        mps.mp_state = KVM_MP_STATE_RUNNABLE;
		vcpu_mp_state_set(source, &mps);

		enter_guest(source);
	}

	TEST_ASSERT(get_ucall(source, &uc) == UCALL_DONE, "Guest did not exit cleanly");
	TEST_ASSERT(nr_shutdowns == 2, "Two shutdown events were expected, but saw %d", nr_shutdowns);
}

int main(void)
{
	TEST_REQUIRE(kvm_has_cap(KVM_CAP_ARM_SYSTEM_SUSPEND));

	host_test_cpu_on();
	host_test_system_suspend();
	host_test_system_off2();
	return 0;
}