[NTOS:MM] Implement proper refcounting of page tables on amd64

CORE-17552

ntoskrnl/mm/ARM3/mdlsup.c

@@ -248,6 +248,7 @@ MiMapLockedPagesInUserSpace(
 
         /* Acquire a share count */
         Pfn1 = MI_PFN_ELEMENT(PointerPde->u.Hard.PageFrameNumber);
+        DPRINT("Incrementing %p from %p\n", Pfn1, _ReturnAddress());
         OldIrql = MiAcquirePfnLock();
         Pfn1->u2.ShareCount++;
         MiReleasePfnLock(OldIrql);
@@ -330,9 +331,6 @@ MiUnmapLockedPagesInUserSpace(
         ASSERT(MiAddressToPte(PointerPte)->u.Hard.Valid == 1);
         ASSERT(PointerPte->u.Hard.Valid == 1);
 
-        /* Dereference the page */
-        MiDecrementPageTableReferences(BaseAddress);
-
         /* Invalidate it */
         MI_ERASE_PTE(PointerPte);
 
@@ -341,28 +339,17 @@ MiUnmapLockedPagesInUserSpace(
         PageTablePage = PointerPde->u.Hard.PageFrameNumber;
         MiDecrementShareCount(MiGetPfnEntry(PageTablePage), PageTablePage);
 
+        if (MiDecrementPageTableReferences(BaseAddress) == 0)
+        {
+            ASSERT(MiIsPteOnPdeBoundary(PointerPte + 1) || (NumberOfPages == 1));
+            MiDeletePde(PointerPde, Process);
+        }
+
         /* Next page */
         PointerPte++;
         NumberOfPages--;
         BaseAddress = (PVOID)((ULONG_PTR)BaseAddress + PAGE_SIZE);
         MdlPages++;
-
-        /* Moving to a new PDE? */
-        if (PointerPde != MiAddressToPde(BaseAddress))
-        {
-            /* See if we should delete it */
-            KeFlushProcessTb();
-            PointerPde = MiPteToPde(PointerPte - 1);
-            ASSERT(PointerPde->u.Hard.Valid == 1);
-            if (MiQueryPageTableReferences(BaseAddress) == 0)
-            {
-                ASSERT(PointerPde->u.Long != 0);
-                MiDeletePte(PointerPde,
-                            MiPteToAddress(PointerPde),
-                            Process,
-                            NULL);
-            }
-        }
     }
 
     KeFlushProcessTb();

ntoskrnl/mm/ARM3/miarm.h

@@ -1823,40 +1823,7 @@ MiReferenceUnusedPageAndBumpLockCount(IN PMMPFN Pfn1)
     }
 }
 
-FORCEINLINE
-VOID
-MiIncrementPageTableReferences(IN PVOID Address)
-{
-    PUSHORT RefCount;
 
-    RefCount = &MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)];
-
-    *RefCount += 1;
-    ASSERT(*RefCount <= PTE_PER_PAGE);
-}
-
-FORCEINLINE
-VOID
-MiDecrementPageTableReferences(IN PVOID Address)
-{
-    PUSHORT RefCount;
-
-    RefCount = &MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)];
-
-    *RefCount -= 1;
-    ASSERT(*RefCount < PTE_PER_PAGE);
-}
-
-FORCEINLINE
-USHORT
-MiQueryPageTableReferences(IN PVOID Address)
-{
-    PUSHORT RefCount;
-
-    RefCount = &MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)];
-
-    return *RefCount;
-}
 
 CODE_SEG("INIT")
 BOOLEAN
@@ -2484,8 +2451,190 @@ MiSynchronizeSystemPde(PMMPDE PointerPde)
 }
 #endif
 
+#if _MI_PAGING_LEVELS == 2
+FORCEINLINE
+USHORT
+MiIncrementPageTableReferences(IN PVOID Address)
+{
+    PUSHORT RefCount;
+
+    RefCount = &MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)];
+
+    *RefCount += 1;
+    ASSERT(*RefCount <= PTE_PER_PAGE);
+    return *RefCount;
+}
+
+FORCEINLINE
+USHORT
+MiDecrementPageTableReferences(IN PVOID Address)
+{
+    PUSHORT RefCount;
+
+    RefCount = &MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)];
+
+    *RefCount -= 1;
+    ASSERT(*RefCount < PTE_PER_PAGE);
+    return *RefCount;
+}
+
+FORCEINLINE
+USHORT
+MiQueryPageTableReferences(IN PVOID Address)
+{
+    PUSHORT RefCount;
+
+    RefCount = &MmWorkingSetList->UsedPageTableEntries[MiGetPdeOffset(Address)];
+
+    return *RefCount;
+}
+#else
+FORCEINLINE
+USHORT
+MiIncrementPageTableReferences(IN PVOID Address)
+{
+    PMMPDE PointerPde = MiAddressToPde(Address);
+    PMMPFN Pfn;
+
+    /* We should not tinker with this one. */
+    ASSERT(PointerPde != (PMMPDE)PXE_SELFMAP);
+    DPRINT("Incrementing %p from %p\n", Address, _ReturnAddress());
+
+    /* Make sure we're locked */
+    ASSERT(PsGetCurrentThread()->OwnsProcessWorkingSetExclusive);
+
+    /* If we're bumping refcount, then it must be valid! */
+    ASSERT(PointerPde->u.Hard.Valid == 1);
+
+    /* This lies on the PFN */
+    Pfn = MiGetPfnEntry(PFN_FROM_PDE(PointerPde));
+    Pfn->OriginalPte.u.Soft.UsedPageTableEntries++;
+
+    ASSERT(Pfn->OriginalPte.u.Soft.UsedPageTableEntries <= PTE_PER_PAGE);
+
+    return Pfn->OriginalPte.u.Soft.UsedPageTableEntries;
+}
+
+FORCEINLINE
+USHORT
+MiDecrementPageTableReferences(IN PVOID Address)
+{
+    PMMPDE PointerPde = MiAddressToPde(Address);
+    PMMPFN Pfn;
+
+    /* We should not tinker with this one. */
+    ASSERT(PointerPde != (PMMPDE)PXE_SELFMAP);
+
+    DPRINT("Decrementing %p from %p\n", PointerPde, _ReturnAddress());
+
+    /* Make sure we're locked */
+    ASSERT(PsGetCurrentThread()->OwnsProcessWorkingSetExclusive);
+
+    /* If we're decreasing refcount, then it must be valid! */
+    ASSERT(PointerPde->u.Hard.Valid == 1);
+
+    /* This lies on the PFN */
+    Pfn = MiGetPfnEntry(PFN_FROM_PDE(PointerPde));
+
+    ASSERT(Pfn->OriginalPte.u.Soft.UsedPageTableEntries != 0);
+    Pfn->OriginalPte.u.Soft.UsedPageTableEntries--;
+
+    ASSERT(Pfn->OriginalPte.u.Soft.UsedPageTableEntries < PTE_PER_PAGE);
+
+    return Pfn->OriginalPte.u.Soft.UsedPageTableEntries;
+}
+
+FORCEINLINE
+USHORT
+MiQueryPageTableReferences(IN PVOID Address)
+{
+    PMMPDE PointerPde;
+    PMMPPE PointerPpe;
+#if _MI_PAGING_LEVELS == 4
+    PMMPXE PointerPxe;
+#endif
+    PMMPFN Pfn;
+
+    /* Make sure we're locked */
+    ASSERT((PsGetCurrentThread()->OwnsProcessWorkingSetExclusive) || (PsGetCurrentThread()->OwnsProcessWorkingSetShared));
+
+    /* Check if PXE or PPE have references first. */
+#if _MI_PAGING_LEVELS == 4
+    PointerPxe = MiAddressToPxe(Address);
+    if ((PointerPxe->u.Hard.Valid == 1) || (PointerPxe->u.Soft.Transition == 1))
+    {
+        Pfn = MiGetPfnEntry(PFN_FROM_PXE(PointerPxe));
+        if (Pfn->OriginalPte.u.Soft.UsedPageTableEntries == 0)
+            return 0;
+    }
+    else if (PointerPxe->u.Soft.UsedPageTableEntries == 0)
+    {
+        return 0;
+    }
+
+    if (PointerPxe->u.Hard.Valid == 0)
+    {
+        MiMakeSystemAddressValid(MiPteToAddress(PointerPxe), PsGetCurrentProcess());
+    }
+#endif
+
+    PointerPpe = MiAddressToPpe(Address);
+    if ((PointerPpe->u.Hard.Valid == 1) || (PointerPpe->u.Soft.Transition == 1))
+    {
+        Pfn = MiGetPfnEntry(PFN_FROM_PPE(PointerPpe));
+        if (Pfn->OriginalPte.u.Soft.UsedPageTableEntries == 0)
+            return 0;
+    }
+    else if (PointerPpe->u.Soft.UsedPageTableEntries == 0)
+    {
+        return 0;
+    }
+
+    if (PointerPpe->u.Hard.Valid == 0)
+    {
+        MiMakeSystemAddressValid(MiPteToAddress(PointerPpe), PsGetCurrentProcess());
+    }
+
+    PointerPde = MiAddressToPde(Address);
+    if ((PointerPde->u.Hard.Valid == 0) && (PointerPde->u.Soft.Transition == 0))
+    {
+        return PointerPde->u.Soft.UsedPageTableEntries;
+    }
+
+    /* This lies on the PFN */
+    Pfn = MiGetPfnEntry(PFN_FROM_PDE(PointerPde));
+    return Pfn->OriginalPte.u.Soft.UsedPageTableEntries;
+}
+#endif
+
 #ifdef __cplusplus
 } // extern "C"
 #endif
 
+FORCEINLINE
+VOID
+MiDeletePde(
+    _In_ PMMPDE PointerPde,
+    _In_ PEPROCESS CurrentProcess)
+{
+    /* Only for user-mode ones */
+    ASSERT(MiIsUserPde(PointerPde));
+
+    /* Kill this one as a PTE */
+    MiDeletePte((PMMPTE)PointerPde, MiPdeToPte(PointerPde), CurrentProcess, NULL);
+#if _MI_PAGING_LEVELS >= 3
+    /* Cascade down */
+    if (MiDecrementPageTableReferences(MiPdeToPte(PointerPde)) == 0)
+    {
+        MiDeletePte(MiPdeToPpe(PointerPde), PointerPde, CurrentProcess, NULL);
+#if _MI_PAGING_LEVELS == 4
+        if (MiDecrementPageTableReferences(PointerPde) == 0)
+        {
+            MiDeletePte(MiPdeToPxe(PointerPde), MiPdeToPpe(PointerPde), CurrentProcess, NULL);
+        }
+#endif
+    }
+#endif
+}
+
 /* EOF */

ntoskrnl/mm/ARM3/pagfault.c

@@ -2145,6 +2145,7 @@ UserFault:
 
         /* We should come back with a valid PPE */
         ASSERT(PointerPpe->u.Hard.Valid == 1);
+        MiIncrementPageTableReferences(PointerPde);
     }
 #endif
 
@@ -2184,6 +2185,10 @@ UserFault:
                                  MM_EXECUTE_READWRITE,
                                  CurrentProcess,
                                  MM_NOIRQL);
+#if _MI_PAGING_LEVELS >= 3
+        MiIncrementPageTableReferences(PointerPte);
+#endif
+
 #if MI_TRACE_PFNS
         UserPdeFault = FALSE;
 #endif

ntoskrnl/mm/ARM3/pfnlist.c

@@ -1027,6 +1027,8 @@ MiInitializePfn(IN PFN_NUMBER PageFrameIndex,
     ASSERT(PageFrameIndex != 0);
     Pfn1->u4.PteFrame = PageFrameIndex;
 
+    DPRINT("Incrementing share count of %lp from %p\n", PageFrameIndex, _ReturnAddress());
+
     /* Increase its share count so we don't get rid of it */
     Pfn1 = MI_PFN_ELEMENT(PageFrameIndex);
     Pfn1->u2.ShareCount++;

ntoskrnl/mm/ARM3/session.c

@@ -477,7 +477,7 @@ MiSessionInitializeWorkingSetList(VOID)
 
     /* Fill out the two pointers */
     MmSessionSpace->Vm.VmWorkingSetList = WorkingSetList;
-    MmSessionSpace->Wsle = (PMMWSLE)WorkingSetList->UsedPageTableEntries;
+    MmSessionSpace->Wsle = (PMMWSLE)((&WorkingSetList->VadBitMapHint) + 1);
 
     /* Get the PDE for the working set, and check if it's already allocated */
     PointerPde = MiAddressToPde(WorkingSetList);

ntoskrnl/mm/ARM3/virtual.c

@@ -727,18 +727,15 @@ MiDeleteVirtualAddresses(IN ULONG_PTR Va,
         /* Check remaining PTE count (go back 1 page due to above loop) */
         if (MiQueryPageTableReferences((PVOID)(Va - PAGE_SIZE)) == 0)
         {
-            if (PointerPde->u.Long != 0)
+            ASSERT(PointerPde->u.Long != 0);
-            {
+
-                /* Delete the PTE proper */
+            /* Delete the PDE proper */
-                MiDeletePte(PointerPde,
+            MiDeletePde(PointerPde, CurrentProcess);
-                            MiPteToAddress(PointerPde),
-                            CurrentProcess,
-                            NULL);
-            }
         }
 
-        /* Release the lock and get out if we're done */
+        /* Release the lock */
         MiReleasePfnLock(OldIrql);
+
         if (Va > EndingAddress) return;
 
         /* Otherwise, we exited because we hit a new PDE boundary, so start over */

ntoskrnl/mm/i386/page.c

@@ -238,11 +238,10 @@ MmDeleteVirtualMapping(PEPROCESS Process, PVOID Address,
     if (Address < MmSystemRangeStart)
     {
         /* Remove PDE reference */
-        MiDecrementPageTableReferences(Address);
+        if (MiDecrementPageTableReferences(Address) == 0)
-        if (MiQueryPageTableReferences(Address) == 0)
         {
             KIRQL OldIrql = MiAcquirePfnLock();
-            MiDeletePte(MiAddressToPte(PointerPte), PointerPte, Process, NULL);
+            MiDeletePde(MiAddressToPde(Address), Process);
             MiReleasePfnLock(OldIrql);
         }
 
@@ -293,8 +292,7 @@ MmDeletePageFileMapping(
     }
 
     /* This used to be a non-zero PTE, now we can let the PDE go. */
-    MiDecrementPageTableReferences(Address);
+    if (MiDecrementPageTableReferences(Address) == 0)
-    if (MiQueryPageTableReferences(Address) == 0)
     {
         /* We can let it go */
         KIRQL OldIrql = MiAcquirePfnLock();

sdk/include/ndk/mmtypes.h

@@ -879,8 +879,19 @@ typedef struct _MMWSL
     PVOID HighestPermittedHashAddress;
     ULONG NumberOfImageWaiters;
     ULONG VadBitMapHint;
+#ifndef _M_AMD64
     USHORT UsedPageTableEntries[768];
     ULONG CommittedPageTables[24];
+#else
+    VOID* HighestUserAddress;
+    ULONG MaximumUserPageTablePages;
+    ULONG MaximumUserPageDirectoryPages;
+    ULONG* CommittedPageTables;
+    ULONG NumberOfCommittedPageDirectories;
+    ULONG* CommittedPageDirectories;
+    ULONG NumberOfCommittedPageDirectoryParents;
+    ULONGLONG CommittedPageDirectoryParents[1];
+#endif
 } MMWSL, *PMMWSL;
 
 //