[Patch v2] mm/page_vma_mapped: revalidate and do proper check before return device-private pmd

[Wei Yang <richard.weiyang@gmail.com>]

For pmd_trans_huge() and pmd_is_migration_entry(), we does following
before return the pmd entry:

  * re-validate pmd entry after PTL
  * check PVMW_MIGRATION
  * check_pmd()
  * handle on pte level if split under us

But for device-private pmd, we just return after pmd_lock().

This may return improper entry, e.g. if we are looking for a migration
entry, device-private entry could still be returned, which leads to data
corruption.

This patch fixes commit 65edfda6f3f2 ("mm/rmap: extend rmap and migration
support device-private entries") by following the same pattern as
pmd_trans_huge() and pmd_is_migration_entry() for device private entry.

While at it, it cleanups the pmd entry handling in page_vma_mapped_walk().

  * Instead of handling trans huge/migration entry/device private entry
    in a mixed manner, we put each case into its own if condition and
    handle with the same pattern.
  * Also we grab PTL and make sure pmd is not changed under us after
    above check instead of do the check with PTL hold.
  * restart the process if pmd is changed under us

Fixes: 65edfda6f3f2 ("mm/rmap: extend rmap and migration support device-private entries")
Signed-off-by: Wei Yang <richard.weiyang@gmail.com>
Suggested-by: David Hildenbrand <david@kernel.org>
Cc: David Hildenbrand <david@kernel.org>
Cc: Balbir Singh <balbirs@nvidia.com>
Cc: SeongJae Park <sj@kernel.org>
Cc: Zi Yan <ziy@nvidia.com>
Cc: Lorenzo Stoakes <lorenzo.stoakes@oracle.com>
Cc: <stable@vger.kernel.org>

---
v2:
  * specify the possible error case of current code and user visible effect
  * besides fix, cleanup the pmd entry handling based on David's suggestion

v1: https://lore.kernel.org/linux-mm/20260508013728.21285-1-richard.weiyang@gmail.com/
---
 mm/page_vma_mapped.c | 63 +++++++++++++++++++++-----------------------
 1 file changed, 30 insertions(+), 33 deletions(-)

diff --git a/mm/page_vma_mapped.c b/mm/page_vma_mapped.c
index 2ccbabfb2cc1..21635fab209c 100644
--- a/mm/page_vma_mapped.c
+++ b/mm/page_vma_mapped.c
@@ -243,40 +243,28 @@ bool page_vma_mapped_walk(struct page_vma_mapped_walk *pvmw)
 		 */
 		pmde = pmdp_get_lockless(pvmw->pmd);
 
-		if (pmd_trans_huge(pmde) || pmd_is_migration_entry(pmde)) {
-			pvmw->ptl = pmd_lock(mm, pvmw->pmd);
-			pmde = *pvmw->pmd;
-			if (!pmd_present(pmde)) {
-				softleaf_t entry;
-
-				if (!thp_migration_supported() ||
-				    !(pvmw->flags & PVMW_MIGRATION))
-					return not_found(pvmw);
-				entry = softleaf_from_pmd(pmde);
-
-				if (!softleaf_is_migration(entry) ||
-				    !check_pmd(softleaf_to_pfn(entry), pvmw))
-					return not_found(pvmw);
-				return true;
-			}
-			if (likely(pmd_trans_huge(pmde))) {
-				if (pvmw->flags & PVMW_MIGRATION)
-					return not_found(pvmw);
-				if (!check_pmd(pmd_pfn(pmde), pvmw))
-					return not_found(pvmw);
-				return true;
-			}
-			/* THP pmd was split under us: handle on pte level */
-			spin_unlock(pvmw->ptl);
-			pvmw->ptl = NULL;
-		} else if (!pmd_present(pmde)) {
-			const softleaf_t entry = softleaf_from_pmd(pmde);
-
-			if (softleaf_is_device_private(entry)) {
-				pvmw->ptl = pmd_lock(mm, pvmw->pmd);
-				return true;
-			}
+		if (pmd_present(pmde)) {
+			if (!pmd_leaf(pmde))
+				goto pte_table;
+			if (pvmw->flags & PVMW_MIGRATION)
+				return not_found(pvmw);
+			if (!check_pmd(pmd_pfn(pmde), pvmw))
+				return not_found(pvmw);
+		} else if (pmd_is_migration_entry(pmde)) {
+			softleaf_t entry = softleaf_from_pmd(pmde);
+
+			if (!(pvmw->flags & PVMW_MIGRATION))
+				return not_found(pvmw);
+			if (!check_pmd(softleaf_to_pfn(entry), pvmw))
+				return not_found(pvmw);
+		} else if (pmd_is_device_private_entry(pmde)) {
+			softleaf_t entry = softleaf_from_pmd(pmde);
 
+			if (pvmw->flags & PVMW_MIGRATION)
+				return not_found(pvmw);
+			if (!check_pmd(softleaf_to_pfn(entry), pvmw))
+				return not_found(pvmw);
+		} else {
 			if ((pvmw->flags & PVMW_SYNC) &&
 			    thp_vma_suitable_order(vma, pvmw->address,
 						   PMD_ORDER) &&
@@ -286,6 +274,15 @@ bool page_vma_mapped_walk(struct page_vma_mapped_walk *pvmw)
 			step_forward(pvmw, PMD_SIZE);
 			continue;
 		}
+
+		/* Double-check under PTL that the PMD didn't change. */
+		pvmw->ptl = pmd_lock(mm, pvmw->pmd);
+		if (pmd_same(pmde, pmdp_get(pvmw->pmd)))
+			return true;
+		spin_unlock(pvmw->ptl);
+		pvmw->ptl = NULL;
+		goto restart;
+pte_table:
 		if (!map_pte(pvmw, &pmde, &ptl)) {
 			if (!pvmw->pte)
 				goto restart;
-- 
2.34.1