ANDROID: KVM: arm64: Move addr_is_allowed_memory() check into host callback

Since host stage-2 mappings are created lazily, we cannot rely on the
pte in order to recover the target physical address when checking a
host-initiated memory transition.

Instead, move the addr_is_allowed_memory() check into the host callback
function where it is passed the physical address directly from the
walker.

Bug: 279739439
Signed-off-by: Will Deacon <willdeacon@google.com>
Change-Id: I84bdc43eded79f1f5e5a489dbc0874604491e5c8
(cherry picked from commit b35a06182451f71cc0543cfe36a3f21fad6f6f02)
Signed-off-by: Lee Jones <joneslee@google.com>
diff --git a/arch/arm64/kvm/hyp/nvhe/mem_protect.c b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
index 821b512..1c902f6 100644
--- a/arch/arm64/kvm/hyp/nvhe/mem_protect.c
+++ b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
@@ -792,9 +792,6 @@
 	struct check_walk_data *d = arg;
 	kvm_pte_t pte = *ptep;
 
-	if (kvm_pte_valid(pte) && !addr_is_allowed_memory(kvm_pte_to_phys(pte)))
-		return -EINVAL;
-
 	return d->get_page_state(pte, addr) == d->desired ? 0 : -EPERM;
 }
 
@@ -812,6 +809,9 @@
 
 static enum pkvm_page_state host_get_page_state(kvm_pte_t pte, u64 addr)
 {
+	if (!addr_is_allowed_memory(addr))
+		return PKVM_NOPAGE;
+
 	if (!kvm_pte_valid(pte) && pte)
 		return PKVM_NOPAGE;