From 54d1b396767b40b7f864644bc8b732861eac0ff7 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Mon, 10 Apr 2023 00:26:34 +0300
Subject: [PATCH] [NTOS:Mm] Fix locking for MmGetSectionAssociation

MmGetSectionAssociation races with _MmSetPageEntrySectionSegment without sharing a lock. So we need to hold the PFN lock, until we have referenced the section segment found in the RMAP. This prevents that a section segment, which still has associated RMAPs from being deleted behind our back.
---
 ntoskrnl/cache/section/sptab.c | 4 ++++
 ntoskrnl/mm/rmap.c             | 9 +++++----
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/ntoskrnl/cache/section/sptab.c b/ntoskrnl/cache/section/sptab.c
index b0125f26b0e..5149ea3b76e 100644
--- a/ntoskrnl/cache/section/sptab.c
+++ b/ntoskrnl/cache/section/sptab.c
@@ -364,6 +364,8 @@ MmGetSectionAssociation(PFN_NUMBER Page,
     PMM_SECTION_SEGMENT Segment = NULL;
     PCACHE_SECTION_PAGE_TABLE PageTable;
 
+    KIRQL OldIrql = MiAcquirePfnLock();
+
     PageTable = MmGetSegmentRmap(Page, &RawOffset);
     if (PageTable)
     {
@@ -374,6 +376,8 @@ MmGetSectionAssociation(PFN_NUMBER Page,
         InterlockedIncrement64(Segment->ReferenceCount);
     }
 
+    MiReleasePfnLock(OldIrql);
+
     return Segment;
 }
 
diff --git a/ntoskrnl/mm/rmap.c b/ntoskrnl/mm/rmap.c
index 3447bda81cd..c8f16300f6d 100644
--- a/ntoskrnl/mm/rmap.c
+++ b/ntoskrnl/mm/rmap.c
@@ -469,7 +469,9 @@ MmGetSegmentRmap(PFN_NUMBER Page, PULONG RawOffset)
 {
     PCACHE_SECTION_PAGE_TABLE Result = NULL;
     PMM_RMAP_ENTRY current_entry;//, previous_entry;
-    KIRQL OldIrql = MiAcquirePfnLock();
+
+    /* Must hold the PFN lock */
+    ASSERT(KeGetCurrentIrql() == DISPATCH_LEVEL);
 
     //previous_entry = NULL;
     current_entry = MmGetRmapListHeadPage(Page);
@@ -481,16 +483,15 @@ MmGetSegmentRmap(PFN_NUMBER Page, PULONG RawOffset)
             *RawOffset = (ULONG_PTR)current_entry->Address & ~RMAP_SEGMENT_MASK;
             if (*Result->Segment->Flags & MM_SEGMENT_INDELETE)
             {
-                MiReleasePfnLock(OldIrql);
                 return NULL;
             }
-            MiReleasePfnLock(OldIrql);
+
             return Result;
         }
         //previous_entry = current_entry;
         current_entry = current_entry->Next;
     }
-    MiReleasePfnLock(OldIrql);
+
     return NULL;
 }