[PATCH v2 2/4] mm: don't report OOM when a huge PMD race blocks a pfn insert

[Eliav Farber <farbere@amazon.com>]

insert_pfn() maps the PTE with get_locked_pte() and treats a NULL return
as an allocation failure. pte_offset_map_lock() also returns NULL when
the PMD is a leaf, which a VM_PFNMAP mapping with both ->fault and
->huge_fault can hit: an order-0 fault that passed handle_pte_fault()'s
pmd_none() check can reach insert_pfn() after a concurrent huge fault
installed a PMD leaf over the same address. The task then takes a
VM_FAULT_OOM for a mapping that is perfectly healthy. Changing an mm's
THP policy while another thread faults the same PMD, for example with
prctl(PR_SET_THP_DISABLE), is enough to produce the mixed orders.

get_locked_pte() cannot tell the two apart, because pte_alloc_map_lock()
folds the allocation and the map into one call. Walk to the PMD,
allocate, and map as three steps, so that only a failed allocation
returns VM_FAULT_OOM. A NULL from the map returns VM_FAULT_NOPAGE and
the access refaults, which is what a PMD leaf, a PMD that is not present
and a bad PMD that __pte_offset_map() cleared all want.

Fixes: f9e54c3a2f5b ("vfio/pci: implement huge_fault support")
Signed-off-by: Eliav Farber <farbere@amazon.com>
Suggested-by: Jose A. Perez de Azpillaga <azpijr@gmail.com>
---
v2: Split pte_alloc_map_lock() into pte_alloc() followed by
    pte_offset_map_lock(), so that only a failed allocation returns
    VM_FAULT_OOM and a NULL from the map returns VM_FAULT_NOPAGE. v1
    called pte_alloc_map_lock() and then re-read the PMD with
    pmd_trans_huge() to tell the two apart, which left the non-present
    and bad-PMD cases reported as OOM. Suggested by Jose A. Perez de
    Azpillaga.

 mm/memory.c | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/mm/memory.c b/mm/memory.c
index 0f54ee0a6d88..05ae87518206 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -2866,10 +2866,24 @@ static vm_fault_t insert_pfn(struct vm_area_struct *vma, unsigned long addr,
 	struct mm_struct *mm = vma->vm_mm;
 	pte_t *pte, entry;
 	spinlock_t *ptl;
+	pmd_t *pmd;
 
-	pte = get_locked_pte(mm, addr, &ptl);
-	if (!pte)
+	pmd = walk_to_pmd(mm, addr);
+	if (!pmd)
+		return VM_FAULT_OOM;
+
+	if (pte_alloc(mm, pmd))
 		return VM_FAULT_OOM;
+
+	/*
+	 * A concurrent ->huge_fault may have installed a PMD leaf over addr,
+	 * which pte_offset_map_lock() refuses to walk. The mapping is fine, so
+	 * let the access refault instead of reporting a failed allocation.
+	 */
+	pte = pte_offset_map_lock(mm, pmd, addr, &ptl);
+	if (!pte)
+		return VM_FAULT_NOPAGE;
+
 	entry = ptep_get(pte);
 	if (!pte_none(entry)) {
 		if (mkwrite) {
-- 
2.47.3