ANDROID: KVM: arm64: pkvm: Don't zero shared pages

The pKVM hypervisor currently zeroes all the pages mapped into guests
when tearing them down for confidentiality reasons. However, for pages
that are shared with the host this is unecessary at best as the content
of memory is already visible. This is particularly bad for non-protected
guests as all their memory is shared with the host by definition.

Add a new flag to distingish pages that solely need to be updated from
an ownership perspective and those that need to be zeroed.

NOTE: We should probably overhaul the teardown procedure at some point
to avoid the proliferation of those flags, but that would require
significant changes so we might not want that in Android 13.

Bug: 223678931
Change-Id: Icefc85a0bdcdf9958e9eb6871c794f68b06a007f
Signed-off-by: Quentin Perret <qperret@google.com>
diff --git a/arch/arm64/kvm/hyp/include/nvhe/memory.h b/arch/arm64/kvm/hyp/include/nvhe/memory.h
index 2723c6ce..caa5b9a 100644
--- a/arch/arm64/kvm/hyp/include/nvhe/memory.h
+++ b/arch/arm64/kvm/hyp/include/nvhe/memory.h
@@ -12,6 +12,7 @@
  * page-table lock due to the lack of atomics at EL2.
  */
 #define HOST_PAGE_NEED_POISONING	BIT(0)
+#define HOST_PAGE_PENDING_RECLAIM	BIT(1)
 
 struct hyp_page {
 	unsigned short refcount;
diff --git a/arch/arm64/kvm/hyp/nvhe/mem_protect.c b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
index 5fb6b4d..3ad4f16 100644
--- a/arch/arm64/kvm/hyp/nvhe/mem_protect.c
+++ b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
@@ -291,7 +291,17 @@
 		return 0;
 
 	page = hyp_phys_to_page(kvm_pte_to_phys(pte));
-	page->flags |= HOST_PAGE_NEED_POISONING;
+	switch (pkvm_getstate(kvm_pgtable_stage2_pte_prot(pte))) {
+	case PKVM_PAGE_OWNED:
+		page->flags |= HOST_PAGE_NEED_POISONING;
+		fallthrough;
+	case PKVM_PAGE_SHARED_BORROWED:
+	case PKVM_PAGE_SHARED_OWNED:
+		page->flags |= HOST_PAGE_PENDING_RECLAIM;
+		break;
+	default:
+		return -EPERM;
+	}
 
 	return 0;
 }
@@ -1748,16 +1758,23 @@
 		goto unlock;
 
 	page = hyp_phys_to_page(addr);
+	if (!(page->flags & HOST_PAGE_PENDING_RECLAIM)) {
+		ret = -EPERM;
+		goto unlock;
+	}
+
 	if (page->flags & HOST_PAGE_NEED_POISONING) {
 		ret = hyp_zero_page(addr);
 		if (ret)
 			goto unlock;
 		page->flags &= ~HOST_PAGE_NEED_POISONING;
-		ret = host_stage2_set_owner_locked(addr, PAGE_SIZE, pkvm_host_id);
-	} else {
-		ret = -EPERM;
 	}
 
+	ret = host_stage2_set_owner_locked(addr, PAGE_SIZE, pkvm_host_id);
+	if (ret)
+		goto unlock;
+	page->flags &= ~HOST_PAGE_PENDING_RECLAIM;
+
 unlock:
 	host_unlock_component();