[PATCH 1/2] mm/huge_memory: take a reference on a folio PUD copied by fork

[Konstantin Khorenko]

Device DAX maps its memory with vmf_insert_folio_pud(), which takes a
folio reference, adds a PUD rmap and accounts the pages to the mm.
zap_huge_pud() drops all three, but copy_huge_pud() copies the PUD into
the child without taking any of them. Each child's unmap then drops
what the parent took:

- the folio is a reference short, so with two children its refcount
  drops to zero and free_zone_device_folio() runs while the parent
  still maps it;
- the child's file counter goes negative, "Bad rss-counter state" when
  its mm is freed;
- the parent's unmap drives the entire mapcount below zero.

fork copies the page tables of a device DAX mapping that has a guard
region: commit f807123d578d ("mm: allow guard regions in file-backed and
read-only mappings") allowed MADV_GUARD_INSTALL there, and it makes fork
copy the VMA.

Take the reference, the rmap and the counter for a PUD that
vm_normal_page_pud() resolves to a page, as copy_huge_pmd() does for a
PMD, with a new folio_dup_file_rmap_pud() helper.

Fixes: aed877c2b425 ("device/dax: properly refcount device dax pages when mapping")
Cc: stable@xxxxxxxxxxxxxxx
Assisted-by: LLM
Signed-off-by: Konstantin Khorenko <khorenko@xxxxxxxxxxxxx>
---
 include/linux/huge_mm.h |  2 +-
 include/linux/rmap.h    | 22 ++++++++++++++++++++++
 mm/huge_memory.c        | 13 +++++++++++--
 mm/memory.c             |  4 ++--
 4 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/include/linux/huge_mm.h b/include/linux/huge_mm.h
index c745f7ad22987..c5094f9d95ba9 100644
--- a/include/linux/huge_mm.h
+++ b/include/linux/huge_mm.h
@@ -14,7 +14,7 @@ int copy_huge_pmd(struct mm_struct *dst_mm, struct mm_struct *src_mm,
 bool huge_pmd_set_accessed(struct vm_fault *vmf);
 int copy_huge_pud(struct mm_struct *dst_mm, struct mm_struct *src_mm,
 		  pud_t *dst_pud, pud_t *src_pud, unsigned long addr,
-		  struct vm_area_struct *vma);
+		  struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma);
 
 #ifdef CONFIG_HAVE_ARCH_TRANSPARENT_HUGEPAGE_PUD
 void huge_pud_set_accessed(struct vm_fault *vmf, pud_t orig_pud);
diff --git a/include/linux/rmap.h b/include/linux/rmap.h
index 0b332770abeed..3c9ed8e27f0e4 100644
--- a/include/linux/rmap.h
+++ b/include/linux/rmap.h
@@ -564,6 +564,28 @@ static inline void folio_dup_file_rmap_pmd(struct folio *folio,
 #endif
 }
 
+/**
+ * folio_dup_file_rmap_pud - duplicate a PUD mapping of a page range of a folio
+ * @folio:	The folio to duplicate the mapping of
+ * @page:	The first page to duplicate the mapping of
+ * @dst_vma:	The destination vm area
+ *
+ * The page range of the folio is defined by [page, page + HPAGE_PUD_NR)
+ *
+ * The caller needs to hold the page table lock.
+ */
+static inline void folio_dup_file_rmap_pud(struct folio *folio,
+		struct page *page, struct vm_area_struct *dst_vma)
+{
+#if defined(CONFIG_TRANSPARENT_HUGEPAGE) && \
+	defined(CONFIG_HAVE_ARCH_TRANSPARENT_HUGEPAGE_PUD)
+	__folio_dup_file_rmap(folio, page, HPAGE_PUD_NR, dst_vma,
+			      PGTABLE_LEVEL_PUD);
+#else
+	WARN_ON_ONCE(true);
+#endif
+}
+
 static __always_inline int __folio_try_dup_anon_rmap(struct folio *folio,
 		struct page *page, int nr_pages, struct vm_area_struct *dst_vma,
 		struct vm_area_struct *src_vma, enum pgtable_level level)
diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index 1e5d68acf62a5..7e8b41f293a49 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -2075,9 +2075,10 @@ void touch_pud(struct vm_area_struct *vma, unsigned long addr,
 
 int copy_huge_pud(struct mm_struct *dst_mm, struct mm_struct *src_mm,
 		  pud_t *dst_pud, pud_t *src_pud, unsigned long addr,
-		  struct vm_area_struct *vma)
+		  struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma)
 {
 	spinlock_t *dst_ptl, *src_ptl;
+	struct page *page;
 	pud_t pud;
 	int ret;
 
@@ -2094,7 +2095,15 @@ int copy_huge_pud(struct mm_struct *dst_mm, struct mm_struct *src_mm,
 	 * TODO: once we support anonymous pages, use
 	 * folio_try_dup_anon_rmap_*() and split if duplicating fails.
 	 */
-	if (vma_is_cow_mapping(vma) && pud_write(pud)) {
+	page = vm_normal_page_pud(src_vma, addr, pud);
+	if (page) {
+		struct folio *folio = page_folio(page);
+
+		folio_get(folio);
+		folio_dup_file_rmap_pud(folio, page, dst_vma);
+		add_mm_counter(dst_mm, mm_counter_file(folio), HPAGE_PUD_NR);
+	}
+	if (vma_is_cow_mapping(src_vma) && pud_write(pud)) {
 		pudp_set_wrprotect(src_mm, addr, src_pud);
 		pud = pud_wrprotect(pud);
 	}
diff --git a/mm/memory.c b/mm/memory.c
index 8b0c2c735d3de..8acd4d5bb2310 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -1514,8 +1514,8 @@ copy_pud_range(struct vm_area_struct *dst_vma, struct vm_area_struct *src_vma,
 			int err;
 
 			VM_BUG_ON_VMA(next-addr != HPAGE_PUD_SIZE, src_vma);
-			err = copy_huge_pud(dst_mm, src_mm,
-					    dst_pud, src_pud, addr, src_vma);
+			err = copy_huge_pud(dst_mm, src_mm, dst_pud, src_pud,
+					    addr, dst_vma, src_vma);
 			if (err == -ENOMEM)
 				return -ENOMEM;
 			if (!err)

-- 
2.43.0