From 0e58b5979431f34e29166517ec22111286929c90 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Sun, 7 Apr 2024 10:56:10 +0300
Subject: [PATCH] [NTOS:MM] Get rid of more calls to
 MmLocateMemoryAreaByAddress

---
 ntoskrnl/mm/ARM3/miarm.h   |  8 ++++----
 ntoskrnl/mm/ARM3/section.c |  2 +-
 ntoskrnl/mm/ARM3/virtual.c |  4 +---
 ntoskrnl/mm/marea.c        |  2 +-
 ntoskrnl/mm/section.c      | 12 ++++++------
 5 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/ntoskrnl/mm/ARM3/miarm.h b/ntoskrnl/mm/ARM3/miarm.h
index 31c19cd3154..45abbea577c 100644
--- a/ntoskrnl/mm/ARM3/miarm.h
+++ b/ntoskrnl/mm/ARM3/miarm.h
@@ -2264,10 +2264,10 @@ MiInsertBasedSection(
 NTSTATUS
 NTAPI
 MiRosUnmapViewOfSection(
-    IN PEPROCESS Process,
-    IN PVOID BaseAddress,
-    IN BOOLEAN SkipDebuggerNotify
-);
+    _In_ PEPROCESS Process,
+    _In_ PMEMORY_AREA MemoryArea,
+    _In_ PVOID BaseAddress,
+    _In_ BOOLEAN SkipDebuggerNotify);
 
 VOID
 NTAPI
diff --git a/ntoskrnl/mm/ARM3/section.c b/ntoskrnl/mm/ARM3/section.c
index 43c925a2d09..7c6d5da6488 100644
--- a/ntoskrnl/mm/ARM3/section.c
+++ b/ntoskrnl/mm/ARM3/section.c
@@ -833,7 +833,7 @@ MiUnmapViewOfSection(IN PEPROCESS Process,
     {
         /* Call Mm API */
         ASSERT(MI_IS_ROSMM_VAD(Vad));
-        NTSTATUS Status = MiRosUnmapViewOfSection(Process, BaseAddress, Process->ProcessExiting);
+        Status = MiRosUnmapViewOfSection(Process, (PMEMORY_AREA)Vad, BaseAddress, Process->ProcessExiting);
         if (!Flags) MmUnlockAddressSpace(&Process->Vm);
         return Status;
     }
diff --git a/ntoskrnl/mm/ARM3/virtual.c b/ntoskrnl/mm/ARM3/virtual.c
index 1adcbdec1ad..67abfc8accd 100644
--- a/ntoskrnl/mm/ARM3/virtual.c
+++ b/ntoskrnl/mm/ARM3/virtual.c
@@ -4460,7 +4460,6 @@ NtAllocateVirtualMemory(IN HANDLE ProcessHandle,
                         IN ULONG Protect)
 {
     PEPROCESS Process;
-    PMEMORY_AREA MemoryArea;
     PMMVAD Vad = NULL, FoundVad;
     NTSTATUS Status;
     PMMSUPPORT AddressSpace;
@@ -4876,8 +4875,7 @@ NtAllocateVirtualMemory(IN HANDLE ProcessHandle,
     //
     // Make sure this is an ARM3 section
     //
-    MemoryArea = MmLocateMemoryAreaByAddress(AddressSpace, (PVOID)PAGE_ROUND_DOWN(PBaseAddress));
-    if (MemoryArea && MemoryArea->Type != MEMORY_AREA_OWNED_BY_ARM3)
+    if (MI_IS_ROSMM_VAD(FoundVad))
     {
         DPRINT1("Illegal commit of non-ARM3 section!\n");
         Status = STATUS_ALREADY_COMMITTED;
diff --git a/ntoskrnl/mm/marea.c b/ntoskrnl/mm/marea.c
index 2b940a1d35a..b7de0df98f2 100644
--- a/ntoskrnl/mm/marea.c
+++ b/ntoskrnl/mm/marea.c
@@ -511,7 +511,7 @@ MiRosCleanupMemoryArea(
 
     if (MemoryArea->Type == MEMORY_AREA_SECTION_VIEW)
     {
-        Status = MiRosUnmapViewOfSection(Process, BaseAddress, Process->ProcessExiting);
+        Status = MiRosUnmapViewOfSection(Process, MemoryArea, BaseAddress, Process->ProcessExiting);
     }
 #ifdef NEWCC
     else if (MemoryArea->Type == MEMORY_AREA_CACHE)
diff --git a/ntoskrnl/mm/section.c b/ntoskrnl/mm/section.c
index abd27e1423f..c01bf3792f3 100644
--- a/ntoskrnl/mm/section.c
+++ b/ntoskrnl/mm/section.c
@@ -3599,12 +3599,13 @@ MmUnmapViewOfSegment(PMMSUPPORT AddressSpace,
 /* This functions must be called with a locked address space */
 NTSTATUS
 NTAPI
-MiRosUnmapViewOfSection(IN PEPROCESS Process,
-                        IN PVOID BaseAddress,
-                        IN BOOLEAN SkipDebuggerNotify)
+MiRosUnmapViewOfSection(
+    _In_ PEPROCESS Process,
+    _In_ PMEMORY_AREA MemoryArea,
+    _In_ PVOID BaseAddress,
+    _In_ BOOLEAN SkipDebuggerNotify)
 {
     NTSTATUS Status;
-    PMEMORY_AREA MemoryArea;
     PMMSUPPORT AddressSpace;
     PVOID ImageBaseAddress = 0;
 
@@ -3612,11 +3613,10 @@ MiRosUnmapViewOfSection(IN PEPROCESS Process,
            Process, BaseAddress);
 
     ASSERT(Process);
+    ASSERT(MemoryArea);
 
     AddressSpace = &Process->Vm;
 
-    MemoryArea = MmLocateMemoryAreaByAddress(AddressSpace,
-                 BaseAddress);
     if (MemoryArea == NULL ||
 #ifdef NEWCC
             ((MemoryArea->Type != MEMORY_AREA_SECTION_VIEW) && (MemoryArea->Type != MEMORY_AREA_CACHE)) ||