[PATCH v2] mm/migrate_device: fix pgtable leak in migrate_vma_insert_huge_pmd_page

[Sunny Patel <nueralspacetech@gmail.com>]

When migrate_vma_insert_huge_pmd_page() jumps to unlock_abort due
to a PMD check failure, the pgtable allocated earlier via
pte_alloc_one() is never freed, causing a memory leak.

Add a pte_free() call in the unlock_abort error path to release
the pgtable before returning.Also included before goto abort in the
folio check path.

Signed-off-by: Sunny Patel <nueralspacetech@gmail.com>
---

Changes in v2:
  - Added pte_free() before goto abort in the folio_is_zone_device()
    check path. The lock is not taken at this point so goto unlock_abort would be incorrect here.
  - v1 only fixed the unlock_abort path, this version fixes both
    leak locations.

 mm/migrate_device.c | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index fbfe5715f635..7e132196856b 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -840,6 +840,7 @@ static int migrate_vma_insert_huge_pmd_page(struct migrate_vma *migrate,
 	} else {
 		if (folio_is_zone_device(folio) &&
 		    !folio_is_device_coherent(folio)) {
+			pte_free(vma->vm_mm, pgtable);
 			goto abort;
 		}
 		entry = folio_mk_pmd(folio, vma->vm_page_prot);
@@ -893,6 +894,7 @@ static int migrate_vma_insert_huge_pmd_page(struct migrate_vma *migrate,
 
 unlock_abort:
 	spin_unlock(ptl);
+	pte_free(vma->vm_mm, pgtable);
 abort:
 	for (i = 0; i < HPAGE_PMD_NR; i++)
 		src[i] &= ~MIGRATE_PFN_MIGRATE;
-- 
2.43.0