[RFC PATCH v2 2/4] arm64: mm: Use fault handler to permit swapper_pg_dir updates

[Ard Biesheuvel <ardb+git@google.com>]

From: Ard Biesheuvel <ardb@kernel.org>

Instead of checking the destination of a page table descriptor update
against the virtual address of swapper_pg_dir in the kernel image, and
taking a dedicated path to perform such updates via a fixmap alias,
rely on a fault handler to do so if the plain store faulted.

This is more efficient for the general case, and makes it more
straight-forward to extend this logic to other page tables that may have
been allocated statically in the .rodata region of the kernel image.

Signed-off-by: Ard Biesheuvel <ardb@kernel.org>
---
 arch/arm64/include/asm/fixmap.h  |  1 +
 arch/arm64/include/asm/pgtable.h | 49 +++++++++-----------
 arch/arm64/mm/mmu.c              | 18 +++----
 3 files changed, 32 insertions(+), 36 deletions(-)

diff --git a/arch/arm64/include/asm/fixmap.h b/arch/arm64/include/asm/fixmap.h
index 9191125738e9..8db791dcb4cb 100644
--- a/arch/arm64/include/asm/fixmap.h
+++ b/arch/arm64/include/asm/fixmap.h
@@ -99,6 +99,7 @@ enum fixed_addresses {
 	FIX_PUD,
 	FIX_P4D,
 	FIX_PGD,
+	FIX_PTVAL,
 
 	__end_of_fixed_addresses
 };
diff --git a/arch/arm64/include/asm/pgtable.h b/arch/arm64/include/asm/pgtable.h
index 6000905a2e86..799cd9d52c28 100644
--- a/arch/arm64/include/asm/pgtable.h
+++ b/arch/arm64/include/asm/pgtable.h
@@ -35,6 +35,7 @@
 #include <asm/cmpxchg.h>
 #include <asm/fixmap.h>
 #include <asm/por.h>
+#include <asm/uaccess.h>
 #include <linux/mmdebug.h>
 #include <linux/mm_types.h>
 #include <linux/sched.h>
@@ -814,7 +815,7 @@ extern pgd_t idmap_pg_dir[];
 extern pgd_t tramp_pg_dir[];
 extern pgd_t reserved_pg_dir[];
 
-extern void set_swapper_pgd(pgd_t *pgdp, pgd_t pgd);
+void try_set_readonly_ptval(ptval_t *ptvalp, ptval_t ptval);
 
 static inline bool in_swapper_pgdir(void *addr)
 {
@@ -824,17 +825,14 @@ static inline bool in_swapper_pgdir(void *addr)
 
 static inline void set_pmd(pmd_t *pmdp, pmd_t pmd)
 {
-#ifdef __PAGETABLE_PMD_FOLDED
-	if (in_swapper_pgdir(pmdp)) {
-		set_swapper_pgd((pgd_t *)pmdp, __pgd(pmd_val(pmd)));
-		return;
-	}
-#endif /* __PAGETABLE_PMD_FOLDED */
-
-	WRITE_ONCE(*pmdp, pmd);
+	__put_kernel_nofault(pmdp, &pmd, pmd_t, fault);
 
 	if (pmd_valid(pmd))
 		queue_pte_barriers();
+	return;
+
+fault:
+	try_set_readonly_ptval(&pmd_val(*pmdp), pmd_val(pmd));
 }
 
 static inline void pmd_clear(pmd_t *pmdp)
@@ -890,15 +888,14 @@ static inline bool pgtable_l4_enabled(void);
 
 static inline void set_pud(pud_t *pudp, pud_t pud)
 {
-	if (!pgtable_l4_enabled() && in_swapper_pgdir(pudp)) {
-		set_swapper_pgd((pgd_t *)pudp, __pgd(pud_val(pud)));
-		return;
-	}
-
-	WRITE_ONCE(*pudp, pud);
+	__put_kernel_nofault(pudp, &pud, pud_t, fault);
 
 	if (pud_valid(pud))
 		queue_pte_barriers();
+	return;
+
+fault:
+	try_set_readonly_ptval(&pud_val(*pudp), pud_val(pud));
 }
 
 static inline void pud_clear(pud_t *pudp)
@@ -971,13 +968,12 @@ static inline bool mm_pud_folded(const struct mm_struct *mm)
 
 static inline void set_p4d(p4d_t *p4dp, p4d_t p4d)
 {
-	if (in_swapper_pgdir(p4dp)) {
-		set_swapper_pgd((pgd_t *)p4dp, __pgd(p4d_val(p4d)));
-		return;
-	}
-
-	WRITE_ONCE(*p4dp, p4d);
+	__put_kernel_nofault(p4dp, &p4d, p4d_t, fault);
 	queue_pte_barriers();
+	return;
+
+fault:
+	try_set_readonly_ptval(&p4d_val(*p4dp), p4d_val(p4d));
 }
 
 static inline void p4d_clear(p4d_t *p4dp)
@@ -1099,13 +1095,12 @@ static inline bool mm_p4d_folded(const struct mm_struct *mm)
 
 static inline void set_pgd(pgd_t *pgdp, pgd_t pgd)
 {
-	if (in_swapper_pgdir(pgdp)) {
-		set_swapper_pgd(pgdp, __pgd(pgd_val(pgd)));
-		return;
-	}
-
-	WRITE_ONCE(*pgdp, pgd);
+	__put_kernel_nofault(pgdp, &pgd, pgd_t, fault);
 	queue_pte_barriers();
+	return;
+
+fault:
+	try_set_readonly_ptval(&pgd_val(*pgdp), pgd_val(pgd));
 }
 
 static inline void pgd_clear(pgd_t *pgdp)
diff --git a/arch/arm64/mm/mmu.c b/arch/arm64/mm/mmu.c
index 9c1aa838e9d5..b97e4bf99ca5 100644
--- a/arch/arm64/mm/mmu.c
+++ b/arch/arm64/mm/mmu.c
@@ -63,34 +63,34 @@ static bool rodata_is_rw __ro_after_init = true;
  */
 long __section(".mmuoff.data.write") __early_cpu_boot_status;
 
-static DEFINE_SPINLOCK(swapper_pgdir_lock);
 static DEFINE_MUTEX(fixmap_lock);
 
-void noinstr set_swapper_pgd(pgd_t *pgdp, pgd_t pgd)
+void noinstr try_set_readonly_ptval(ptval_t *ptvalp, ptval_t ptval)
 {
-	pgd_t *fixmap_pgdp;
+	static DEFINE_SPINLOCK(lock);
+
+	BUG_ON(!in_swapper_pgdir(ptvalp));
 
 	/*
 	 * Don't bother with the fixmap if swapper_pg_dir is still mapped
 	 * writable in the kernel mapping.
 	 */
 	if (rodata_is_rw) {
-		WRITE_ONCE(*pgdp, pgd);
+		WRITE_ONCE(*ptvalp, ptval);
 		dsb(ishst);
 		isb();
 		return;
 	}
 
-	spin_lock(&swapper_pgdir_lock);
-	fixmap_pgdp = pgd_set_fixmap(__pa_symbol(pgdp));
-	WRITE_ONCE(*fixmap_pgdp, pgd);
+	guard(spinlock)(&lock);
+	ptvalp = (ptval_t *)set_fixmap_offset(FIX_PTVAL, __pa_symbol(ptvalp));
+	WRITE_ONCE(*ptvalp, ptval);
 	/*
 	 * We need dsb(ishst) here to ensure the page-table-walker sees
 	 * our new entry before set_p?d() returns. The fixmap's
 	 * flush_tlb_kernel_range() via clear_fixmap() does this for us.
 	 */
-	pgd_clear_fixmap();
-	spin_unlock(&swapper_pgdir_lock);
+	clear_fixmap(FIX_PTVAL);
 }
 
 pgprot_t phys_mem_access_prot(struct file *file, unsigned long pfn,
-- 
2.55.0.887.g758fc8c411-goog