ANDROID: KVM: arm64: Eagerly map pages when setting host as owner

Rather than annotating and mapping on demand, simply map as soon
as ownership changes.

Bug: 278749606
Bug: 308373293
Bug: 311571169
Bug: 357781595
Change-Id: I966058fcb71ce0c9a151dd61d8f6729cc64e04b0
Signed-off-by: Keir Fraser <keirf@google.com>

arch/arm64/kvm/hyp/nvhe/mem_protect.c

diff --git a/arch/arm64/kvm/hyp/nvhe/mem_protect.c b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
index 9d15e26..38a82ab 100644
--- a/arch/arm64/kvm/hyp/nvhe/mem_protect.c
+++ b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
@@ -587,14 +587,22 @@ static kvm_pte_t kvm_init_invalid_leaf_owner(u8 owner_id)
 int host_stage2_set_owner_locked(phys_addr_t addr, u64 size, u8 owner_id)
 {
 	kvm_pte_t annotation;
+	enum kvm_pgtable_prot prot;
+	int ret;
 
 	if (owner_id > KVM_MAX_OWNER_ID)
 		return -EINVAL;
 
-	annotation = kvm_init_invalid_leaf_owner(owner_id);
-
-	return host_stage2_try(kvm_pgtable_stage2_annotate, &host_mmu.pgt,
-			       addr, size, &host_s2_pool, annotation);
+	if (owner_id == PKVM_ID_HOST) {
+		prot = default_host_prot(addr_is_memory(addr));
+		ret = host_stage2_idmap_locked(addr, size, prot);
+	} else {
+		annotation = kvm_init_invalid_leaf_owner(owner_id);
+		ret = host_stage2_try(kvm_pgtable_stage2_annotate,
+				      &host_mmu.pgt,
+				      addr, size, &host_s2_pool, annotation);
+	}
+	return ret;
 }
 
 static bool host_stage2_force_pte_cb(u64 addr, u64 end, enum kvm_pgtable_prot prot)