From ec45b8fde8bb56c79af8fcf1544b3b7174408a67 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Sat, 6 Apr 2024 13:29:30 +0300
Subject: [PATCH] [NTOS:MM] Stop using memory areas in MmAccessFault
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Implement MiLocateVad and make MiLocateAddress call it
- Use MiLocateVad in MmAccessFault to get the VAD
- Check VadFlags for whether this is an ARM³ VAD/MEMORY_AREA
---
 ntoskrnl/mm/ARM3/miarm.h   |  6 ++++++
 ntoskrnl/mm/ARM3/vadnode.c | 11 +++++++++--
 ntoskrnl/mm/mmfault.c      | 18 ++++++++++--------
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/ntoskrnl/mm/ARM3/miarm.h b/ntoskrnl/mm/ARM3/miarm.h
index 8e5fc41b89c..f4721abc5b6 100644
--- a/ntoskrnl/mm/ARM3/miarm.h
+++ b/ntoskrnl/mm/ARM3/miarm.h
@@ -2178,6 +2178,12 @@ MiIsPfnInUse(
     IN PMMPFN Pfn1
 );
 
+PMMVAD
+NTAPI
+MiLocateVad(
+    _In_ PMM_AVL_TABLE Table,
+    _In_ PVOID VirtualAddress);
+
 PMMVAD
 NTAPI
 MiLocateAddress(
diff --git a/ntoskrnl/mm/ARM3/vadnode.c b/ntoskrnl/mm/ARM3/vadnode.c
index a4820dac5c8..ad7d4c90af8 100644
--- a/ntoskrnl/mm/ARM3/vadnode.c
+++ b/ntoskrnl/mm/ARM3/vadnode.c
@@ -113,11 +113,10 @@ MiDbgAssertIsLockedForWrite(_In_ PMM_AVL_TABLE Table)
 
 PMMVAD
 NTAPI
-MiLocateAddress(IN PVOID VirtualAddress)
+MiLocateVad(_In_ PMM_AVL_TABLE Table, _In_ PVOID VirtualAddress)
 {
     PMMVAD FoundVad;
     ULONG_PTR Vpn;
-    PMM_AVL_TABLE Table = &PsGetCurrentProcess()->VadRoot;
     TABLE_SEARCH_RESULT SearchResult;
 
     ASSERT_LOCKED_FOR_READ(Table);
@@ -145,6 +144,14 @@ MiLocateAddress(IN PVOID VirtualAddress)
     return FoundVad;
 }
 
+PMMVAD
+NTAPI
+MiLocateAddress(_In_ PVOID VirtualAddress)
+{
+    PMM_AVL_TABLE Table = &PsGetCurrentProcess()->VadRoot;
+    return MiLocateVad(Table, VirtualAddress);
+}
+
 TABLE_SEARCH_RESULT
 NTAPI
 MiCheckForConflictingNode(IN ULONG_PTR StartVpn,
diff --git a/ntoskrnl/mm/mmfault.c b/ntoskrnl/mm/mmfault.c
index 5dbdc7a8ac2..ae332ec9d4e 100644
--- a/ntoskrnl/mm/mmfault.c
+++ b/ntoskrnl/mm/mmfault.c
@@ -16,6 +16,8 @@
 #define MODULE_INVOLVED_IN_ARM3
 #include "ARM3/miarm.h"
 
+extern MM_AVL_TABLE MiRosKernelVadRoot;
+
 /* PRIVATE FUNCTIONS **********************************************************/
 
 NTSTATUS
@@ -220,7 +222,7 @@ MmAccessFault(IN ULONG FaultCode,
               IN KPROCESSOR_MODE Mode,
               IN PVOID TrapInformation)
 {
-    PMEMORY_AREA MemoryArea = NULL;
+    PMMVAD Vad = NULL;
     NTSTATUS Status;
     BOOLEAN IsArm3Fault = FALSE;
 
@@ -253,9 +255,9 @@ MmAccessFault(IN ULONG FaultCode,
         {
             /* Check if this is an ARM3 memory area */
             MiLockWorkingSetShared(PsGetCurrentThread(), &MmSystemCacheWs);
-            MemoryArea = MmLocateMemoryAreaByAddress(MmGetKernelAddressSpace(), Address);
+            Vad = MiLocateVad(&MiRosKernelVadRoot, Address);
 
-            if ((MemoryArea != NULL) && (MemoryArea->Type == MEMORY_AREA_OWNED_BY_ARM3))
+            if ((Vad != NULL) && !MI_IS_ROSMM_VAD(Vad))
             {
                 IsArm3Fault = TRUE;
             }
@@ -266,9 +268,9 @@ MmAccessFault(IN ULONG FaultCode,
         {
             /* Could this be a VAD fault from user-mode? */
             MiLockProcessWorkingSetShared(PsGetCurrentProcess(), PsGetCurrentThread());
-            MemoryArea = MmLocateMemoryAreaByAddress(MmGetCurrentAddressSpace(), Address);
+            Vad = MiLocateVad(&PsGetCurrentProcess()->VadRoot, Address);
 
-            if ((MemoryArea != NULL) && (MemoryArea->Type == MEMORY_AREA_OWNED_BY_ARM3))
+            if ((Vad != NULL) && !MI_IS_ROSMM_VAD(Vad))
             {
                 IsArm3Fault = TRUE;
             }
@@ -277,15 +279,15 @@ MmAccessFault(IN ULONG FaultCode,
         }
     }
 
-    /* Is this an ARM3 memory area, or is there no address space yet? */
+    /* Is this an ARM3 VAD, or is there no address space yet? */
     if (IsArm3Fault ||
-        ((MemoryArea == NULL) &&
+        ((Vad == 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);
+        DPRINT("ARM3 fault %p\n", Vad);
         return MmArmAccessFault(FaultCode, Address, Mode, TrapInformation);
     }