[reactos] 09/14: [NTOS:Mm] Handle failure to allocate pages in the page fault handler

[Timo Kreuzer]

https://git.reactos.org/?p=reactos.git;a=commitdiff;h=7c6e4d38c7718a7ecfb92…

commit 7c6e4d38c7718a7ecfb927025cf968356ccc6d42
Author:     Timo Kreuzer <timo.kreuzer(a)reactos.org>
AuthorDate: Wed Apr 5 22:14:09 2023 +0300
Commit:     Timo Kreuzer <timo.kreuzer(a)reactos.org>
CommitDate: Sat Jul 29 14:00:44 2023 +0300

    [NTOS:Mm] Handle failure to allocate pages in the page fault handler
---
 ntoskrnl/mm/ARM3/pagfault.c | 63 ++++++++++++++++++++++++++++++++++++++-------
 ntoskrnl/mm/section.c       | 59 +++++++++++++++++++++++++++++-------------
 2 files changed, 95 insertions(+), 27 deletions(-)

diff --git a/ntoskrnl/mm/ARM3/pagfault.c b/ntoskrnl/mm/ARM3/pagfault.c
index a5dd04560d3..a312f4a8bdd 100644
--- a/ntoskrnl/mm/ARM3/pagfault.c
+++ b/ntoskrnl/mm/ARM3/pagfault.c
@@ -653,7 +653,7 @@ MiResolveDemandZeroFault(IN PVOID Address,
     ASSERT(PointerPte->u.Hard.Valid == 0);
 
     /* Assert we have enough pages */
-    ASSERT(MmAvailablePages >= 32);
+    //ASSERT(MmAvailablePages >= 32);
 
 #if MI_TRACE_PFNS
     if (UserPdeFault) MI_SET_USAGE(MI_USAGE_PAGE_TABLE);
@@ -698,6 +698,12 @@ MiResolveDemandZeroFault(IN PVOID Address,
         }
     }
 
+    if (PageFrameNumber == 0)
+    {
+        MiReleasePfnLock(OldIrql);
+        return STATUS_NO_MEMORY;
+    }
+
     /* Initialize it */
     MiInitializePfn(PageFrameNumber, PointerPte, TRUE);
 
@@ -920,6 +926,10 @@ MiResolvePageFileFault(_In_ BOOLEAN StoreInstruction,
     /* Get any page, it will be overwritten */
     Color = MI_GET_NEXT_PROCESS_COLOR(CurrentProcess);
     Page = MiRemoveAnyPage(Color);
+    if (Page == 0)
+    {
+        return STATUS_NO_MEMORY;
+    }
 
     /* Initialize this PFN */
     MiInitializePfn(Page, PointerPte, StoreInstruction);
@@ -1230,6 +1240,11 @@ MiResolveProtoPteFault(IN BOOLEAN StoreInstruction,
             Color = MI_GET_NEXT_COLOR();
 
         PageFrameIndex = MiRemoveAnyPage(Color);
+        if (PageFrameIndex == 0)
+        {
+            MiReleasePfnLock(OldIrql);
+            return STATUS_NO_MEMORY;
+        }
 
         /* Perform the copy */
         MiCopyPfn(PageFrameIndex, ProtoPageFrameIndex);
@@ -1669,9 +1684,9 @@ MiDispatchFault(IN ULONG FaultCode,
     }
 
     //
-    // Generate an access fault
+    // Return status
     //
-    return STATUS_ACCESS_VIOLATION;
+    return Status;
 }
 
 NTSTATUS
@@ -2126,11 +2141,15 @@ UserFault:
         }
 
         /* Resolve a demand zero fault */
-        MiResolveDemandZeroFault(PointerPpe,
+        Status = MiResolveDemandZeroFault(PointerPpe,
                                  PointerPxe,
                                  MM_EXECUTE_READWRITE,
                                  CurrentProcess,
                                  MM_NOIRQL);
+        if (!NT_SUCCESS(Status))
+        {
+            goto ExitUser;
+        }
 
         /* We should come back with a valid PXE */
         ASSERT(PointerPxe->u.Hard.Valid == 1);
@@ -2160,11 +2179,15 @@ UserFault:
         }
 
         /* Resolve a demand zero fault */
-        MiResolveDemandZeroFault(PointerPde,
+        Status = MiResolveDemandZeroFault(PointerPde,
                                  PointerPpe,
                                  MM_EXECUTE_READWRITE,
                                  CurrentProcess,
                                  MM_NOIRQL);
+        if (!NT_SUCCESS(Status))
+        {
+            goto ExitUser;
+        }
 
         /* We should come back with a valid PPE */
         ASSERT(PointerPpe->u.Hard.Valid == 1);
@@ -2203,11 +2226,16 @@ UserFault:
         }
 
         /* Resolve a demand zero fault */
-        MiResolveDemandZeroFault(PointerPte,
+        Status = MiResolveDemandZeroFault(PointerPte,
                                  PointerPde,
                                  MM_EXECUTE_READWRITE,
                                  CurrentProcess,
                                  MM_NOIRQL);
+        if (!NT_SUCCESS(Status))
+        {
+            goto ExitUser;
+        }
+
 #if _MI_PAGING_LEVELS >= 3
         MiIncrementPageTableReferences(PointerPte);
 #endif
@@ -2254,6 +2282,12 @@ UserFault:
 
                 /* Allocate a new page and copy it */
                 PageFrameIndex =
MiRemoveAnyPage(MI_GET_NEXT_PROCESS_COLOR(CurrentProcess));
+                if (PageFrameIndex == 0)
+                {
+                    MiReleasePfnLock(LockIrql);
+                    Status = STATUS_NO_MEMORY;
+                    goto ExitUser;
+                }
                 OldPageFrameIndex = PFN_FROM_PTE(&TempPte);
 
                 MiCopyPfn(PageFrameIndex, OldPageFrameIndex);
@@ -2308,11 +2342,15 @@ UserFault:
         (TempPte.u.Long == (MM_EXECUTE_READWRITE <<
MM_PTE_SOFTWARE_PROTECTION_BITS)))
     {
         /* Resolve the fault */
-        MiResolveDemandZeroFault(Address,
+        Status = MiResolveDemandZeroFault(Address,
                                  PointerPte,
                                  TempPte.u.Soft.Protection,
                                  CurrentProcess,
                                  MM_NOIRQL);
+        if (!NT_SUCCESS(Status))
+        {
+            goto ExitUser;
+        }
 
 #if MI_TRACE_PFNS
         /* Update debug info */
@@ -2410,7 +2448,7 @@ UserFault:
             OldIrql = MiAcquirePfnLock();
 
             /* Make sure we have enough pages */
-            ASSERT(MmAvailablePages >= 32);
+            //ASSERT(MmAvailablePages >= 32);
 
             /* Try to get a zero page */
             MI_SET_USAGE(MI_USAGE_PEB_TEB);
@@ -2421,11 +2459,16 @@ UserFault:
             {
                 /* Grab a page out of there. Later we should grab a colored zero page */
                 PageFrameIndex = MiRemoveAnyPage(Color);
-                ASSERT(PageFrameIndex);
 
                 /* Release the lock since we need to do some zeroing */
                 MiReleasePfnLock(OldIrql);
 
+                if (PageFrameIndex == 0)
+                {
+                    Status = STATUS_NO_MEMORY;
+                    goto ExitUser;
+                }
+
                 /* Zero out the page, since it's for user-mode */
                 MiZeroPfn(PageFrameIndex);
 
@@ -2568,6 +2611,8 @@ UserFault:
                              TrapInformation,
                              Vad);
 
+ExitUser:
+
     /* Return the status */
     ASSERT(KeGetCurrentIrql() <= APC_LEVEL);
     MiUnlockProcessWorkingSet(CurrentProcess, CurrentThread);
diff --git a/ntoskrnl/mm/section.c b/ntoskrnl/mm/section.c
index 7b2e9d0d694..653b4759b27 100644
--- a/ntoskrnl/mm/section.c
+++ b/ntoskrnl/mm/section.c
@@ -1335,9 +1335,16 @@ MmMakeSegmentResident(
             RtlZeroMemory(Pages, BYTES_TO_PAGES(ReadLength) * sizeof(PFN_NUMBER));
             for (UINT i = 0; i < BYTES_TO_PAGES(ReadLength); i++)
             {
-                /* MmRequestPageMemoryConsumer succeeds or bugchecks */
-                (void)MmRequestPageMemoryConsumer(MC_USER, FALSE, &Pages[i]);
+                Status = MmRequestPageMemoryConsumer(MC_USER, FALSE, &Pages[i]);
+                if (!NT_SUCCESS(Status))
+                {
+                    /* Damn. Roll-back. */
+                    for (UINT j = 0; j < i; j++)
+                        MmReleasePageMemoryConsumer(MC_USER, Pages[j]);
+                    goto Failed;
+                }
             }
+
             Mdl->MdlFlags |= MDL_PAGES_LOCKED | MDL_IO_PAGE_READ;
 
             LARGE_INTEGER FileOffset;
@@ -1389,6 +1396,7 @@ MmMakeSegmentResident(
                 for (UINT i = 0; i < BYTES_TO_PAGES(ReadLength); i++)
                     MmReleasePageMemoryConsumer(MC_USER, Pages[i]);
 
+Failed:
                 MmLockSectionSegment(Segment);
                 while (ChunkOffset < ChunkEnd)
                 {
@@ -1613,6 +1621,15 @@ MmNotPresentFaultSectionView(PMMSUPPORT AddressSpace,
             return STATUS_MM_RESTART_OPERATION;
         }
 
+        MI_SET_USAGE(MI_USAGE_SECTION);
+        if (Process) MI_SET_PROCESS2(Process->ImageFileName);
+        if (!Process) MI_SET_PROCESS2("Kernel Section");
+        Status = MmRequestPageMemoryConsumer(MC_USER, TRUE, &Page);
+        if (!NT_SUCCESS(Status))
+        {
+            return STATUS_NO_MEMORY;
+        }
+
         /*
          * Must be private page we have swapped out.
          */
@@ -1627,14 +1644,6 @@ MmNotPresentFaultSectionView(PMMSUPPORT AddressSpace,
         MmCreatePageFileMapping(Process, Address, MM_WAIT_ENTRY);
 
         MmUnlockAddressSpace(AddressSpace);
-        MI_SET_USAGE(MI_USAGE_SECTION);
-        if (Process) MI_SET_PROCESS2(Process->ImageFileName);
-        if (!Process) MI_SET_PROCESS2("Kernel Section");
-        Status = MmRequestPageMemoryConsumer(MC_USER, TRUE, &Page);
-        if (!NT_SUCCESS(Status))
-        {
-            KeBugCheck(MEMORY_MANAGEMENT);
-        }
 
         Status = MmReadFromSwapPage(SwapEntry, Page);
         if (!NT_SUCCESS(Status))
@@ -1736,7 +1745,12 @@ MmNotPresentFaultSectionView(PMMSUPPORT AddressSpace,
             MI_SET_USAGE(MI_USAGE_SECTION);
             if (Process) MI_SET_PROCESS2(Process->ImageFileName);
             if (!Process) MI_SET_PROCESS2("Kernel Section");
-            MmRequestPageMemoryConsumer(MC_USER, FALSE, &Page);
+            Status = MmRequestPageMemoryConsumer(MC_USER, FALSE, &Page);
+            if (!NT_SUCCESS(Status))
+            {
+                MmUnlockSectionSegment(Segment);
+                return STATUS_NO_MEMORY;
+            }
             MmSetPageEntrySectionSegment(Segment, &Offset, MAKE_SSE(Page <<
PAGE_SHIFT, 1));
             MmUnlockSectionSegment(Segment);
 
@@ -1770,6 +1784,10 @@ MmNotPresentFaultSectionView(PMMSUPPORT AddressSpace,
         MmLockAddressSpace(AddressSpace);
         if (!NT_SUCCESS(Status))
         {
+            if (Status == STATUS_NO_MEMORY)
+            {
+                return Status;
+            }
             /* Damn */
             DPRINT1("Failed to page data in!\n");
             return STATUS_IN_PAGE_ERROR;
@@ -1794,6 +1812,13 @@ MmNotPresentFaultSectionView(PMMSUPPORT AddressSpace,
             return STATUS_MM_RESTART_OPERATION;
         }
 
+        Status = MmRequestPageMemoryConsumer(MC_USER, TRUE, &Page);
+        if (!NT_SUCCESS(Status))
+        {
+            MmUnlockSectionSegment(Segment);
+            return STATUS_NO_MEMORY;
+        }
+
         /*
         * Release all our locks and read in the page from disk
         */
@@ -1801,11 +1826,6 @@ MmNotPresentFaultSectionView(PMMSUPPORT AddressSpace,
         MmUnlockSectionSegment(Segment);
 
         MmUnlockAddressSpace(AddressSpace);
-        Status = MmRequestPageMemoryConsumer(MC_USER, TRUE, &Page);
-        if (!NT_SUCCESS(Status))
-        {
-            KeBugCheck(MEMORY_MANAGEMENT);
-        }
 
         Status = MmReadFromSwapPage(SwapEntry, Page);
         if (!NT_SUCCESS(Status))
@@ -1905,6 +1925,7 @@ MmAccessFaultSectionView(PMMSUPPORT AddressSpace,
     BOOLEAN Cow = FALSE;
     ULONG NewProtect;
     BOOLEAN Unmapped;
+    NTSTATUS Status;
 
     DPRINT("MmAccessFaultSectionView(%p, %p, %p)\n", AddressSpace, MemoryArea,
Address);
 
@@ -1990,9 +2011,11 @@ MmAccessFaultSectionView(PMMSUPPORT AddressSpace,
     /*
      * Allocate a page
      */
-    if (!NT_SUCCESS(MmRequestPageMemoryConsumer(MC_USER, TRUE, &NewPage)))
+    Status = MmRequestPageMemoryConsumer(MC_USER, TRUE, &NewPage);
+    if (!NT_SUCCESS(Status))
     {
-        KeBugCheck(MEMORY_MANAGEMENT);
+        MmUnlockSectionSegment(Segment);
+        return STATUS_NO_MEMORY;
     }
 
     /*