[PATCH RFC v2 04/12] mm/hugetlb: Add pgtable walker lock

[Peter Xu]

huge_pte_offset() is potentially a pgtable walker, looking up pte_t* for a
hugetlb address.

Normally, it's always safe to walk a generic pgtable as long as we're with
the mmap lock held for either read or write, because that guarantees the
pgtable pages will always be valid during the process.

But it's not true for hugetlbfs, especially shared: hugetlbfs can have its
pgtable freed by pmd unsharing, it means that even with mmap lock held for
current mm, the PMD pgtable page can still go away from under us if pmd
unsharing is possible during the walk.

So we have three ways to make it safe:

  (1) If the mapping is private we're not prone to pmd sharing or
      unsharing, so it's okay.

  (2) If we're with the hugetlb vma lock held for either read/write, it's
      okay too because pmd unshare cannot happen at all.

  (3) Otherwise we may need the pgtable walker lock

The pgtable walker is implemented in different ways depending on the
config:

  - if !ARCH_WANT_HUGE_PMD_SHARE:      it's no-op
  - else if MMU_GATHER_RCU_TABLE_FREE: it should be rcu_read_lock()
  - else:                              it should be local_irq_disable()

A more efficient way to take the lock is only taking them with valid vmas
that are possible to have pmd sharing (aka, want_pmd_share() returns true),
but since the lock is mostly lightweighted (only riscv will use irq
disable, x86 & arm will use rcu) just take them unconditionally for now.

Document all these explicitly for huge_pte_offset() too, because it's not
that obvious.

Signed-off-by: Peter Xu <peterx@xxxxxxxxxx>
---
 include/linux/hugetlb.h | 99 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 99 insertions(+)

diff --git a/include/linux/hugetlb.h b/include/linux/hugetlb.h
index 551834cd5299..8f85ad0d5bdb 100644
--- a/include/linux/hugetlb.h
+++ b/include/linux/hugetlb.h
@@ -12,6 +12,8 @@
 #include <linux/pgtable.h>
 #include <linux/gfp.h>
 #include <linux/userfaultfd_k.h>
+#include <linux/irqflags.h>
+#include <linux/rcupdate.h>
 
 struct ctl_table;
 struct user_struct;
@@ -24,6 +26,71 @@ typedef struct { unsigned long pd; } hugepd_t;
 #define __hugepd(x) ((hugepd_t) { (x) })
 #endif
 
+/**
+ * @hugetlb_walker_[un]lock(): protects software hugetlb pgtable walkers
+ *
+ * The hugetlb walker lock needs to be taken before huge_pte_offset() and
+ * needs to be released after finishing using the pte_t* returned.  It's
+ * used to guarantee we won't access a freed pgtable page.  Normally in
+ * this case we already have the mmap lock held for read.
+ *
+ * When holding the hugetlb walker lock, the thread cannot sleep, nor can
+ * it already in irq context (just to simplify unlock with no-save for
+ * !RCU_TABLE_TREE).  Before the thread yields itself, it should release
+ * the pgtable lock.  After the lock released, pte_t* can become invalid
+ * anytime so it cannot be accessed anymore.
+ *
+ * The only reason for the lock is to protect concurrent pmd unsharing
+ * followed by e.g. munmap() to drop the pgtable page.  For no-pmd-sharing
+ * archs, it's no-op because it's always safe to access hugetlb pgtables
+ * just like generic pgtables.
+ */
+#if !defined(CONFIG_ARCH_WANT_HUGE_PMD_SHARE)
+static inline void hugetlb_walker_lock(void)
+{
+}
+static inline void hugetlb_walker_unlock(void)
+{
+}
+static inline bool __hugetlb_walker_locked(void)
+{
+	return true;
+}
+#elif defined(CONFIG_MMU_GATHER_RCU_TABLE_FREE)
+static inline void hugetlb_walker_lock(void)
+{
+	rcu_read_lock();
+}
+static inline void hugetlb_walker_unlock(void)
+{
+	rcu_read_unlock();
+}
+static inline bool __hugetlb_walker_locked(void)
+{
+	return rcu_read_lock_held();
+}
+#else
+static inline void hugetlb_walker_lock(void)
+{
+	WARN_ON_ONCE(irqs_disabled());
+	local_irq_disable();
+}
+static inline void hugetlb_walker_unlock(void)
+{
+	local_irq_enable();
+}
+static inline bool __hugetlb_walker_locked(void)
+{
+	return irqs_disabled();
+}
+#endif
+
+#ifdef CONFIG_LOCKDEP
+#define  hugetlb_walker_locked()  __hugetlb_walker_locked()
+#else
+#define  hugetlb_walker_locked()  true
+#endif
+
 #ifdef CONFIG_HUGETLB_PAGE
 
 #include <linux/mempolicy.h>
@@ -192,6 +259,38 @@ extern struct list_head huge_boot_pages;
 
 pte_t *huge_pte_alloc(struct mm_struct *mm, struct vm_area_struct *vma,
 			unsigned long addr, unsigned long sz);
+/*
+ * huge_pte_offset(): Walk the hugetlb pgtable until the last level PTE.
+ * Returns the pte_t* if found, or NULL if the address is not mapped.
+ *
+ * NOTE: since this function will walk all the pgtable pages (including not
+ * only high-level pgtable page, but also PUD entry that can be unshared
+ * concurrently for VM_SHARED), the caller of this function should be
+ * responsible of its thread safety.  One can follow this rule:
+ *
+ *  (1) For private mappings: pmd unsharing is not possible, so it'll
+ *      always be safe if we're with the mmap sem for either read or write.
+ *      This is normally always the case, IOW we don't need to do anything
+ *      special.
+ *
+ *  (2) For shared mappings: pmd unsharing is possible (so the PUD-ranged
+ *      pgtable page can go away from under us!  It can be done by a pmd
+ *      unshare with a follow up munmap() on the other process), then we
+ *      need either:
+ *
+ *     (2.1) hugetlb vma lock read or write held, to make sure pmd unshare
+ *           won't happen upon the range (it also makes sure the pte_t we
+ *           read is the right and stable one), or,
+ *
+ *     (2.2) pgtable walker lock, to make sure even pmd unsharing happened,
+ *           the old shared PUD page won't get freed from under us, so even
+ *           the pteval can be obsolete, at least it's still always safe to
+ *           access the pgtable page (e.g., de-referencing pte_t* would not
+ *           cause use-after-free).
+ *
+ * PS: from the regard of (2.2), it's the same logic of fast-gup being safe
+ * for generic mm when walking the pgtables even without mmap lock.
+ */
 pte_t *huge_pte_offset(struct mm_struct *mm,
 		       unsigned long addr, unsigned long sz);
 unsigned long hugetlb_mask_last_page(struct hstate *h);
-- 
2.37.3