ANDROID: KVM: arm64: Fix missing poison for huge-mapping

The entire huge stage-2 mapping must be poisonned when reclaiming.

Bug: 278749606
Bug: 409228797
Change-Id: Idd103a80be3f6c110b5d6605c213c8393a57535a
Signed-off-by: Vincent Donnefort <vdonnefort@google.com>
(cherry picked from commit c56c645543cbb2d40f24125da142d26a754b9a19)
Signed-off-by: Lee Jones <joneslee@google.com>

arch/arm64/kvm/hyp/include/nvhe/mm.h

diff --git a/arch/arm64/kvm/hyp/include/nvhe/mm.h b/arch/arm64/kvm/hyp/include/nvhe/mm.h
index f0d7262..66bc4b1 100644
--- a/arch/arm64/kvm/hyp/include/nvhe/mm.h
+++ b/arch/arm64/kvm/hyp/include/nvhe/mm.h
@@ -19,7 +19,7 @@ void *hyp_fixmap_map(phys_addr_t phys);
 void hyp_fixmap_unmap(void);
 void *hyp_fixblock_map(phys_addr_t phys);
 void hyp_fixblock_unmap(void);
-void hyp_poison_page(phys_addr_t phys);
+void hyp_poison_page(phys_addr_t phys, size_t page_size);
 
 int hyp_create_idmap(u32 hyp_va_bits);
 int hyp_map_vectors(void);

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 442569c..9c53190 100644
--- a/arch/arm64/kvm/hyp/nvhe/mem_protect.c
+++ b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
@@ -401,7 +401,7 @@ static int relinquish_walker(const struct kvm_pgtable_visit_ctx *ctx,
 	phys += ctx->addr - addr;
 
 	if (state == PKVM_PAGE_OWNED) {
-		hyp_poison_page(phys);
+		hyp_poison_page(phys, PAGE_SIZE);
 		psci_mem_protect_dec(1);
 	}
 
@@ -2770,20 +2770,29 @@ int __pkvm_host_donate_guest(struct pkvm_hyp_vcpu *vcpu, u64 pfn, u64 gfn,
 	return ret;
 }
 
-void hyp_poison_page(phys_addr_t phys)
+void hyp_poison_page(phys_addr_t phys, size_t size)
 {
-	void *addr = hyp_fixmap_map(phys);
+	WARN_ON(!PAGE_ALIGNED(size));
 
-	memset(addr, 0, PAGE_SIZE);
-	/*
-	 * Prefer kvm_flush_dcache_to_poc() over __clean_dcache_guest_page()
-	 * here as the latter may elide the CMO under the assumption that FWB
-	 * will be enabled on CPUs that support it. This is incorrect for the
-	 * host stage-2 and would otherwise lead to a malicious host potentially
-	 * being able to read the contents of newly reclaimed guest pages.
-	 */
-	kvm_flush_dcache_to_poc(addr, PAGE_SIZE);
-	hyp_fixmap_unmap();
+	while (size) {
+		size_t __size = size == PMD_SIZE ? size : PAGE_SIZE;
+		void *addr = __fixmap_guest_page(__hyp_va(phys), &__size);
+
+		memset(addr, 0, __size);
+
+		/*
+		 * Prefer kvm_flush_dcache_to_poc() over __clean_dcache_guest_page()
+		 * here as the latter may elide the CMO under the assumption that FWB
+		 * will be enabled on CPUs that support it. This is incorrect for the
+		 * host stage-2 and would otherwise lead to a malicious host potentially
+		 * being able to read the contents of newly reclaimed guest pages.
+		 */
+		kvm_flush_dcache_to_poc(addr, __size);
+		__fixunmap_guest_page(__size);
+
+		size -= __size;
+		phys += __size;
+	}
 }
 
 void destroy_hyp_vm_pgt(struct pkvm_hyp_vm *vm)
@@ -2825,7 +2834,7 @@ int __pkvm_host_reclaim_page(struct pkvm_hyp_vm *vm, u64 pfn, u64 ipa, u8 order)
 	switch((int)guest_get_page_state(pte, ipa)) {
 	case PKVM_PAGE_OWNED:
 		WARN_ON(__host_check_page_state_range(phys, page_size, PKVM_NOPAGE));
-		hyp_poison_page(phys);
+		hyp_poison_page(phys, page_size);
 		psci_mem_protect_dec(order);
 		break;
 	case PKVM_PAGE_SHARED_BORROWED:

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

diff --git a/arch/arm64/kvm/hyp/nvhe/pkvm.c b/arch/arm64/kvm/hyp/nvhe/pkvm.c
index f7b6548..caf7c5e 100644
--- a/arch/arm64/kvm/hyp/nvhe/pkvm.c
+++ b/arch/arm64/kvm/hyp/nvhe/pkvm.c
@@ -1086,7 +1086,7 @@ void pkvm_poison_pvmfw_pages(void)
 	phys_addr_t addr = pvmfw_base;
 
 	while (npages--) {
-		hyp_poison_page(addr);
+		hyp_poison_page(addr, PAGE_SIZE);
 		addr += PAGE_SIZE;
 	}
 }