[PATCH] mm/migrate_device: avoid overflowing migrate_vma collection arrays

[Zi Yan]

migrate_vma_collect_pmd() can drop pte lock to split a large folio and
restart. But the code does not handle restart properly when pmd becomes
huge or cleared. It can overflow migrate->dst and migrate->src arrays
during the hole or skip collection. Fix it by:
1. avoiding migrate_vma_collect_huge_pmd() if some collection is done,
2. skipping the rest of the range if pmd no longer points to a pte page
   table.

Fixes: a30b48bf1b244 ("mm/migrate_device: implement THP migration of zone device pages")
Reported-by: Sashiko <sashiko-bot@xxxxxxxxxx>
Closes: https://sashiko.dev/#/patchset/20260706111958.3649651-1-wangkefeng.wang@xxxxxxxxxx
Assisted-by: Claude:claude-opus-4-8
Signed-off-by: Zi Yan <ziy@xxxxxxxxxx>
---
The issue is spot by Sashiko during a patch[1] review as a pre-existing one.
This patch has the minimal change. An alternative is to reset
migrate->cpages and migrate->npages and restart the whole range from the
beginning, but that also requires a restoration of no-longer-present PTEs.

Link: https://lore.kernel.org/all/20260706111958.3649651-1-wangkefeng.wang@xxxxxxxxxx/ [1]
---
 mm/migrate_device.c | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index 2fffeb1f99694..6ceb47ec1da24 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -257,7 +257,12 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
 	pte_t *ptep;
 
 again:
-	if (pmd_trans_huge(*pmdp) || !pmd_present(*pmdp)) {
+	/*
+	 * Only check pmd when addr is at start, namely no pte is collected.
+	 * It avoids collecting the same address range [start, addr) twice
+	 * and overflowing the collection arrays.
+	 */
+	if (addr == start && (pmd_trans_huge(*pmdp) || !pmd_present(*pmdp))) {
 		int ret = migrate_vma_collect_huge_pmd(pmdp, start, end, walk, fault_folio);
 
 		if (ret == -EAGAIN)
@@ -267,8 +272,18 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
 	}
 
 	ptep = pte_offset_map_lock(mm, pmdp, start, &ptl);
-	if (!ptep)
+	if (!ptep) {
+		/*
+		 * Skip the rest if pmd becomes huge or cleared. Flush if any
+		 * pte is modified
+		 */
+		if (addr != start) {
+			if (unmapped)
+				flush_tlb_range(walk->vma, start, end);
+			return migrate_vma_collect_skip(addr, end, walk);
+		}
 		goto again;
+	}
 	lazy_mmu_mode_enable();
 	ptep += (addr - start) / PAGE_SIZE;
 

---
base-commit: 302dfbff3e73be2ba5723237c6303244ec27cebb
change-id: 20260707-fix-array-overflow-in-migrate_vma_collect_pmd-c504b92613de

Best regards,
-- 
Yan, Zi