RFC: KVM: arm64: Manage FPSIMD state at EL2 for protected vCPUs

Protected VM's FPSIMD state must not be exposed to the host. Since the
FPSIMD state is switched lazily, hyp must take precautions to prevent
leaks. Do this by trapping FP access to EL2 to lazily save a protected
guest's and restore the hosts's state.

This is a little ahead of its time since it requires knowledge which
vcpus are protected or not (see the TODO).

Signed-off-by: Andrew Scull <ascull@google.com>
diff --git a/arch/arm64/include/asm/kvm_host.h b/arch/arm64/include/asm/kvm_host.h
index d43b23a..447b46e 100644
--- a/arch/arm64/include/asm/kvm_host.h
+++ b/arch/arm64/include/asm/kvm_host.h
@@ -237,6 +237,8 @@
 
 struct kvm_host_data {
 	struct kvm_cpu_context host_ctxt;
+	struct kvm_vcpu* fpsimd_last_vcpu;
+	struct user_fpsimd_state fpsimd_state;
 	struct kvm_pmu_events pmu_events;
 };
 
@@ -275,6 +277,9 @@
  * cannot interfere with the hyp while it is running.
  */
 struct kvm_vcpu_arch_run {
+	/* Whether the vcpu is running as part of a protected vm */
+	bool protected;
+
 	/* Miscellaneous vcpu run state flags */
 	u64 flags;
 };
diff --git a/arch/arm64/kvm/hyp/include/hyp/switch.h b/arch/arm64/kvm/hyp/include/hyp/switch.h
index e13fb6d..c82102a 100644
--- a/arch/arm64/kvm/hyp/include/hyp/switch.h
+++ b/arch/arm64/kvm/hyp/include/hyp/switch.h
@@ -228,12 +228,21 @@
 
 	isb();
 
+	/* A protected vcpu's state might already be in registers. */
+	if (run->protected &&
+	    this_cpu_ptr(&kvm_host_data)->fpsimd_last_vcpu == vcpu &&
+	    vcpu->arch.fpsimd_cpu == smp_processor_id()) {
+		goto out;
+	}
+
 	if (run->flags & KVM_ARM64_RUN_FP_HOST) {
 		/*
 		 * In the SVE case, VHE is assumed: it is enforced by
 		 * Kconfig and kvm_arch_init().
 		 */
-		if (sve_host) {
+		if (run->protected) {
+			__fpsimd_save_state(&this_cpu_ptr(&kvm_host_data)->fpsimd_state);
+		} else if (sve_host) {
 			struct thread_struct *thread = container_of(
 				vcpu->arch.host_fpsimd_state,
 				struct thread_struct, uw.fpsimd_state);
@@ -260,6 +269,12 @@
 	if (!(read_sysreg(hcr_el2) & HCR_RW))
 		write_sysreg(__vcpu_sys_reg(vcpu, FPEXC32_EL2), fpexc32_el2);
 
+	if (run->protected) {
+		this_cpu_ptr(&kvm_host_data)->fpsimd_last_vcpu = vcpu;
+		vcpu->arch.fpsimd_cpu = smp_processor_id();
+	}
+
+out:
 	run->flags |= KVM_ARM64_RUN_FP_ENABLED;
 
 	return true;
diff --git a/arch/arm64/kvm/hyp/nvhe/hyp-main.c b/arch/arm64/kvm/hyp/nvhe/hyp-main.c
index a906f9e..361d33c 100644
--- a/arch/arm64/kvm/hyp/nvhe/hyp-main.c
+++ b/arch/arm64/kvm/hyp/nvhe/hyp-main.c
@@ -169,6 +169,22 @@
 	kvm_skip_host_instr();
 }
 
+static void handle_host_fpsimd(struct kvm_cpu_context *host_ctxt)
+{
+	/*
+	 * An FPSIMD trap from the host means the host's state has been saved
+	 * by hyp and needs to be restored.
+	 */
+	if (this_cpu_ptr(&kvm_host_data)->fpsimd_last_vcpu == NULL)
+		goto out;
+
+	this_cpu_ptr(&kvm_host_data)->fpsimd_last_vcpu = NULL;
+	__fpsimd_restore_state(&this_cpu_ptr(&kvm_host_data)->fpsimd_state);
+
+out:
+	write_sysreg(read_sysreg(cptr_el2) & ~(u64)CPTR_EL2_TFP, cptr_el2);
+}
+
 void handle_trap(struct kvm_cpu_context *host_ctxt)
 {
 	u64 esr = read_sysreg_el2(SYS_ESR);
@@ -180,6 +196,9 @@
 	case ESR_ELx_EC_SMC64:
 		handle_host_smc(host_ctxt);
 		break;
+	case ESR_ELx_EC_FP_ASIMD:
+		handle_host_fpsimd(host_ctxt);
+		break;
 	default:
 		hyp_panic();
 	}
diff --git a/arch/arm64/kvm/hyp/nvhe/switch.c b/arch/arm64/kvm/hyp/nvhe/switch.c
index a0fbaf0..bb34712 100644
--- a/arch/arm64/kvm/hyp/nvhe/switch.c
+++ b/arch/arm64/kvm/hyp/nvhe/switch.c
@@ -32,6 +32,8 @@
 DEFINE_PER_CPU(struct kvm_cpu_context, kvm_hyp_ctxt);
 DEFINE_PER_CPU(unsigned long, kvm_hyp_vector);
 
+static DEFINE_PER_CPU(struct kvm_vcpu_arch_run *, kvm_hyp_vcpu_run);
+
 static void __activate_traps(struct kvm_vcpu *vcpu, struct kvm_vcpu_arch_run *run)
 {
 	u64 val;
@@ -64,10 +66,12 @@
 	}
 }
 
-static void __deactivate_traps(struct kvm_vcpu *vcpu)
+static void __deactivate_traps(struct kvm_vcpu *vcpu, struct kvm_vcpu_arch_run *run)
 {
 	extern char __kvm_hyp_host_vector[];
 	u64 mdcr_el2;
+	u64 hcr_el2;
+	u64 cptr_el2;
 
 	___deactivate_traps(vcpu);
 
@@ -95,12 +99,18 @@
 	mdcr_el2 &= MDCR_EL2_HPMN_MASK;
 	mdcr_el2 |= MDCR_EL2_E2PB_MASK << MDCR_EL2_E2PB_SHIFT;
 
-	write_sysreg(mdcr_el2, mdcr_el2);
 	if (is_protected_kvm_enabled())
-		write_sysreg(HCR_HOST_NVHE_PROTECTED_FLAGS, hcr_el2);
+		hcr_el2 = HCR_HOST_NVHE_PROTECTED_FLAGS;
 	else
-		write_sysreg(HCR_HOST_NVHE_FLAGS, hcr_el2);
-	write_sysreg(CPTR_EL2_DEFAULT, cptr_el2);
+		hcr_el2 = HCR_HOST_NVHE_FLAGS;
+
+	cptr_el2 = CPTR_EL2_DEFAULT;
+	if (run->protected)
+		cptr_el2 |= CPTR_EL2_TFP;
+
+	write_sysreg(mdcr_el2, mdcr_el2);
+	write_sysreg(hcr_el2, hcr_el2);
+	write_sysreg(cptr_el2, cptr_el2);
 	write_sysreg(__kvm_hyp_host_vector, vbar_el2);
 }
 
@@ -172,11 +182,36 @@
 
 	/* Clear host state to make misuse apparent. */
 	vcpu->arch.run.flags = 0;
+
+	if (run->protected) {
+		/*
+		 * For protected vCPUs, always initially disable FPSIMD so we
+		 * can avoid saving the state if it isn't used, but if it is
+		 * used, only save the state for the host if the host state is
+		 * loaded.
+		 */
+		run->flags &= ~(KVM_ARM64_RUN_FP_ENABLED |
+				KVM_ARM64_RUN_FP_HOST);
+		if (this_cpu_ptr(&kvm_host_data)->fpsimd_last_vcpu == NULL)
+			run->flags |= KVM_ARM64_RUN_FP_HOST;
+	} else {
+		/*
+		 * For non-protecetd vCPUs on a system that can also host
+		 * protected vCPUs, ensure protected vCPU FPSIMD state isn't
+		 * used by another vCPU or saved as the host state.
+		 */
+		if (this_cpu_ptr(&kvm_host_data)->fpsimd_last_vcpu != NULL)
+			run->flags &= ~(KVM_ARM64_RUN_FP_ENABLED |
+					KVM_ARM64_RUN_FP_HOST);
+	}
 }
 
 /* Sanitize the run state before writing it back to the host. */
 void __sync_vcpu_after_run(struct kvm_vcpu *vcpu, struct kvm_vcpu_arch_run *run)
 {
+	if (run->protected)
+		return;
+
 	vcpu->arch.run.flags = run->flags;
 }
 
@@ -203,10 +238,12 @@
 
 	if (is_protected_kvm_enabled()) {
 		run = &protected_run;
+		/* TODO: safely check vcpu and set run->protected accordingly. */
 		__sync_vcpu_before_run(vcpu, run);
 	} else {
 		run = &vcpu->arch.run;
 	}
+	__this_cpu_write(kvm_hyp_vcpu_run, run);
 
 	host_ctxt = &this_cpu_ptr(&kvm_host_data)->host_ctxt;
 	host_ctxt->__hyp_running_vcpu = vcpu;
@@ -249,13 +286,16 @@
 	__timer_disable_traps(vcpu);
 	__hyp_vgic_save_state(vcpu);
 
-	__deactivate_traps(vcpu);
+	__deactivate_traps(vcpu, run);
 	__load_host_stage2();
 
 	__sysreg_restore_state_nvhe(host_ctxt);
 
-	if (run->flags & KVM_ARM64_RUN_FP_ENABLED)
+	if (run->flags & KVM_ARM64_RUN_FP_ENABLED) {
 		__fpsimd_save_fpexc32(vcpu);
+		if (run->protected)
+			__fpsimd_save_state(&vcpu->arch.ctxt.fp_regs);
+	}
 
 	/*
 	 * This must come after restoring the host sysregs, since a non-VHE
@@ -284,15 +324,17 @@
 	u64 elr = read_sysreg_el2(SYS_ELR);
 	u64 par = read_sysreg_par();
 	bool restore_host = true;
+	struct kvm_vcpu_arch_run *run;
 	struct kvm_cpu_context *host_ctxt;
 	struct kvm_vcpu *vcpu;
 
+	run = __this_cpu_read(kvm_hyp_vcpu_run);
 	host_ctxt = &this_cpu_ptr(&kvm_host_data)->host_ctxt;
 	vcpu = host_ctxt->__hyp_running_vcpu;
 
 	if (vcpu) {
 		__timer_disable_traps(vcpu);
-		__deactivate_traps(vcpu);
+		__deactivate_traps(vcpu, run);
 		__load_host_stage2();
 		__sysreg_restore_state_nvhe(host_ctxt);
 	}