From 0726cfce1022af4dcfbedb210642fd74dc54e6d1 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Mon, 26 Oct 2020 11:12:04 +0100
Subject: [PATCH] [NTOS:MM:X64] Handle PPEs and PXEs as well in
 MmInitializeProcessAddressSpace()

But don't set up the PFNs for the initial process on x64, as these have already been set up.
---
 ntoskrnl/mm/ARM3/procsup.c | 36 ++++++++++++++++++++++++++++++------
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/ntoskrnl/mm/ARM3/procsup.c b/ntoskrnl/mm/ARM3/procsup.c
index 809f36e824d..e3aa8336bbd 100644
--- a/ntoskrnl/mm/ARM3/procsup.c
+++ b/ntoskrnl/mm/ARM3/procsup.c
@@ -951,6 +951,12 @@ MmInitializeProcessAddressSpace(IN PEPROCESS Process,
     PCHAR Destination;
     USHORT Length = 0;
     MMPTE TempPte;
+#if (_MI_PAGING_LEVELS >= 3)
+    PMMPPE PointerPpe;
+#endif
+#if (_MI_PAGING_LEVELS == 4)
+    PMMPXE PointerPxe;
+#endif
 
     /* We should have a PDE */
     ASSERT(Process->Pcb.DirectoryTableBase[0] != 0);
@@ -971,12 +977,19 @@ MmInitializeProcessAddressSpace(IN PEPROCESS Process,
     ASSERT(Process->VadRoot.NumberGenericTableElements == 0);
     Process->VadRoot.BalancedRoot.u1.Parent = &Process->VadRoot.BalancedRoot;
 
+#ifdef _M_AMD64
+    /* On x64 the PFNs for the initial process are already set up */
+    if (Process != &KiInitialProcess) {
+#endif
+
     /* Lock PFN database */
     OldIrql = MiAcquirePfnLock();
 
     /* Setup the PFN for the PDE base of this process */
-#ifdef _M_AMD64
+#if (_MI_PAGING_LEVELS == 4)
     PointerPte = MiAddressToPte(PXE_BASE);
+#elif (_MI_PAGING_LEVELS == 3)
+    PointerPte = MiAddressToPte(PPE_BASE);
 #else
     PointerPte = MiAddressToPte(PDE_BASE);
 #endif
@@ -985,15 +998,22 @@ MmInitializeProcessAddressSpace(IN PEPROCESS Process,
     MiInitializePfn(PageFrameNumber, PointerPte, TRUE);
 
     /* Do the same for hyperspace */
-#ifdef _M_AMD64
-    PointerPde = MiAddressToPxe((PVOID)HYPER_SPACE);
-#else
-    PointerPde = MiAddressToPde(HYPER_SPACE);
-#endif
+    PointerPde = MiAddressToPde((PVOID)HYPER_SPACE);
     PageFrameNumber = PFN_FROM_PTE(PointerPde);
     //ASSERT(Process->Pcb.DirectoryTableBase[0] == PageFrameNumber * PAGE_SIZE); // we're not lucky
     MiInitializePfn(PageFrameNumber, (PMMPTE)PointerPde, TRUE);
 
+#if (_MI_PAGING_LEVELS >= 3)
+    PointerPpe = MiAddressToPpe((PVOID)HYPER_SPACE);
+    PageFrameNumber = PFN_FROM_PTE(PointerPpe);
+    MiInitializePfn(PageFrameNumber, PointerPpe, TRUE);
+#endif
+#if (_MI_PAGING_LEVELS == 4)
+    PointerPxe = MiAddressToPxe((PVOID)HYPER_SPACE);
+    PageFrameNumber = PFN_FROM_PTE(PointerPxe);
+    MiInitializePfn(PageFrameNumber, PointerPxe, TRUE);
+#endif
+
     /* Setup the PFN for the PTE for the working set */
     PointerPte = MiAddressToPte(MI_WORKING_SET_LIST);
     MI_MAKE_HARDWARE_PTE(&TempPte, PointerPte, MM_READWRITE, 0);
@@ -1013,6 +1033,10 @@ MmInitializeProcessAddressSpace(IN PEPROCESS Process,
     /* Release PFN lock */
     MiReleasePfnLock(OldIrql);
 
+#ifdef _M_AMD64
+   } /* On x64 the PFNs for the initial process are already set up */
+#endif
+
 #ifdef _M_AMD64
     /* On x64 we need a VAD for the shared user page */
     Status = MiInsertSharedUserPageVad();