From c845bb9e8de54be07af0654c6e157d7d6d4e4f22 Mon Sep 17 00:00:00 2001
From: Will Deacon <willdeacon@google.com>
Date: Thu, 7 Jul 2022 15:48:47 +0100
Subject: [PATCH] Revert "ANDROID: KVM: arm64: Remove stale shadow_handle
 field"

This reverts commit 10bb2d3f934ef78eba177ccd291fb8c3da914343.

Bug: 233587962
Signed-off-by: Will Deacon <willdeacon@google.com>
Change-Id: I8cbfe4020bc6349887ccacf8fde477db6b076686
---
 arch/arm64/include/asm/kvm_host.h |  3 +++
 arch/arm64/kvm/hyp/nvhe/pkvm.c    |  1 +
 arch/arm64/kvm/pkvm.c             | 14 +++++++++++++-
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/arch/arm64/include/asm/kvm_host.h b/arch/arm64/include/asm/kvm_host.h
index 20393bd8f73c..0e04f3c2ba4f 100644
--- a/arch/arm64/include/asm/kvm_host.h
+++ b/arch/arm64/include/asm/kvm_host.h
@@ -214,6 +214,9 @@ struct kvm_arch {
 };
 
 struct kvm_protected_vcpu {
+	/* A unique id to the shadow structs in the hyp shadow area. */
+	int shadow_handle;
+
 	/* A pointer to the host's vcpu. */
 	struct kvm_vcpu *host_vcpu;
 
diff --git a/arch/arm64/kvm/hyp/nvhe/pkvm.c b/arch/arm64/kvm/hyp/nvhe/pkvm.c
index de191428cccf..511b5dd5a6fb 100644
--- a/arch/arm64/kvm/hyp/nvhe/pkvm.c
+++ b/arch/arm64/kvm/hyp/nvhe/pkvm.c
@@ -445,6 +445,7 @@ static int init_shadow_structs(struct kvm *kvm, struct kvm_shadow_vm *vm,
 		shadow_state->vm = vm;
 
 		shadow_vcpu->arch.hw_mmu = &vm->arch.mmu;
+		shadow_vcpu->arch.pkvm.shadow_handle = vm->shadow_handle;
 		shadow_vcpu->arch.pkvm.shadow_vm = vm;
 		shadow_vcpu->arch.power_off = true;
 
diff --git a/arch/arm64/kvm/pkvm.c b/arch/arm64/kvm/pkvm.c
index ceabdc9f2ba6..23a6b954657f 100644
--- a/arch/arm64/kvm/pkvm.c
+++ b/arch/arm64/kvm/pkvm.c
@@ -105,6 +105,14 @@ void __init kvm_hyp_reserve(void)
 		 hyp_mem_base);
 }
 
+/*
+ * Updates the state of the host's version of the vcpu state.
+ */
+static void update_vcpu_state(struct kvm_vcpu *vcpu, int shadow_handle)
+{
+	vcpu->arch.pkvm.shadow_handle = shadow_handle;
+}
+
 /*
  * Allocates and donates memory for EL2 shadow structs.
  *
@@ -122,7 +130,7 @@ static int __create_el2_shadow(struct kvm *kvm)
 	void *pgd, *shadow_addr;
 	unsigned long idx;
 	int shadow_handle;
-	int ret;
+	int ret, i;
 
 	if (kvm->created_vcpus < 1)
 		return -EINVAL;
@@ -163,6 +171,10 @@ static int __create_el2_shadow(struct kvm *kvm)
 	/* Store the shadow handle given by hyp for future call reference. */
 	kvm->arch.pkvm.shadow_handle = shadow_handle;
 
+	/* Adjust host's vcpu state as it doesn't control it anymore. */
+	for (i = 0; i < kvm->created_vcpus; i++)
+		update_vcpu_state(kvm->vcpus[i], shadow_handle);
+
 	return 0;
 
 free_shadow: