KVM: arm64: Only drop references on empty tables in stage2_free_walker

A subsequent change to the way KVM frees stage-2s will invoke the free
walker on sub-ranges of the VM's IPA space, meaning there's potential
for only partially visiting a table's PTEs.

Split the leaf and table visitors and only drop references on a table
when the page count reaches 1, implying there are no valid PTEs that
need to be visited. Invalidate the table PTE to avoid traversing the
stale reference.

Link: https://msgid.link/20251113052452.975081-2-rananta@google.com
Signed-off-by: Oliver Upton <oupton@kernel.org>

diff --git a/arch/arm64/kvm/hyp/pgtable.c b/arch/arm64/kvm/hyp/pgtable.c
index c351b4abd5dbfbcbe738ba0f9efc41a1c99ecc05..6d6a23f7dedb6933830dd791c1a740cacef32b15 100644 (file)
--- a/arch/arm64/kvm/hyp/pgtable.c
+++ b/arch/arm64/kvm/hyp/pgtable.c
@@ -1535,20 +1535,46 @@ size_t kvm_pgtable_stage2_pgd_size(u64 vtcr)
        return kvm_pgd_pages(ia_bits, start_level) * PAGE_SIZE;
 }
 
-static int stage2_free_walker(const struct kvm_pgtable_visit_ctx *ctx,
-                             enum kvm_pgtable_walk_flags visit)
+static int stage2_free_leaf(const struct kvm_pgtable_visit_ctx *ctx)
 {
        struct kvm_pgtable_mm_ops *mm_ops = ctx->mm_ops;
 
-       if (!stage2_pte_is_counted(ctx->old))
+       mm_ops->put_page(ctx->ptep);
+       return 0;
+}
+
+static int stage2_free_table_post(const struct kvm_pgtable_visit_ctx *ctx)
+{
+       struct kvm_pgtable_mm_ops *mm_ops = ctx->mm_ops;
+       kvm_pte_t *childp = kvm_pte_follow(ctx->old, mm_ops);
+
+       if (mm_ops->page_count(childp) != 1)
                return 0;
 
+       /*
+        * Drop references and clear the now stale PTE to avoid rewalking the
+        * freed page table.
+        */
        mm_ops->put_page(ctx->ptep);
+       mm_ops->put_page(childp);
+       kvm_clear_pte(ctx->ptep);
+       return 0;
+}
 
-       if (kvm_pte_table(ctx->old, ctx->level))
-               mm_ops->put_page(kvm_pte_follow(ctx->old, mm_ops));
+static int stage2_free_walker(const struct kvm_pgtable_visit_ctx *ctx,
+                             enum kvm_pgtable_walk_flags visit)
+{
+       if (!stage2_pte_is_counted(ctx->old))
+               return 0;
 
-       return 0;
+       switch (visit) {
+       case KVM_PGTABLE_WALK_LEAF:
+               return stage2_free_leaf(ctx);
+       case KVM_PGTABLE_WALK_TABLE_POST:
+               return stage2_free_table_post(ctx);
+       default:
+               return -EINVAL;
+       }
 }
 
 void kvm_pgtable_stage2_destroy(struct kvm_pgtable *pgt)