From ab4f5f55f78b1bc58b34fe933ce7688e1717e33f Mon Sep 17 00:00:00 2001
From: Will Deacon <willdeacon@google.com>
Date: Thu, 7 Jul 2022 15:51:31 +0100
Subject: [PATCH] Revert "ANDROID: BACKPORT: KVM: arm64: Make hyp stage-1
 refcnt correct on the whole range"

This reverts commit 71a2987234164cc5df69e3881263310d38782a92.

Bug: 233587962
Signed-off-by: Will Deacon <willdeacon@google.com>
Change-Id: Ieb10fd1a85c2f4e496da234f499511cd3c342ed2
---
 arch/arm64/kvm/hyp/nvhe/setup.c | 62 ++++++++++-----------------------
 1 file changed, 19 insertions(+), 43 deletions(-)

diff --git a/arch/arm64/kvm/hyp/nvhe/setup.c b/arch/arm64/kvm/hyp/nvhe/setup.c
index 8f1162c20c48..a0e1b4f33b16 100644
--- a/arch/arm64/kvm/hyp/nvhe/setup.c
+++ b/arch/arm64/kvm/hyp/nvhe/setup.c
@@ -165,11 +165,12 @@ static void hpool_put_page(void *addr)
 	hyp_put_page(&hpool, addr);
 }
 
-static int fix_host_ownership_walker(u64 addr, u64 end, u32 level,
-				     kvm_pte_t *ptep,
-				     enum kvm_pgtable_walk_flags flag,
-				     void * const arg)
+static int finalize_host_mappings_walker(u64 addr, u64 end, u32 level,
+					 kvm_pte_t *ptep,
+					 enum kvm_pgtable_walk_flags flag,
+					 void * const arg)
 {
+	struct kvm_pgtable_mm_ops *mm_ops = arg;
 	enum kvm_pgtable_prot prot;
 	enum pkvm_page_state state;
 	kvm_pte_t pte = *ptep;
@@ -178,6 +179,15 @@ static int fix_host_ownership_walker(u64 addr, u64 end, u32 level,
 	if (!kvm_pte_valid(pte))
 		return 0;
 
+	/*
+	 * Fix-up the refcount for the page-table pages as the early allocator
+	 * was unable to access the hyp_vmemmap and so the buddy allocator has
+	 * initialised the refcount to '1'.
+	 */
+	mm_ops->get_page(ptep);
+	if (flag != KVM_PGTABLE_WALK_LEAF)
+		return 0;
+
 	if (level != (KVM_PGTABLE_MAX_LEVELS - 1))
 		return -EINVAL;
 
@@ -206,30 +216,12 @@ static int fix_host_ownership_walker(u64 addr, u64 end, u32 level,
 	return host_stage2_idmap_locked(phys, PAGE_SIZE, prot);
 }
 
-static int fix_hyp_pgtable_refcnt_walker(u64 addr, u64 end, u32 level,
-					 kvm_pte_t *ptep,
-					 enum kvm_pgtable_walk_flags flag,
-					 void * const arg)
-{
-	struct kvm_pgtable_mm_ops *mm_ops = arg;
-	kvm_pte_t pte = *ptep;
-
-	/*
-	 * Fix-up the refcount for the page-table pages as the early allocator
-	 * was unable to access the hyp_vmemmap and so the buddy allocator has
-	 * initialised the refcount to '1'.
-	 */
-	if (kvm_pte_valid(pte))
-		mm_ops->get_page(ptep);
-
-	return 0;
-}
-
-static int fix_host_ownership(void)
+static int finalize_host_mappings(void)
 {
 	struct kvm_pgtable_walker walker = {
-		.cb	= fix_host_ownership_walker,
-		.flags	= KVM_PGTABLE_WALK_LEAF,
+		.cb	= finalize_host_mappings_walker,
+		.flags	= KVM_PGTABLE_WALK_LEAF | KVM_PGTABLE_WALK_TABLE_POST,
+		.arg	= pkvm_pgtable.mm_ops,
 	};
 	int i, ret;
 
@@ -245,18 +237,6 @@ static int fix_host_ownership(void)
 	return 0;
 }
 
-static int fix_hyp_pgtable_refcnt(void)
-{
-	struct kvm_pgtable_walker walker = {
-		.cb	= fix_hyp_pgtable_refcnt_walker,
-		.flags	= KVM_PGTABLE_WALK_LEAF | KVM_PGTABLE_WALK_TABLE_POST,
-		.arg	= pkvm_pgtable.mm_ops,
-	};
-
-	return kvm_pgtable_walk(&pkvm_pgtable, 0, BIT(pkvm_pgtable.ia_bits),
-				&walker);
-}
-
 static int select_iommu_ops(enum kvm_iommu_driver driver)
 {
 	switch (driver) {
@@ -308,11 +288,7 @@ void __noreturn __pkvm_init_finalise(void)
 			goto out;
 	}
 
-	ret = fix_host_ownership();
-	if (ret)
-		goto out;
-
-	ret = fix_hyp_pgtable_refcnt();
+	ret = finalize_host_mappings();
 	if (ret)
 		goto out;