From a662bedab8310ac4c15e1337e584e5f5835d14f5 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Tue, 10 Oct 2023 02:10:52 +0300
Subject: [PATCH] [NTOS:MM] Fix bugs in MmAccessFault

- Acquire the appropriate working set lock when calling MmLocateMemoryAreaByAddress
- Do not access MemoryArea without holding the lock (otherwise it can be pulled away under our feet)
- Fix range check for paged pool
---
 ntoskrnl/mm/mmfault.c | 34 ++++++++++++++++++++++++++++------
 1 file changed, 28 insertions(+), 6 deletions(-)

diff --git a/ntoskrnl/mm/mmfault.c b/ntoskrnl/mm/mmfault.c
index 68e31ba1281..cfd6d756d15 100644
--- a/ntoskrnl/mm/mmfault.c
+++ b/ntoskrnl/mm/mmfault.c
@@ -213,6 +213,7 @@ MmAccessFault(IN ULONG FaultCode,
 {
     PMEMORY_AREA MemoryArea = NULL;
     NTSTATUS Status;
+    BOOLEAN IsArm3Fault = FALSE;
 
     /* Cute little hack for ROS */
     if ((ULONG_PTR)Address >= (ULONG_PTR)MmSystemRangeStart)
@@ -239,19 +240,40 @@ MmAccessFault(IN ULONG FaultCode,
     /* Is there a ReactOS address space yet? */
     if (MmGetKernelAddressSpace())
     {
-        /* Check if this is an ARM3 memory area */
-        MemoryArea = MmLocateMemoryAreaByAddress(MmGetKernelAddressSpace(), Address);
-        if (!(MemoryArea) && (Address <= MM_HIGHEST_USER_ADDRESS))
+        if (Address > MM_HIGHEST_USER_ADDRESS)
+        {
+            /* Check if this is an ARM3 memory area */
+            MiLockWorkingSetShared(PsGetCurrentThread(), &MmSystemCacheWs);
+            MemoryArea = MmLocateMemoryAreaByAddress(MmGetKernelAddressSpace(), Address);
+
+            if ((MemoryArea != NULL) && (MemoryArea->Type == MEMORY_AREA_OWNED_BY_ARM3))
+            {
+                IsArm3Fault = TRUE;
+            }
+
+            MiUnlockWorkingSetShared(PsGetCurrentThread(), &MmSystemCacheWs);
+        }
+        else
         {
             /* Could this be a VAD fault from user-mode? */
+            MiLockProcessWorkingSetShared(PsGetCurrentProcess(), PsGetCurrentThread());
             MemoryArea = MmLocateMemoryAreaByAddress(MmGetCurrentAddressSpace(), Address);
+
+            if ((MemoryArea != NULL) && (MemoryArea->Type == MEMORY_AREA_OWNED_BY_ARM3))
+            {
+                IsArm3Fault = TRUE;
+            }
+
+            MiUnlockProcessWorkingSetShared(PsGetCurrentProcess(), PsGetCurrentThread());
         }
     }
 
     /* Is this an ARM3 memory area, or is there no address space yet? */
-    if (((MemoryArea) && (MemoryArea->Type == MEMORY_AREA_OWNED_BY_ARM3)) ||
-            (!(MemoryArea) && ((ULONG_PTR)Address >= (ULONG_PTR)MmPagedPoolStart)) ||
-            (!MmGetKernelAddressSpace()))
+    if (IsArm3Fault ||
+        ((MemoryArea == NULL) &&
+         ((ULONG_PTR)Address >= (ULONG_PTR)MmPagedPoolStart) &&
+         ((ULONG_PTR)Address < (ULONG_PTR)MmPagedPoolEnd)) ||
+        (!MmGetKernelAddressSpace()))
     {
         /* This is an ARM3 fault */
         DPRINT("ARM3 fault %p\n", MemoryArea);