KVM: arm64: Avoid pointless annotation when mapping host-owned pages

When a page is transitioned to host ownership, we can eagerly map it
into the host stage-2 page-table rather than going via the convoluted
step of a faulting annotation to trigger the mapping.

Call host_stage2_idmap_locked() directly when transitioning a page to
be owned by the host.

Tested-by: Fuad Tabba <tabba@google.com>
Tested-by: Mostafa Saleh <smostafa@google.com>
Signed-off-by: Will Deacon <will@kernel.org>
Link: https://patch.msgid.link/20260330144841.26181-21-will@kernel.org
Signed-off-by: Marc Zyngier <maz@kernel.org>

diff --git a/arch/arm64/kvm/hyp/nvhe/mem_protect.c b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
index 0e57dc1881e02c23e3101732cb710add8d6069aa..bf5102594fc82a0ea23e93aba8b7eed1b749a6a9 100644 (file)
--- a/arch/arm64/kvm/hyp/nvhe/mem_protect.c
+++ b/arch/arm64/kvm/hyp/nvhe/mem_protect.c
@@ -551,23 +551,27 @@ static void __host_update_page_state(phys_addr_t addr, u64 size, enum pkvm_page_
 
 int host_stage2_set_owner_locked(phys_addr_t addr, u64 size, u8 owner_id)
 {
-       int ret;
+       int ret = -EINVAL;
 
        if (!range_is_memory(addr, addr + size))
                return -EPERM;
 
-       ret = host_stage2_try(kvm_pgtable_stage2_set_owner, &host_mmu.pgt,
-                             addr, size, &host_s2_pool, owner_id);
-       if (ret)
-               return ret;
-
-       /* Don't forget to update the vmemmap tracking for the host */
-       if (owner_id == PKVM_ID_HOST)
-               __host_update_page_state(addr, size, PKVM_PAGE_OWNED);
-       else
-               __host_update_page_state(addr, size, PKVM_NOPAGE);
+       switch (owner_id) {
+       case PKVM_ID_HOST:
+               ret = host_stage2_idmap_locked(addr, size, PKVM_HOST_MEM_PROT);
+               if (!ret)
+                       __host_update_page_state(addr, size, PKVM_PAGE_OWNED);
+               break;
+       case PKVM_ID_GUEST:
+       case PKVM_ID_HYP:
+               ret = host_stage2_try(kvm_pgtable_stage2_set_owner, &host_mmu.pgt,
+                                     addr, size, &host_s2_pool, owner_id);
+               if (!ret)
+                       __host_update_page_state(addr, size, PKVM_NOPAGE);
+               break;
+       }
 
-       return 0;
+       return ret;
 }
 
 static bool host_stage2_force_pte_cb(u64 addr, u64 end, enum kvm_pgtable_prot prot)