From ed1038ff9c41e640a04ad63b7c98f726b3f0ecad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A9r=C3=B4me=20Gardou?= <jerome.gardou@reactos.org>
Date: Tue, 25 May 2021 10:44:11 +0200
Subject: [PATCH] [NTOS:MM] Properly initialize initial process address space
 on amd64

By applying the same kludge on the relevant MMPFNs than on x86
---
 ntoskrnl/mm/ARM3/procsup.c |  7 -------
 ntoskrnl/mm/amd64/init.c   | 24 +++++++++++++++++++++++-
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/ntoskrnl/mm/ARM3/procsup.c b/ntoskrnl/mm/ARM3/procsup.c
index 2fc8dccae54..0d46aa9c05f 100644
--- a/ntoskrnl/mm/ARM3/procsup.c
+++ b/ntoskrnl/mm/ARM3/procsup.c
@@ -943,10 +943,6 @@ 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 our working set */
     MiLockProcessWorkingSet(Process, PsGetCurrentThread());
 
@@ -1014,9 +1010,6 @@ MmInitializeProcessAddressSpace(IN PEPROCESS Process,
 
     /* Release the process working set */
     MiUnlockProcessWorkingSet(Process, PsGetCurrentThread());
-#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 */
diff --git a/ntoskrnl/mm/amd64/init.c b/ntoskrnl/mm/amd64/init.c
index eea9cc2c1bb..0c6a868b144 100644
--- a/ntoskrnl/mm/amd64/init.c
+++ b/ntoskrnl/mm/amd64/init.c
@@ -252,6 +252,8 @@ MiInitializePageTable(VOID)
             RtlZeroMemory(MiPteToAddress(PointerPxe), PAGE_SIZE);
         }
     }
+    PxePfn = PFN_FROM_PXE(MiAddressToPxe((PVOID)HYPER_SPACE));
+    PsGetCurrentProcess()->Pcb.DirectoryTableBase[1] = PxePfn << PAGE_SHIFT;
 
     /* Map PPEs for paged pool */
     MiMapPPEs(MmPagedPoolStart, MmPagedPoolEnd);
@@ -713,7 +715,6 @@ MiInitMachineDependent(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
 //    PrototypePte.u.ProtoAddress
     PrototypePte.u.Soft.PageFileHigh = MI_PTE_LOOKUP_NEEDED;
 
-
     MiInitializePageTable();
 
     MiBuildNonPagedPool();
@@ -723,6 +724,27 @@ MiInitMachineDependent(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
     /* Map the PFN database pages */
     MiBuildPfnDatabase(LoaderBlock);
 
+    /* Reset the ref/share count so that MmInitializeProcessAddressSpace works */
+    PMMPFN Pfn = MiGetPfnEntry(PFN_FROM_PTE((PMMPTE)PXE_SELFMAP));
+    Pfn->u2.ShareCount = 0;
+    Pfn->u3.e2.ReferenceCount = 0;
+
+    Pfn = MiGetPfnEntry(PFN_FROM_PDE(MiAddressToPde((PVOID)HYPER_SPACE)));
+    Pfn->u2.ShareCount = 0;
+    Pfn->u3.e2.ReferenceCount = 0;
+
+    Pfn = MiGetPfnEntry(PFN_FROM_PPE(MiAddressToPpe((PVOID)HYPER_SPACE)));
+    Pfn->u2.ShareCount = 0;
+    Pfn->u3.e2.ReferenceCount = 0;
+
+    Pfn = MiGetPfnEntry(PFN_FROM_PXE(MiAddressToPxe((PVOID)HYPER_SPACE)));
+    Pfn->u2.ShareCount = 0;
+    Pfn->u3.e2.ReferenceCount = 0;
+
+    Pfn = MiGetPfnEntry(PFN_FROM_PTE(MiAddressToPte(MmWorkingSetList)));
+    Pfn->u2.ShareCount = 0;
+    Pfn->u3.e2.ReferenceCount = 0;
+
     /* Initialize the nonpaged pool */
     InitializePool(NonPagedPool, 0);