arm64: Enable vmalloc-huge with ptdump

Our goal is to move towards enabling vmalloc-huge by default on arm64 so as to reduce TLB pressure. Therefore, we need a way to analyze the portion of block mappings in vmalloc space we can get on a production system; this can be done through ptdump, but currently we disable vmalloc-huge if CONFIG_PTDUMP_DEBUGFS is on. The reason is that lazy freeing of kernel pagetables via vmap_try_huge_pxd() may race with ptdump, so ptdump may dereference a bogus address.

To solve this, we need to synchronize ptdump_walk_pgd() with pud_free_pmd_page() and pmd_free_pte_page().

Since this race is very unlikely to happen in practice, we do not want to penalize other code paths taking the init_mm mmap_lock. Therefore, we use static keys. ptdump will enable the static key - upon observing that, the vmalloc table freeing path will get patched in with an mmap_read_lock/unlock sequence. A code comment explains in detail, how a combination of acquire sematics of static_branch_enable() and the barriers in __flush_tlb_kernel_pgtable() ensures that ptdump will never get a hold on the address of a freed PMD or PTE table.

For an (almost) rigorous proof of correctness, one may look at:
https://lore.kernel.org/all/[email protected]/

---
Rebased on 6.16-rc7.

mm-selfests pass. No issues were observed while parallelly running test_vmalloc.sh (which stresses the vmalloc subsystem) and dumping the kernel pagetable through sysfs in a loop.

v4->v5:
- The arch_vmap_pxd_supported() changes were dropped by mistake in between the revisions, add them back (Anshuman)
- Rewrite cover letter, drop stray change, add arm64_ptdump_walk_pgd() helper, rename ptdump_lock_key -> arm64_ptdump_lock_key, add comment above __pmd_free_pte_page() to explain when the lock won't be taken (Anshuman)
- Rewrite the comment explaining the synchronization logic (Catalin)

v3->v4:
- Lock-unlock immediately
- Simplify includes

v2->v3:
- Use static key mechanism

v1->v2:
- Take lock only when CONFIG_PTDUMP_DEBUGFS is on
- In case of pud_free_pmd_page(), isolate the PMD table to avoid taking the lock 512 times again via pmd_free_pte_page()

 arch/arm64/include/asm/ptdump.h  |  2 +
 arch/arm64/include/asm/vmalloc.h |  6 +--
 arch/arm64/mm/mmu.c              | 86 ++++++++++++++++++++++++++++++--
 arch/arm64/mm/ptdump.c           | 12 ++++-
 4 files changed, 96 insertions(+), 10 deletions(-)

index fded5358641f..baff24004459 100644
@@ -7,6 +7,8 @@
 
 #include <linux/ptdump.h>
 
+DECLARE_STATIC_KEY_FALSE(arm64_ptdump_lock_key);
+
 #ifdef CONFIG_PTDUMP
 
 #include <linux/mm_types.h>

index 12f534e8f3ed..e835fd437ae0 100644
@@ -12,15 +12,13 @@ static inline bool arch_vmap_pud_supported(pgprot_t prot)
 	/*
 	 * SW table walks can't handle removal of intermediate entries.
 	 */
-	return pud_sect_supported() &&
-	       !IS_ENABLED(CONFIG_PTDUMP_DEBUGFS);
+	return pud_sect_supported();
 }
 
 #define arch_vmap_pmd_supported arch_vmap_pmd_supported
 static inline bool arch_vmap_pmd_supported(pgprot_t prot)
 {
-	/* See arch_vmap_pud_supported() */
-	return !IS_ENABLED(CONFIG_PTDUMP_DEBUGFS);
+	return true;
 }
 
 #define arch_vmap_pte_range_map_size arch_vmap_pte_range_map_size

index 421a5de806c6..507792fa22de 100644
@@ -283,6 +283,14 @@ void note_page_flush(struct ptdump_state *pt_st)
 	note_page(pt_st, 0, -1, pte_val(pte_zero));
 }
 
+static void arm64_ptdump_walk_pgd(struct ptdump_state *st,
+				  struct mm_struct *mm, pgd_t *pgd)
+{
+	static_branch_enable(&arm64_ptdump_lock_key);
+	ptdump_walk_pgd(&st.ptdump, &init_mm, NULL);
+	static_branch_disable(&arm64_ptdump_lock_key);
+}
+
 void ptdump_walk(struct seq_file *s, struct ptdump_info *info)
 {
 	unsigned long end = ~0UL;

@@ -311,7 +319,7 @@ void ptdump_walk(struct seq_file *s, struct ptdump_info *info)
 		}
 	};
 
-	ptdump_walk_pgd(&st.ptdump, info->mm, NULL);
+	arm64_ptdump_walk_pgd(&st.ptdump, info->mm, NULL);
 }
 
 static void __init ptdump_initialize(void)

@@ -353,7 +361,7 @@ bool ptdump_check_wx(void)
 		}
 	};
 
-	ptdump_walk_pgd(&st.ptdump, &init_mm, NULL);
+	arm64_ptdump_walk_pgd(&st.ptdump, info->mm, NULL);
 
 	if (st.wx_pages || st.uxn_pages) {
 		pr_warn("Checked W+X mappings: FAILED, %lu W+X pages found, %lu non-UXN pages found\n",

2.30.2