mm/userfaultfd: fix and cleanup for migration entries with uffd-wp

One fix and one cleanup.

uffd-wp migration entry handling for PTE/PMDs should now be fairly similar code-wise.

v1 -> v2:
- Add RB's
- "mm/userfaultfd: fix uffd-wp handling for THP migration entries" -> pmd_swp_uffd_wp(*pvmw->pmd) -> pmd_uffd_wp(pmdval)

David Hildenbrand (2):
mm/userfaultfd: fix uffd-wp handling for THP migration entries mm/userfaultfd: don't consider uffd-wp bit of writable migration
entries

 mm/huge_memory.c | 16 ++++++++++++----
 mm/mprotect.c    |  2 --
 2 files changed, 12 insertions(+), 6 deletions(-)

2.39.2

Looks like what we fixed for hugetlb in commit 44f86392bdd1 ("mm/hugetlb:
fix uffd-wp handling for migration entries in hugetlb_change_protection()") similarly applies to THP.

Setting/clearing uffd-wp on THP migration entries is not implemented properly. Further, while removing migration PMDs considers the uffd-wp bit, inserting migration PMDs does not consider the uffd-wp bit.

We have to set/clear independently of the migration entry type in change_huge_pmd() and properly copy the uffd-wp bit in set_pmd_migration_entry().

Verified using a simple reproducer that triggers migration of a THP, that the set_pmd_migration_entry() no longer loses the uffd-wp bit.

---
 mm/huge_memory.c | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

index 032fb0ef9cd1..e3706a2b34b2 100644
@@ -1838,10 +1838,10 @@ int change_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
 	if (is_swap_pmd(*pmd)) {
 		swp_entry_t entry = pmd_to_swp_entry(*pmd);
 		struct page *page = pfn_swap_entry_to_page(entry);
+		pmd_t newpmd;
 
 		VM_BUG_ON(!is_pmd_migration_entry(*pmd));
 		if (is_writable_migration_entry(entry)) {
-			pmd_t newpmd;
 			/*
 			 * A protection check is difficult so
 			 * just be safe and disable write

@@ -1855,8 +1855,16 @@ int change_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
 				newpmd = pmd_swp_mksoft_dirty(newpmd);
 			if (pmd_swp_uffd_wp(*pmd))
 				newpmd = pmd_swp_mkuffd_wp(newpmd);
-			set_pmd_at(mm, addr, pmd, newpmd);
+		} else {
+			newpmd = *pmd;
 		}
+
+		if (uffd_wp)
+			newpmd = pmd_swp_mkuffd_wp(newpmd);
+		else if (uffd_wp_resolve)
+			newpmd = pmd_swp_clear_uffd_wp(newpmd);
+		if (!pmd_same(*pmd, newpmd))
+			set_pmd_at(mm, addr, pmd, newpmd);
 		goto unlock;
 	}
 #endif

@@ -3251,6 +3259,8 @@ int set_pmd_migration_entry(struct page_vma_mapped_walk *pvmw,
 	pmdswp = swp_entry_to_pmd(entry);
 	if (pmd_soft_dirty(pmdval))
 		pmdswp = pmd_swp_mksoft_dirty(pmdswp);
+	if (pmd_uffd_wp(pmdval))
+		pmdswp = pmd_swp_mkuffd_wp(pmdswp);
 	set_pmd_at(mm, address, pvmw->pmd, pmdswp);
 	page_remove_rmap(page, vma, true);
 	put_page(page);

2.39.2

If we end up with a writable migration entry that has the uffd-wp bit set, we already messed up: the source PTE/PMD was writable, which means we could have modified the page without notifying uffd first. Setting the uffd-wp bit always implies converting migration entries to !writable migration entries.

Commit 8f34f1eac382 ("mm/userfaultfd: fix uffd-wp special cases for fork()") documents that "3. Forget to carry over uffd-wp bit for a write migration huge pmd entry", but it doesn't really say why that should be relevant.

So let's remove that code to avoid hiding an eventual underlying issue (in the future, we might want to warn when creating writable migration
entries that have the uffd-wp bit set -- or even better when turning a PTE writable that still has the uffd-wp bit set).

This now matches the handling for hugetlb migration entries in hugetlb_change_protection().

In copy_huge_pmd()/copy_nonpresent_pte()/copy_hugetlb_page_range(), we still transfer the uffd-bit also for writable migration entries, but simply because we have unified handling for "writable" and "readable-exclusive" migration entries, and we care about transferring the uffd-wp bit for the latter.

---
 mm/huge_memory.c | 2 --
 mm/mprotect.c    | 2 --
 2 files changed, 4 deletions(-)

index e3706a2b34b2..fffc953fa6ea 100644
@@ -1853,8 +1853,6 @@ int change_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
 			newpmd = swp_entry_to_pmd(entry);
 			if (pmd_swp_soft_dirty(*pmd))
 				newpmd = pmd_swp_mksoft_dirty(newpmd);
-			if (pmd_swp_uffd_wp(*pmd))
-				newpmd = pmd_swp_mkuffd_wp(newpmd);
 		} else {
 			newpmd = *pmd;
 		}

index 13e84d8c0797..e04e9ea62ae7 100644
@@ -223,8 +223,6 @@ static long change_pte_range(struct mmu_gather *tlb,
 				newpte = swp_entry_to_pte(entry);
 				if (pte_swp_soft_dirty(oldpte))
 					newpte = pte_swp_mksoft_dirty(newpte);
-				if (pte_swp_uffd_wp(oldpte))
-					newpte = pte_swp_mkuffd_wp(newpte);
 			} else if (is_writable_device_private_entry(entry)) {
 				/*
 				 * We do not preserve soft-dirtiness. See

2.39.2