// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) 2020 - Google Inc
 * Author: Andrew Scull <ascull@google.com>
 */

#include <hyp/adjust_pc.h>
#include <hyp/switch.h>

#include <linux/irqchip/arm-gic-v3.h>

#include <asm/pgtable-types.h>
#include <asm/kvm_asm.h>
#include <asm/kvm_emulate.h>
#include <asm/kvm_host.h>
#include <asm/kvm_hyp.h>
#include <asm/kvm_hypevents.h>
#include <asm/kvm_mmu.h>

#include <nvhe/ffa.h>
#include <nvhe/mem_protect.h>
#include <nvhe/mm.h>
#include <nvhe/pkvm.h>
#include <nvhe/trace.h>
#include <nvhe/trap_handler.h>

DEFINE_PER_CPU(struct kvm_nvhe_init_params, kvm_init_params);

/* Number of implemented GICv3 LRs. Used by flush_hyp_vcpu(). */
unsigned int hyp_gicv3_nr_lr;

void __kvm_hyp_host_forward_smc(struct kvm_cpu_context *host_ctxt);

static void __hyp_sve_save_guest(struct kvm_vcpu *vcpu)
{
	__vcpu_assign_sys_reg(vcpu, ZCR_EL1, read_sysreg_el1(SYS_ZCR));
	/*
	 * On saving/restoring guest sve state, always use the maximum VL for
	 * the guest. The layout of the data when saving the sve state depends
	 * on the VL, so use a consistent (i.e., the maximum) guest VL.
	 */
	sve_cond_update_zcr_vq(vcpu_sve_max_vq(vcpu) - 1, SYS_ZCR_EL2);
	sve_save_state(kern_hyp_va(vcpu->arch.sve_state), true);
	fpsimd_save_common(&vcpu->arch.ctxt.fp_regs);
	write_sysreg_s(sve_vq_from_vl(kvm_host_sve_max_vl) - 1, SYS_ZCR_EL2);
}

static void __hyp_sve_restore_host(void)
{
	struct kvm_cpu_context *hctxt = host_data_ptr(host_ctxt);
	struct arm64_sve_state *sve_regs = *host_data_ptr(sve_regs);

	/*
	 * On saving/restoring host sve state, always use the maximum VL for
	 * the host. The layout of the data when saving the sve state depends
	 * on the VL, so use a consistent (i.e., the maximum) host VL.
	 *
	 * Note that this constrains the PE to the maximum shared VL
	 * that was discovered, if we wish to use larger VLs this will
	 * need to be revisited.
	 */
	write_sysreg_s(sve_vq_from_vl(kvm_host_sve_max_vl) - 1, SYS_ZCR_EL2);
	sve_load_state(sve_regs, true);
	fpsimd_load_common(&hctxt->fp_regs);
	write_sysreg_el1(ctxt_sys_reg(hctxt, ZCR_EL1), SYS_ZCR);
}

static void fpsimd_sve_flush(void)
{
	*host_data_ptr(fp_owner) = FP_STATE_HOST_OWNED;
}

static void fpsimd_sve_sync(struct kvm_vcpu *vcpu)
{
	struct kvm_cpu_context *hctxt = host_data_ptr(host_ctxt);
	bool has_fpmr;

	if (!guest_owns_fp_regs())
		return;

	/*
	 * Traps have been disabled by __deactivate_cptr_traps(), but there
	 * hasn't necessarily been a context synchronization event yet.
	 */
	isb();

	if (vcpu_has_sve(vcpu))
		__hyp_sve_save_guest(vcpu);
	else
		fpsimd_save_state(&vcpu->arch.ctxt.fp_regs);

	has_fpmr = kvm_has_fpmr(kern_hyp_va(vcpu->kvm));
	if (has_fpmr)
		__vcpu_assign_sys_reg(vcpu, FPMR, read_sysreg_s(SYS_FPMR));

	if (system_supports_sve())
		__hyp_sve_restore_host();
	else
		fpsimd_load_state(&hctxt->fp_regs);

	if (has_fpmr)
		write_sysreg_s(ctxt_sys_reg(hctxt, FPMR), SYS_FPMR);

	*host_data_ptr(fp_owner) = FP_STATE_HOST_OWNED;
}

static void flush_hyp_vgic_state(struct pkvm_hyp_vcpu *hyp_vcpu)
{
	struct kvm_vcpu *host_vcpu = hyp_vcpu->host_vcpu;
	struct vgic_v3_cpu_if *host_cpu_if, *hyp_cpu_if;
	unsigned int used_lrs, i;

	host_cpu_if	= &host_vcpu->arch.vgic_cpu.vgic_v3;
	hyp_cpu_if	= &hyp_vcpu->vcpu.arch.vgic_cpu.vgic_v3;

	used_lrs	= host_cpu_if->used_lrs;
	used_lrs	= min(used_lrs, hyp_gicv3_nr_lr);

	hyp_cpu_if->vgic_hcr	= host_cpu_if->vgic_hcr;
	/* Should be a one-off */
	hyp_cpu_if->vgic_sre	= (ICC_SRE_EL1_DIB |
				   ICC_SRE_EL1_DFB |
				   ICC_SRE_EL1_SRE);
	hyp_cpu_if->used_lrs	= used_lrs;

	for (i = 0; i < used_lrs; i++)
		hyp_cpu_if->vgic_lr[i] = host_cpu_if->vgic_lr[i];
}

static void sync_hyp_vgic_state(struct pkvm_hyp_vcpu *hyp_vcpu)
{
	struct kvm_vcpu *host_vcpu = hyp_vcpu->host_vcpu;
	struct vgic_v3_cpu_if *host_cpu_if, *hyp_cpu_if;
	unsigned int i;

	host_cpu_if	= &host_vcpu->arch.vgic_cpu.vgic_v3;
	hyp_cpu_if	= &hyp_vcpu->vcpu.arch.vgic_cpu.vgic_v3;

	host_cpu_if->vgic_hcr = hyp_cpu_if->vgic_hcr;
	host_cpu_if->vgic_vmcr = hyp_cpu_if->vgic_vmcr;

	for (i = 0; i < hyp_cpu_if->used_lrs; i++)
		host_cpu_if->vgic_lr[i] = hyp_cpu_if->vgic_lr[i];
}

static void __copy_vcpu_state(const struct kvm_vcpu *from_vcpu,
			      struct kvm_vcpu *to_vcpu)
{
	int i;

	to_vcpu->arch.ctxt.regs		= from_vcpu->arch.ctxt.regs;
	to_vcpu->arch.ctxt.spsr_abt	= from_vcpu->arch.ctxt.spsr_abt;
	to_vcpu->arch.ctxt.spsr_und	= from_vcpu->arch.ctxt.spsr_und;
	to_vcpu->arch.ctxt.spsr_irq	= from_vcpu->arch.ctxt.spsr_irq;
	to_vcpu->arch.ctxt.spsr_fiq	= from_vcpu->arch.ctxt.spsr_fiq;
	to_vcpu->arch.ctxt.fp_regs	= from_vcpu->arch.ctxt.fp_regs;

	/*
	 * Copy the sysregs, but don't mess with the timer state which
	 * is directly handled by EL1 and is expected to be preserved.
	 * enum vcpu_sysreg is sparse: VNCR-mapped registers take values
	 * derived from their VNCR page offset, so the timer registers do
	 * not form a contiguous numeric range and must be skipped by name.
	 */
	for (i = 1; i < NR_SYS_REGS; i++) {
		switch (i) {
		case CNTVOFF_EL2:
		case CNTV_CVAL_EL0:
		case CNTV_CTL_EL0:
		case CNTP_CVAL_EL0:
		case CNTP_CTL_EL0:
			continue;
		}
		to_vcpu->arch.ctxt.sys_regs[i] = from_vcpu->arch.ctxt.sys_regs[i];
	}
}

static void sync_hyp_vcpu_state(struct pkvm_hyp_vcpu *hyp_vcpu)
{
	__copy_vcpu_state(&hyp_vcpu->vcpu, hyp_vcpu->host_vcpu);
}

static void flush_hyp_vcpu_state(struct pkvm_hyp_vcpu *hyp_vcpu)
{
	__copy_vcpu_state(hyp_vcpu->host_vcpu, &hyp_vcpu->vcpu);
}

static void flush_debug_state(struct pkvm_hyp_vcpu *hyp_vcpu)
{
	struct kvm_vcpu *host_vcpu = hyp_vcpu->host_vcpu;

	hyp_vcpu->vcpu.arch.debug_owner = host_vcpu->arch.debug_owner;

	if (kvm_guest_owns_debug_regs(&hyp_vcpu->vcpu)) {
		hyp_vcpu->vcpu.arch.vcpu_debug_state = host_vcpu->arch.vcpu_debug_state;
	} else if (kvm_host_owns_debug_regs(&hyp_vcpu->vcpu)) {
		hyp_vcpu->vcpu.arch.external_debug_state = host_vcpu->arch.external_debug_state;
		/*
		 * The world switch loads MDSCR_EL1 from external_mdscr_el1
		 * (ctxt_mdscr_el1()).
		 */
		hyp_vcpu->vcpu.arch.external_mdscr_el1 = host_vcpu->arch.external_mdscr_el1;
	}
}

static void sync_debug_state(struct pkvm_hyp_vcpu *hyp_vcpu)
{
	struct kvm_vcpu *host_vcpu = hyp_vcpu->host_vcpu;

	if (kvm_guest_owns_debug_regs(&hyp_vcpu->vcpu))
		host_vcpu->arch.vcpu_debug_state = hyp_vcpu->vcpu.arch.vcpu_debug_state;
	else if (kvm_host_owns_debug_regs(&hyp_vcpu->vcpu))
		host_vcpu->arch.external_debug_state = hyp_vcpu->vcpu.arch.external_debug_state;
}

static void flush_hyp_vcpu(struct pkvm_hyp_vcpu *hyp_vcpu)
{
	struct kvm_vcpu *host_vcpu = hyp_vcpu->host_vcpu;

	fpsimd_sve_flush();
	flush_debug_state(hyp_vcpu);

	/*
	 * If we deal with a non-protected guest and the state is potentially
	 * dirty (from a host perspective), copy the state back into the hyp
	 * vcpu.
	 */
	if (!pkvm_hyp_vcpu_is_protected(hyp_vcpu)) {
		if (vcpu_get_flag(host_vcpu, PKVM_HOST_STATE_DIRTY))
			flush_hyp_vcpu_state(hyp_vcpu);
	} else {
		hyp_vcpu->vcpu.arch.ctxt = host_vcpu->arch.ctxt;
	}

	/* __hyp_running_vcpu must be NULL in a guest context. */
	hyp_vcpu->vcpu.arch.ctxt.__hyp_running_vcpu = NULL;

	hyp_vcpu->vcpu.arch.mdcr_el2	= host_vcpu->arch.mdcr_el2;
	/*
	 * HCR_EL2.VSE is host-owned (a pending virtual SError to inject), not a
	 * trap-control bit, so it must flow to the hyp vCPU alongside TWI/TWE
	 * for the vSError to be delivered. sync_hyp_vcpu() reflects it back.
	 */
	hyp_vcpu->vcpu.arch.hcr_el2 &= ~(HCR_TWI | HCR_TWE | HCR_VSE);
	hyp_vcpu->vcpu.arch.hcr_el2 |= READ_ONCE(host_vcpu->arch.hcr_el2) &
						 (HCR_TWI | HCR_TWE | HCR_VSE);

	hyp_vcpu->vcpu.arch.iflags	= host_vcpu->arch.iflags;

	hyp_vcpu->vcpu.arch.vsesr_el2	= host_vcpu->arch.vsesr_el2;

	flush_hyp_vgic_state(hyp_vcpu);

	hyp_vcpu->vcpu.arch.pid = host_vcpu->arch.pid;
}

static void sync_hyp_vcpu(struct pkvm_hyp_vcpu *hyp_vcpu)
{
	struct kvm_vcpu *host_vcpu = hyp_vcpu->host_vcpu;

	fpsimd_sve_sync(&hyp_vcpu->vcpu);
	sync_debug_state(hyp_vcpu);

	if (pkvm_hyp_vcpu_is_protected(hyp_vcpu)) {
		host_vcpu->arch.ctxt = hyp_vcpu->vcpu.arch.ctxt;
	} else {
		/*
		 * PC feeds trace_kvm_exit(), PSTATE.SS the host software-step
		 * machine, and both run before the next on-demand ctxt sync.
		 */
		host_vcpu->arch.ctxt.regs.pc = hyp_vcpu->vcpu.arch.ctxt.regs.pc;
		host_vcpu->arch.ctxt.regs.pstate = hyp_vcpu->vcpu.arch.ctxt.regs.pstate;
	}

	host_vcpu->arch.fault		= hyp_vcpu->vcpu.arch.fault;

	host_vcpu->arch.iflags		= hyp_vcpu->vcpu.arch.iflags;

	sync_hyp_vgic_state(hyp_vcpu);
}

static void handle___pkvm_vcpu_load(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(pkvm_handle_t, handle, host_ctxt, 1);
	DECLARE_REG(unsigned int, vcpu_idx, host_ctxt, 2);
	DECLARE_REG(u64, hcr_el2, host_ctxt, 3);
	struct pkvm_hyp_vcpu *hyp_vcpu;

	hyp_vcpu = pkvm_load_hyp_vcpu(handle, vcpu_idx);
	if (!hyp_vcpu)
		return;

	if (pkvm_hyp_vcpu_is_protected(hyp_vcpu)) {
		/* Propagate WFx trapping flags */
		hyp_vcpu->vcpu.arch.hcr_el2 &= ~(HCR_TWE | HCR_TWI);
		hyp_vcpu->vcpu.arch.hcr_el2 |= hcr_el2 & (HCR_TWE | HCR_TWI);
	} else {
		memcpy(&hyp_vcpu->vcpu.arch.fgt, hyp_vcpu->host_vcpu->arch.fgt,
		       sizeof(hyp_vcpu->vcpu.arch.fgt));
	}
}

static void handle___pkvm_vcpu_put(struct kvm_cpu_context *host_ctxt)
{
	struct pkvm_hyp_vcpu *hyp_vcpu = pkvm_get_loaded_hyp_vcpu();

	if (hyp_vcpu) {
		struct kvm_vcpu *host_vcpu = hyp_vcpu->host_vcpu;

		if (!pkvm_hyp_vcpu_is_protected(hyp_vcpu) &&
		    !vcpu_get_flag(host_vcpu, PKVM_HOST_STATE_DIRTY)) {
			sync_hyp_vcpu_state(hyp_vcpu);
		}

		pkvm_put_hyp_vcpu(hyp_vcpu);
	}
}

static void handle___pkvm_vcpu_sync_state(struct kvm_cpu_context *host_ctxt)
{
	struct pkvm_hyp_vcpu *hyp_vcpu;

	hyp_vcpu = pkvm_get_loaded_hyp_vcpu();
	if (!hyp_vcpu || pkvm_hyp_vcpu_is_protected(hyp_vcpu))
		return;

	sync_hyp_vcpu_state(hyp_vcpu);
}

static struct kvm_vcpu *__get_host_hyp_vcpus(struct kvm_vcpu *arg,
					     struct pkvm_hyp_vcpu **hyp_vcpup)
{
	struct kvm_vcpu *host_vcpu = kern_hyp_va(arg);
	struct pkvm_hyp_vcpu *hyp_vcpu = NULL;

	if (unlikely(is_protected_kvm_enabled())) {
		hyp_vcpu = pkvm_get_loaded_hyp_vcpu();

		if (!hyp_vcpu || hyp_vcpu->host_vcpu != host_vcpu) {
			hyp_vcpu = NULL;
			host_vcpu = NULL;
		}
	}

	*hyp_vcpup = hyp_vcpu;
	return host_vcpu;
}

#define get_host_hyp_vcpus(ctxt, regnr, hyp_vcpup)			\
	({								\
		DECLARE_REG(struct kvm_vcpu *, __vcpu, ctxt, regnr);	\
		__get_host_hyp_vcpus(__vcpu, hyp_vcpup);		\
	})

#define get_host_hyp_vcpus_from_vgic_v3_cpu_if(ctxt, regnr, hyp_vcpup)		\
	({									\
		DECLARE_REG(struct vgic_v3_cpu_if *, cif, ctxt, regnr);\
		struct kvm_vcpu *__vcpu = container_of(cif,			\
						       struct kvm_vcpu,		\
						       arch.vgic_cpu.vgic_v3);	\
										\
		__get_host_hyp_vcpus(__vcpu, hyp_vcpup);			\
	})

static void handle___kvm_vcpu_run(struct kvm_cpu_context *host_ctxt)
{
	struct pkvm_hyp_vcpu *hyp_vcpu;
	struct kvm_vcpu *host_vcpu;
	int ret;

	host_vcpu = get_host_hyp_vcpus(host_ctxt, 1, &hyp_vcpu);

	if (!host_vcpu) {
		ret = -EINVAL;
		goto out;
	}

	if (unlikely(hyp_vcpu)) {
		/*
		 * KVM (and pKVM) doesn't support SME guests for now, and
		 * ensures that SME features aren't enabled in pstate when
		 * loading a vcpu. Therefore, if SME features enabled the host
		 * is misbehaving.
		 */
		if (unlikely(system_supports_sme() && read_sysreg_s(SYS_SVCR))) {
			ret = -EINVAL;
			goto out;
		}

		flush_hyp_vcpu(hyp_vcpu);

		ret = __kvm_vcpu_run(&hyp_vcpu->vcpu);

		sync_hyp_vcpu(hyp_vcpu);
	} else {
		/* The host is fully trusted, run its vCPU directly. */
		fpsimd_lazy_switch_to_guest(host_vcpu);
		ret = __kvm_vcpu_run(host_vcpu);
		fpsimd_lazy_switch_to_host(host_vcpu);
	}
out:
	cpu_reg(host_ctxt, 1) =  ret;
}

static int pkvm_refill_memcache(struct pkvm_hyp_vcpu *hyp_vcpu)
{
	struct kvm_vcpu *host_vcpu = hyp_vcpu->host_vcpu;

	return refill_memcache(&hyp_vcpu->vcpu.arch.pkvm_memcache,
			       host_vcpu->arch.pkvm_memcache.nr_pages,
			       &host_vcpu->arch.pkvm_memcache);
}

static void handle___pkvm_host_donate_guest(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(u64, pfn, host_ctxt, 1);
	DECLARE_REG(u64, gfn, host_ctxt, 2);
	struct pkvm_hyp_vcpu *hyp_vcpu;
	int ret = -EINVAL;

	hyp_vcpu = pkvm_get_loaded_hyp_vcpu();
	if (!hyp_vcpu || !pkvm_hyp_vcpu_is_protected(hyp_vcpu))
		goto out;

	ret = pkvm_refill_memcache(hyp_vcpu);
	if (ret)
		goto out;

	ret = __pkvm_host_donate_guest(pfn, gfn, hyp_vcpu);
out:
	cpu_reg(host_ctxt, 1) =  ret;
}

static void handle___pkvm_host_share_guest(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(u64, pfn, host_ctxt, 1);
	DECLARE_REG(u64, gfn, host_ctxt, 2);
	DECLARE_REG(u64, nr_pages, host_ctxt, 3);
	DECLARE_REG(enum kvm_pgtable_prot, prot, host_ctxt, 4);
	struct pkvm_hyp_vcpu *hyp_vcpu;
	int ret = -EINVAL;

	hyp_vcpu = pkvm_get_loaded_hyp_vcpu();
	if (!hyp_vcpu || pkvm_hyp_vcpu_is_protected(hyp_vcpu))
		goto out;

	ret = pkvm_refill_memcache(hyp_vcpu);
	if (ret)
		goto out;

	ret = __pkvm_host_share_guest(pfn, gfn, nr_pages, hyp_vcpu, prot);
out:
	cpu_reg(host_ctxt, 1) =  ret;
}

static void handle___pkvm_host_unshare_guest(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(pkvm_handle_t, handle, host_ctxt, 1);
	DECLARE_REG(u64, gfn, host_ctxt, 2);
	DECLARE_REG(u64, nr_pages, host_ctxt, 3);
	struct pkvm_hyp_vm *hyp_vm;
	int ret = -EINVAL;

	hyp_vm = get_np_pkvm_hyp_vm(handle);
	if (!hyp_vm)
		goto out;

	ret = __pkvm_host_unshare_guest(gfn, nr_pages, hyp_vm);
	put_pkvm_hyp_vm(hyp_vm);
out:
	cpu_reg(host_ctxt, 1) =  ret;
}

static void handle___pkvm_host_relax_perms_guest(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(u64, gfn, host_ctxt, 1);
	DECLARE_REG(enum kvm_pgtable_prot, prot, host_ctxt, 2);
	struct pkvm_hyp_vcpu *hyp_vcpu;
	int ret = -EINVAL;

	hyp_vcpu = pkvm_get_loaded_hyp_vcpu();
	if (!hyp_vcpu || pkvm_hyp_vcpu_is_protected(hyp_vcpu))
		goto out;

	ret = __pkvm_host_relax_perms_guest(gfn, hyp_vcpu, prot);
out:
	cpu_reg(host_ctxt, 1) = ret;
}

static void handle___pkvm_host_wrprotect_guest(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(pkvm_handle_t, handle, host_ctxt, 1);
	DECLARE_REG(u64, gfn, host_ctxt, 2);
	DECLARE_REG(u64, nr_pages, host_ctxt, 3);
	struct pkvm_hyp_vm *hyp_vm;
	int ret = -EINVAL;

	hyp_vm = get_np_pkvm_hyp_vm(handle);
	if (!hyp_vm)
		goto out;

	ret = __pkvm_host_wrprotect_guest(gfn, nr_pages, hyp_vm);
	put_pkvm_hyp_vm(hyp_vm);
out:
	cpu_reg(host_ctxt, 1) = ret;
}

static void handle___pkvm_host_test_clear_young_guest(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(pkvm_handle_t, handle, host_ctxt, 1);
	DECLARE_REG(u64, gfn, host_ctxt, 2);
	DECLARE_REG(u64, nr_pages, host_ctxt, 3);
	DECLARE_REG(bool, mkold, host_ctxt, 4);
	struct pkvm_hyp_vm *hyp_vm;
	int ret = -EINVAL;

	hyp_vm = get_np_pkvm_hyp_vm(handle);
	if (!hyp_vm)
		goto out;

	ret = __pkvm_host_test_clear_young_guest(gfn, nr_pages, mkold, hyp_vm);
	put_pkvm_hyp_vm(hyp_vm);
out:
	cpu_reg(host_ctxt, 1) = ret;
}

static void handle___pkvm_host_mkyoung_guest(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(u64, gfn, host_ctxt, 1);
	struct pkvm_hyp_vcpu *hyp_vcpu;
	int ret = -EINVAL;

	hyp_vcpu = pkvm_get_loaded_hyp_vcpu();
	if (!hyp_vcpu || pkvm_hyp_vcpu_is_protected(hyp_vcpu))
		goto out;

	ret = __pkvm_host_mkyoung_guest(gfn, hyp_vcpu);
out:
	cpu_reg(host_ctxt, 1) =  ret;
}

static void handle___kvm_adjust_pc(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(struct kvm_vcpu *, vcpu, host_ctxt, 1);

	__kvm_adjust_pc(kern_hyp_va(vcpu));
}

static void handle___kvm_flush_vm_context(struct kvm_cpu_context *host_ctxt)
{
	__kvm_flush_vm_context();
}

static void handle___kvm_tlb_flush_vmid_ipa(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(struct kvm_s2_mmu *, mmu, host_ctxt, 1);
	DECLARE_REG(phys_addr_t, ipa, host_ctxt, 2);
	DECLARE_REG(int, level, host_ctxt, 3);

	__kvm_tlb_flush_vmid_ipa(kern_hyp_va(mmu), ipa, level);
}

static void handle___kvm_tlb_flush_vmid_ipa_nsh(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(struct kvm_s2_mmu *, mmu, host_ctxt, 1);
	DECLARE_REG(phys_addr_t, ipa, host_ctxt, 2);
	DECLARE_REG(int, level, host_ctxt, 3);

	__kvm_tlb_flush_vmid_ipa_nsh(kern_hyp_va(mmu), ipa, level);
}

static void
handle___kvm_tlb_flush_vmid_range(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(struct kvm_s2_mmu *, mmu, host_ctxt, 1);
	DECLARE_REG(phys_addr_t, start, host_ctxt, 2);
	DECLARE_REG(unsigned long, pages, host_ctxt, 3);

	__kvm_tlb_flush_vmid_range(kern_hyp_va(mmu), start, pages);
}

static void handle___kvm_tlb_flush_vmid(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(struct kvm_s2_mmu *, mmu, host_ctxt, 1);

	__kvm_tlb_flush_vmid(kern_hyp_va(mmu));
}

static void handle___pkvm_tlb_flush_vmid(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(pkvm_handle_t, handle, host_ctxt, 1);
	struct pkvm_hyp_vm *hyp_vm = get_np_pkvm_hyp_vm(handle);

	if (!hyp_vm)
		return;

	__kvm_tlb_flush_vmid(&hyp_vm->kvm.arch.mmu);
	put_pkvm_hyp_vm(hyp_vm);
}

static void handle___kvm_flush_cpu_context(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(struct kvm_s2_mmu *, mmu, host_ctxt, 1);

	__kvm_flush_cpu_context(kern_hyp_va(mmu));
}

static void handle___kvm_timer_set_cntvoff(struct kvm_cpu_context *host_ctxt)
{
	__kvm_timer_set_cntvoff(cpu_reg(host_ctxt, 1));
}

static void handle___kvm_enable_ssbs(struct kvm_cpu_context *host_ctxt)
{
	u64 tmp;

	tmp = read_sysreg_el2(SYS_SCTLR);
	tmp |= SCTLR_ELx_DSSBS;
	write_sysreg_el2(tmp, SYS_SCTLR);
}

static void handle___vgic_v3_get_gic_config(struct kvm_cpu_context *host_ctxt)
{
	cpu_reg(host_ctxt, 1) = __vgic_v3_get_gic_config();
}

static void handle___vgic_v3_init_lrs(struct kvm_cpu_context *host_ctxt)
{
	__vgic_v3_init_lrs();
}

static void handle___vgic_v3_save_aprs(struct kvm_cpu_context *host_ctxt)
{
	struct pkvm_hyp_vcpu *hyp_vcpu;
	struct kvm_vcpu *host_vcpu;

	host_vcpu = get_host_hyp_vcpus_from_vgic_v3_cpu_if(host_ctxt, 1,
							   &hyp_vcpu);
	if (!host_vcpu)
		return;

	if (unlikely(hyp_vcpu)) {
		struct vgic_v3_cpu_if *hyp_cpu_if, *host_cpu_if;
		int i;

		hyp_cpu_if = &hyp_vcpu->vcpu.arch.vgic_cpu.vgic_v3;
		__vgic_v3_save_aprs(hyp_cpu_if);

		host_cpu_if = &host_vcpu->arch.vgic_cpu.vgic_v3;
		host_cpu_if->vgic_vmcr = hyp_cpu_if->vgic_vmcr;
		for (i = 0; i < ARRAY_SIZE(host_cpu_if->vgic_ap0r); i++) {
			host_cpu_if->vgic_ap0r[i] = hyp_cpu_if->vgic_ap0r[i];
			host_cpu_if->vgic_ap1r[i] = hyp_cpu_if->vgic_ap1r[i];
		}
	} else {
		__vgic_v3_save_aprs(&host_vcpu->arch.vgic_cpu.vgic_v3);
	}
}

static void handle___vgic_v3_restore_vmcr_aprs(struct kvm_cpu_context *host_ctxt)
{
	struct pkvm_hyp_vcpu *hyp_vcpu;
	struct kvm_vcpu *host_vcpu;

	host_vcpu = get_host_hyp_vcpus_from_vgic_v3_cpu_if(host_ctxt, 1,
							   &hyp_vcpu);
	if (!host_vcpu)
		return;

	if (unlikely(hyp_vcpu)) {
		struct vgic_v3_cpu_if *hyp_cpu_if, *host_cpu_if;
		int i;

		hyp_cpu_if = &hyp_vcpu->vcpu.arch.vgic_cpu.vgic_v3;
		host_cpu_if = &host_vcpu->arch.vgic_cpu.vgic_v3;

		hyp_cpu_if->vgic_vmcr = host_cpu_if->vgic_vmcr;
		/* Should be a one-off */
		hyp_cpu_if->vgic_sre = (ICC_SRE_EL1_DIB |
					ICC_SRE_EL1_DFB |
					ICC_SRE_EL1_SRE);
		for (i = 0; i < ARRAY_SIZE(host_cpu_if->vgic_ap0r); i++) {
			hyp_cpu_if->vgic_ap0r[i] = host_cpu_if->vgic_ap0r[i];
			hyp_cpu_if->vgic_ap1r[i] = host_cpu_if->vgic_ap1r[i];
		}

		__vgic_v3_restore_vmcr_aprs(hyp_cpu_if);
	} else {
		__vgic_v3_restore_vmcr_aprs(&host_vcpu->arch.vgic_cpu.vgic_v3);
	}
}

static void handle___pkvm_init(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(phys_addr_t, phys, host_ctxt, 1);
	DECLARE_REG(unsigned long, size, host_ctxt, 2);
	DECLARE_REG(unsigned long *, per_cpu_base, host_ctxt, 3);
	DECLARE_REG(u32, hyp_va_bits, host_ctxt, 4);

	/*
	 * __pkvm_init() will return only if an error occurred, otherwise it
	 * will tail-call in __pkvm_init_finalise() which will have to deal
	 * with the host context directly.
	 */
	cpu_reg(host_ctxt, 1) = __pkvm_init(phys, size, per_cpu_base, hyp_va_bits);
}

static void handle___pkvm_cpu_set_vector(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(enum arm64_hyp_spectre_vector, slot, host_ctxt, 1);

	cpu_reg(host_ctxt, 1) = pkvm_cpu_set_vector(slot);
}

static void handle___pkvm_host_share_hyp(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(u64, pfn, host_ctxt, 1);

	cpu_reg(host_ctxt, 1) = __pkvm_host_share_hyp(pfn);
}

static void handle___pkvm_host_unshare_hyp(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(u64, pfn, host_ctxt, 1);

	cpu_reg(host_ctxt, 1) = __pkvm_host_unshare_hyp(pfn);
}

static void handle___pkvm_create_private_mapping(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(phys_addr_t, phys, host_ctxt, 1);
	DECLARE_REG(size_t, size, host_ctxt, 2);
	DECLARE_REG(enum kvm_pgtable_prot, prot, host_ctxt, 3);

	/*
	 * __pkvm_create_private_mapping() populates a pointer with the
	 * hypervisor start address of the allocation.
	 *
	 * However, handle___pkvm_create_private_mapping() hypercall crosses the
	 * EL1/EL2 boundary so the pointer would not be valid in this context.
	 *
	 * Instead pass the allocation address as the return value (or return
	 * ERR_PTR() on failure).
	 */
	unsigned long haddr;
	int err = __pkvm_create_private_mapping(phys, size, prot, &haddr);

	if (err)
		haddr = (unsigned long)ERR_PTR(err);

	cpu_reg(host_ctxt, 1) = haddr;
}

static void handle___pkvm_prot_finalize(struct kvm_cpu_context *host_ctxt)
{
	cpu_reg(host_ctxt, 1) = __pkvm_prot_finalize();
}

static void handle___pkvm_reserve_vm(struct kvm_cpu_context *host_ctxt)
{
	cpu_reg(host_ctxt, 1) = __pkvm_reserve_vm();
}

static void handle___pkvm_unreserve_vm(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(pkvm_handle_t, handle, host_ctxt, 1);

	__pkvm_unreserve_vm(handle);
}

static void handle___pkvm_init_vm(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(struct kvm *, host_kvm, host_ctxt, 1);
	DECLARE_REG(unsigned long, vm_hva, host_ctxt, 2);
	DECLARE_REG(unsigned long, pgd_hva, host_ctxt, 3);

	host_kvm = kern_hyp_va(host_kvm);
	cpu_reg(host_ctxt, 1) = __pkvm_init_vm(host_kvm, vm_hva, pgd_hva);
}

static void handle___pkvm_init_vcpu(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(pkvm_handle_t, handle, host_ctxt, 1);
	DECLARE_REG(struct kvm_vcpu *, host_vcpu, host_ctxt, 2);
	DECLARE_REG(unsigned long, vcpu_hva, host_ctxt, 3);

	host_vcpu = kern_hyp_va(host_vcpu);
	cpu_reg(host_ctxt, 1) = __pkvm_init_vcpu(handle, host_vcpu, vcpu_hva);
}

static void handle___pkvm_vcpu_in_poison_fault(struct kvm_cpu_context *host_ctxt)
{
	int ret;
	struct pkvm_hyp_vcpu *hyp_vcpu = pkvm_get_loaded_hyp_vcpu();

	ret = hyp_vcpu ? __pkvm_vcpu_in_poison_fault(hyp_vcpu) : -EINVAL;
	cpu_reg(host_ctxt, 1) = ret;
}

static void handle___pkvm_force_reclaim_guest_page(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(phys_addr_t, phys, host_ctxt, 1);

	cpu_reg(host_ctxt, 1) = __pkvm_host_force_reclaim_page_guest(phys);
}

static void handle___pkvm_reclaim_dying_guest_page(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(pkvm_handle_t, handle, host_ctxt, 1);
	DECLARE_REG(u64, gfn, host_ctxt, 2);

	cpu_reg(host_ctxt, 1) = __pkvm_reclaim_dying_guest_page(handle, gfn);
}

static void handle___pkvm_start_teardown_vm(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(pkvm_handle_t, handle, host_ctxt, 1);

	cpu_reg(host_ctxt, 1) = __pkvm_start_teardown_vm(handle);
}

static void handle___pkvm_finalize_teardown_vm(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(pkvm_handle_t, handle, host_ctxt, 1);

	cpu_reg(host_ctxt, 1) = __pkvm_finalize_teardown_vm(handle);
}

static void handle___tracing_load(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(unsigned long, desc_hva, host_ctxt, 1);
	DECLARE_REG(size_t, desc_size, host_ctxt, 2);

	cpu_reg(host_ctxt, 1) = __tracing_load(desc_hva, desc_size);
}

static void handle___tracing_unload(struct kvm_cpu_context *host_ctxt)
{
	__tracing_unload();
}

static void handle___tracing_enable(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(bool, enable, host_ctxt, 1);

	cpu_reg(host_ctxt, 1) = __tracing_enable(enable);
}

static void handle___tracing_swap_reader(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(unsigned int, cpu, host_ctxt, 1);

	cpu_reg(host_ctxt, 1) = __tracing_swap_reader(cpu);
}

static void handle___tracing_update_clock(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(u32, mult, host_ctxt, 1);
	DECLARE_REG(u32, shift, host_ctxt, 2);
	DECLARE_REG(u64, epoch_ns, host_ctxt, 3);
	DECLARE_REG(u64, epoch_cyc, host_ctxt, 4);

	__tracing_update_clock(mult, shift, epoch_ns, epoch_cyc);
}

static void handle___tracing_reset(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(unsigned int, cpu, host_ctxt, 1);

	cpu_reg(host_ctxt, 1) = __tracing_reset(cpu);
}

static void handle___tracing_enable_event(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(unsigned short, id, host_ctxt, 1);
	DECLARE_REG(bool, enable, host_ctxt, 2);

	cpu_reg(host_ctxt, 1) = __tracing_enable_event(id, enable);
}

static void handle___tracing_write_event(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(u64, id, host_ctxt, 1);

	trace_selftest(id);
}

static void handle___vgic_v5_save_apr(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(struct vgic_v5_cpu_if *, cpu_if, host_ctxt, 1);

	__vgic_v5_save_apr(kern_hyp_va(cpu_if));
}

static void handle___vgic_v5_restore_vmcr_apr(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(struct vgic_v5_cpu_if *, cpu_if, host_ctxt, 1);

	__vgic_v5_restore_vmcr_apr(kern_hyp_va(cpu_if));
}

typedef void (*hcall_t)(struct kvm_cpu_context *);

#define HANDLE_FUNC(x)	[__KVM_HOST_SMCCC_FUNC_##x] = (hcall_t)handle_##x

static const hcall_t host_hcall[] = {
	/* ___kvm_hyp_init */
	HANDLE_FUNC(__pkvm_init),
	HANDLE_FUNC(__pkvm_create_private_mapping),
	HANDLE_FUNC(__pkvm_cpu_set_vector),
	HANDLE_FUNC(__kvm_enable_ssbs),
	HANDLE_FUNC(__vgic_v3_init_lrs),
	HANDLE_FUNC(__vgic_v3_get_gic_config),
	HANDLE_FUNC(__pkvm_prot_finalize),

	HANDLE_FUNC(__kvm_adjust_pc),
	HANDLE_FUNC(__kvm_vcpu_run),
	HANDLE_FUNC(__kvm_flush_vm_context),
	HANDLE_FUNC(__kvm_tlb_flush_vmid_ipa),
	HANDLE_FUNC(__kvm_tlb_flush_vmid_ipa_nsh),
	HANDLE_FUNC(__kvm_tlb_flush_vmid),
	HANDLE_FUNC(__kvm_tlb_flush_vmid_range),
	HANDLE_FUNC(__kvm_flush_cpu_context),
	HANDLE_FUNC(__kvm_timer_set_cntvoff),
	HANDLE_FUNC(__tracing_load),
	HANDLE_FUNC(__tracing_unload),
	HANDLE_FUNC(__tracing_enable),
	HANDLE_FUNC(__tracing_swap_reader),
	HANDLE_FUNC(__tracing_update_clock),
	HANDLE_FUNC(__tracing_reset),
	HANDLE_FUNC(__tracing_enable_event),
	HANDLE_FUNC(__tracing_write_event),
	HANDLE_FUNC(__vgic_v3_save_aprs),
	HANDLE_FUNC(__vgic_v3_restore_vmcr_aprs),
	HANDLE_FUNC(__vgic_v5_save_apr),
	HANDLE_FUNC(__vgic_v5_restore_vmcr_apr),

	HANDLE_FUNC(__pkvm_host_share_hyp),
	HANDLE_FUNC(__pkvm_host_unshare_hyp),
	HANDLE_FUNC(__pkvm_host_donate_guest),
	HANDLE_FUNC(__pkvm_host_share_guest),
	HANDLE_FUNC(__pkvm_host_unshare_guest),
	HANDLE_FUNC(__pkvm_host_relax_perms_guest),
	HANDLE_FUNC(__pkvm_host_wrprotect_guest),
	HANDLE_FUNC(__pkvm_host_test_clear_young_guest),
	HANDLE_FUNC(__pkvm_host_mkyoung_guest),
	HANDLE_FUNC(__pkvm_reserve_vm),
	HANDLE_FUNC(__pkvm_unreserve_vm),
	HANDLE_FUNC(__pkvm_init_vm),
	HANDLE_FUNC(__pkvm_init_vcpu),
	HANDLE_FUNC(__pkvm_vcpu_in_poison_fault),
	HANDLE_FUNC(__pkvm_force_reclaim_guest_page),
	HANDLE_FUNC(__pkvm_reclaim_dying_guest_page),
	HANDLE_FUNC(__pkvm_start_teardown_vm),
	HANDLE_FUNC(__pkvm_finalize_teardown_vm),
	HANDLE_FUNC(__pkvm_vcpu_load),
	HANDLE_FUNC(__pkvm_vcpu_put),
	HANDLE_FUNC(__pkvm_vcpu_sync_state),
	HANDLE_FUNC(__pkvm_tlb_flush_vmid),
};

static void handle_host_hcall(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(unsigned long, id, host_ctxt, 0);
	unsigned long hcall_min = 0, hcall_max = __KVM_HOST_SMCCC_FUNC_MAX;
	hcall_t hfn;

	BUILD_BUG_ON(ARRAY_SIZE(host_hcall) != __KVM_HOST_SMCCC_FUNC_MAX);

	/*
	 * If pKVM has been initialised then reject any calls to the
	 * early "privileged" hypercalls. Note that we cannot reject
	 * calls to __pkvm_prot_finalize for two reasons: (1) The static
	 * key used to determine initialisation must be toggled prior to
	 * finalisation and (2) finalisation is performed on a per-CPU
	 * basis. This is all fine, however, since __pkvm_prot_finalize
	 * returns -EPERM after the first call for a given CPU.
	 */
	if (static_branch_unlikely(&kvm_protected_mode_initialized)) {
		hcall_min = __KVM_HOST_SMCCC_FUNC_MIN_PKVM;
	} else {
		hcall_max = __KVM_HOST_SMCCC_FUNC_PKVM_ONLY;
	}

	id &= ~ARM_SMCCC_CALL_HINTS;
	id -= KVM_HOST_SMCCC_ID(0);

	if (unlikely(id < hcall_min || id >= hcall_max))
		goto inval;

	hfn = host_hcall[id];
	if (unlikely(!hfn))
		goto inval;

	cpu_reg(host_ctxt, 0) = SMCCC_RET_SUCCESS;
	hfn(host_ctxt);

	return;
inval:
	cpu_reg(host_ctxt, 0) = SMCCC_RET_NOT_SUPPORTED;
}

static void default_host_smc_handler(struct kvm_cpu_context *host_ctxt)
{
	trace_hyp_exit(host_ctxt, HYP_REASON_SMC);
	__kvm_hyp_host_forward_smc(host_ctxt);
	trace_hyp_enter(host_ctxt, HYP_REASON_SMC);
}

static void handle_host_smc(struct kvm_cpu_context *host_ctxt)
{
	DECLARE_REG(u64, func_id, host_ctxt, 0);
	u64 esr = read_sysreg_el2(SYS_ESR);
	bool handled;

	if (esr & ESR_ELx_xVC_IMM_MASK) {
		cpu_reg(host_ctxt, 0) = SMCCC_RET_NOT_SUPPORTED;
		goto exit_skip_instr;
	}

	func_id &= ~ARM_SMCCC_CALL_HINTS;
	if (upper_32_bits(func_id)) {
		cpu_reg(host_ctxt, 0) = SMCCC_RET_NOT_SUPPORTED;
		goto exit_skip_instr;
	}

	handled = kvm_host_psci_handler(host_ctxt, func_id);
	if (!handled)
		handled = kvm_host_ffa_handler(host_ctxt, func_id);
	if (!handled)
		default_host_smc_handler(host_ctxt);

exit_skip_instr:
	/* SMC was trapped, move ELR past the current PC. */
	kvm_skip_host_instr();
}

void inject_host_exception(u64 esr)
{
	u64 sctlr, spsr_el1, spsr_el2, exc_offset = except_type_sync;
	const u64 spsr_mask = PSR_N_BIT | PSR_Z_BIT | PSR_C_BIT |
			      PSR_V_BIT | PSR_DIT_BIT | PSR_PAN_BIT;

	spsr_el1 = spsr_el2 = read_sysreg_el2(SYS_SPSR);
	switch (spsr_el1 & (PSR_MODE_MASK | PSR_MODE32_BIT)) {
	case PSR_MODE_EL0t:
		exc_offset += LOWER_EL_AArch64_VECTOR;
		break;
	case PSR_MODE_EL0t | PSR_MODE32_BIT:
		exc_offset += LOWER_EL_AArch32_VECTOR;
		break;
	default:
		exc_offset += CURRENT_EL_SP_ELx_VECTOR;
	}

	spsr_el2 &= spsr_mask;
	spsr_el2 |= PSR_D_BIT | PSR_A_BIT | PSR_I_BIT | PSR_F_BIT |
		    PSR_MODE_EL1h;

	sctlr = read_sysreg_el1(SYS_SCTLR);
	if (!(sctlr & SCTLR_EL1_SPAN))
		spsr_el2 |= PSR_PAN_BIT;

	if (sctlr & SCTLR_ELx_DSSBS)
		spsr_el2 |= PSR_SSBS_BIT;

	if (system_supports_mte())
		spsr_el2 |= PSR_TCO_BIT;

	if (esr_fsc_is_translation_fault(esr))
		write_sysreg_el1(read_sysreg_el2(SYS_FAR), SYS_FAR);

	write_sysreg_el1(esr, SYS_ESR);
	write_sysreg_el1(read_sysreg_el2(SYS_ELR), SYS_ELR);
	write_sysreg_el1(spsr_el1, SYS_SPSR);
	write_sysreg_el2(read_sysreg_el1(SYS_VBAR) + exc_offset, SYS_ELR);
	write_sysreg_el2(spsr_el2, SYS_SPSR);
}

static void inject_host_undef64(void)
{
	inject_host_exception((ESR_ELx_EC_UNKNOWN << ESR_ELx_EC_SHIFT) |
			       ESR_ELx_IL);
}

static bool handle_host_mte(u64 esr)
{
	switch (esr_sys64_to_sysreg(esr)) {
	case SYS_RGSR_EL1:
	case SYS_GCR_EL1:
	case SYS_TFSR_EL1:
	case SYS_TFSRE0_EL1:
		/* If we're here for any reason other than MTE, it's a bug. */
		if (read_sysreg(HCR_EL2) & HCR_ATA)
			return false;
		break;
	case SYS_GMID_EL1:
		/* If we're here for any reason other than MTE, it's a bug. */
		if (!(read_sysreg(HCR_EL2) & HCR_TID5))
			return false;
		break;
	default:
		return false;
	}

	inject_host_undef64();
	return true;
}

void handle_trap(struct kvm_cpu_context *host_ctxt)
{
	u64 esr = read_sysreg_el2(SYS_ESR);


	switch (ESR_ELx_EC(esr)) {
	case ESR_ELx_EC_HVC64:
		trace_hyp_enter(host_ctxt, HYP_REASON_HVC);
		handle_host_hcall(host_ctxt);
		break;
	case ESR_ELx_EC_SMC64:
		trace_hyp_enter(host_ctxt, HYP_REASON_SMC);
		handle_host_smc(host_ctxt);
		break;
	case ESR_ELx_EC_IABT_LOW:
	case ESR_ELx_EC_DABT_LOW:
		trace_hyp_enter(host_ctxt, HYP_REASON_HOST_ABORT);
		handle_host_mem_abort(host_ctxt);
		break;
	case ESR_ELx_EC_SYS64:
		trace_hyp_enter(host_ctxt, HYP_REASON_SYS);
		if (handle_host_mte(esr))
			break;
		fallthrough;
	default:
		BUG();
	}

	trace_hyp_exit(host_ctxt, HYP_REASON_ERET_HOST);
}