From acf28dbc1eec4748f8bc6b01d8862783894d551b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Gardou?= <jerome.gardou@reactos.org>
Date: Wed, 4 Aug 2021 20:35:14 +0200
Subject: [PATCH] [NTOS:MM] Fix the mess created by the "balancer thread"

When processing:
  Make sure that the process is not terminating.
  Make sure that the process WorkingSet is still valid
  Protect accessing & writing to PTEs by acquiring the working set lock

CORE-17595 CORE-17642
---
 ntoskrnl/mm/balance.c | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/ntoskrnl/mm/balance.c b/ntoskrnl/mm/balance.c
index 0eec51fbb88..4558deeff18 100644
--- a/ntoskrnl/mm/balance.c
+++ b/ntoskrnl/mm/balance.c
@@ -207,16 +207,24 @@ MmTrimUserMemory(ULONG Target, ULONG Priority, PULONG NrFreedPages)
                 Process = Entry->Process;
                 Address = Entry->Address;
 
+                ObReferenceObject(Process);
+
+                if (!ExAcquireRundownProtection(&Process->RundownProtect))
+                {
+                    ObDereferenceObject(Process);
+                    MiReleasePfnLock(OldIrql);
+                    continue;
+                }
+
                 MiReleasePfnLock(OldIrql);
 
                 KeStackAttachProcess(&Process->Pcb, &ApcState);
-
-                MmLockAddressSpace(&Process->Vm);
+                MiLockProcessWorkingSet(Process, PsGetCurrentThread());
 
                 /* Be sure this is still valid. */
-                PMMPTE Pte = MiAddressToPte(Address);
-                if (Pte->u.Hard.Valid)
+                if (MmIsAddressValid(Address))
                 {
+                    PMMPTE Pte = MiAddressToPte(Address);
                     Accessed = Accessed || Pte->u.Hard.Accessed;
                     Pte->u.Hard.Accessed = 0;
 
@@ -224,9 +232,11 @@ MmTrimUserMemory(ULONG Target, ULONG Priority, PULONG NrFreedPages)
                     //KeInvalidateTlbEntry(Address);
                 }
 
-                MmUnlockAddressSpace(&Process->Vm);
+                MiUnlockProcessWorkingSet(Process, PsGetCurrentThread());
 
                 KeUnstackDetachProcess(&ApcState);
+                ExReleaseRundownProtection(&Process->RundownProtect);
+                ObDereferenceObject(Process);
             }
 
             if (!Accessed)