From b3e9c8972553b040d1b43dd96406d28ebd5b953f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Gardou?= <jerome.gardou@reactos.org>
Date: Mon, 7 Jun 2021 19:28:51 +0200
Subject: [PATCH] [NTOS:MM] Simplify refcounting when adding/deleting a page in
 a shared segment

CORE-17544
---
 ntoskrnl/cache/section/sptab.c | 13 +++++++++++--
 ntoskrnl/mm/rmap.c             |  3 ---
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/ntoskrnl/cache/section/sptab.c b/ntoskrnl/cache/section/sptab.c
index 0877fb085f6..b0125f26b0e 100644
--- a/ntoskrnl/cache/section/sptab.c
+++ b/ntoskrnl/cache/section/sptab.c
@@ -204,6 +204,16 @@ _MmSetPageEntrySectionSegment(PMM_SECTION_SEGMENT Segment,
             OldEntry,
             Entry);
 
+    /* Manage ref on segment */
+    if (Entry && !OldEntry)
+    {
+        InterlockedIncrement64(Segment->ReferenceCount);
+    }
+    if (OldEntry && !Entry)
+    {
+        MmDereferenceSegment(Segment);
+    }
+
     if (Entry && !IS_SWAP_FROM_SSE(Entry))
     {
         /* We have a valid entry. See if we must do something */
@@ -223,7 +233,6 @@ _MmSetPageEntrySectionSegment(PMM_SECTION_SEGMENT Segment,
              * Add the Rmap and take a ref on the segment.
              */
             MmSetSectionAssociation(PFN_FROM_SSE(Entry), Segment, Offset);
-            InterlockedIncrement64(Segment->ReferenceCount);
 
             if (Offset->QuadPart >= (Segment->LastPage << PAGE_SHIFT))
                 Segment->LastPage = (Offset->QuadPart >> PAGE_SHIFT) + 1;
@@ -233,7 +242,6 @@ _MmSetPageEntrySectionSegment(PMM_SECTION_SEGMENT Segment,
     {
         /* We're switching to an invalid entry from a valid one */
         MmDeleteSectionAssociation(PFN_FROM_SSE(OldEntry));
-        MmDereferenceSegment(Segment);
 
         if (Offset->QuadPart == ((Segment->LastPage - 1ULL) << PAGE_SHIFT))
         {
@@ -363,6 +371,7 @@ MmGetSectionAssociation(PFN_NUMBER Page,
         Offset->QuadPart = PageTable->FileOffset.QuadPart +
                            ((ULONG64)RawOffset << PAGE_SHIFT);
         ASSERT(PFN_FROM_SSE(PageTable->PageEntries[RawOffset]) == Page);
+        InterlockedIncrement64(Segment->ReferenceCount);
     }
 
     return Segment;
diff --git a/ntoskrnl/mm/rmap.c b/ntoskrnl/mm/rmap.c
index 64758cc5b89..e27e31de820 100644
--- a/ntoskrnl/mm/rmap.c
+++ b/ntoskrnl/mm/rmap.c
@@ -300,7 +300,6 @@ WriteSegment:
         Released = MmCheckDirtySegment(Segment, &SegmentOffset, FALSE, TRUE);
 
         MmUnlockSectionSegment(Segment);
-
         MmDereferenceSegment(Segment);
 
         if (Released)
@@ -475,8 +474,6 @@ MmGetSegmentRmap(PFN_NUMBER Page, PULONG RawOffset)
                 MiReleasePfnLock(OldIrql);
                 return NULL;
             }
-
-            InterlockedIncrement64(Result->Segment->ReferenceCount);
             MiReleasePfnLock(OldIrql);
             return Result;
         }