[NTOS:MM] Fix PFN tracing

ntoskrnl/cc/view.c

@@ -806,25 +806,6 @@ Retry:
     InsertTailList(&VacbLruListHead, &current->VacbLruListEntry);
     KeReleaseQueuedSpinLock(LockQueueMasterLock, oldIrql);
 
-    MI_SET_USAGE(MI_USAGE_CACHE);
-#if MI_TRACE_PFNS
-    if ((SharedCacheMap->FileObject) && (SharedCacheMap->FileObject->FileName.Buffer))
-    {
-        PWCHAR pos;
-        ULONG len = 0;
-        pos = wcsrchr(SharedCacheMap->FileObject->FileName.Buffer, '\\');
-        if (pos)
-        {
-            len = wcslen(pos) * sizeof(WCHAR);
-            snprintf(MI_PFN_CURRENT_PROCESS_NAME, min(16, len), "%S", pos);
-        }
-        else
-        {
-            snprintf(MI_PFN_CURRENT_PROCESS_NAME, min(16, len), "%wZ", &SharedCacheMap->FileObject->FileName);
-        }
-    }
-#endif
-
     /* Reference it to allow release */
     CcRosVacbIncRefCount(current);
 

ntoskrnl/include/internal/mm.h

@@ -245,9 +245,45 @@ MM_RMAP_ENTRY, *PMM_RMAP_ENTRY;
 extern ULONG MI_PFN_CURRENT_USAGE;
 extern CHAR MI_PFN_CURRENT_PROCESS_NAME[16];
 #define MI_SET_USAGE(x)     MI_PFN_CURRENT_USAGE = x
-#define MI_SET_PROCESS2(x)  memcpy(MI_PFN_CURRENT_PROCESS_NAME, x, 16)
+#define MI_SET_PROCESS2(x)  memcpy(MI_PFN_CURRENT_PROCESS_NAME, x, min(sizeof(x), sizeof(MI_PFN_CURRENT_PROCESS_NAME)))
+FORCEINLINE
+void
+MI_SET_PROCESS(PEPROCESS Process)
+{
+    if (!Process)
+        MI_SET_PROCESS2("Kernel");
+    else if (Process == (PEPROCESS)1)
+        MI_SET_PROCESS2("Hydra");
+    else
+        MI_SET_PROCESS2(Process->ImageFileName);
+}
+
+FORCEINLINE
+void
+MI_SET_PROCESS_USTR(PUNICODE_STRING ustr)
+{
+    PWSTR pos, strEnd;
+    int i;
+
+    if (!ustr->Buffer || ustr->Length == 0)
+    {
+        MI_PFN_CURRENT_PROCESS_NAME[0] = 0;
+        return;
+    }
+
+    pos = strEnd = &ustr->Buffer[ustr->Length / sizeof(WCHAR)];
+    while ((*pos != L'\\') && (pos >  ustr->Buffer))
+        pos--;
+
+    if (*pos == L'\\')
+        pos++;
+
+    for (i = 0; i < sizeof(MI_PFN_CURRENT_PROCESS_NAME) && pos <= strEnd; i++, pos++)
+        MI_PFN_CURRENT_PROCESS_NAME[i] = (CHAR)*pos;
+}
 #else
 #define MI_SET_USAGE(x)
+#define MI_SET_PROCESS(x)
 #define MI_SET_PROCESS2(x)
 #endif
 
@@ -275,6 +311,9 @@ typedef enum _MI_PFN_USAGES
     MI_USAGE_PFN_DATABASE,
     MI_USAGE_BOOT_DRIVER,
     MI_USAGE_INIT_MEMORY,
+    MI_USAGE_PAGE_FILE,
+    MI_USAGE_COW,
+    MI_USAGE_WSLE,
     MI_USAGE_FREE_PAGE
 } MI_PFN_USAGES;
 
@@ -355,6 +394,7 @@ typedef struct _MMPFN
 #if MI_TRACE_PFNS
     MI_PFN_USAGES PfnUsage;
     CHAR ProcessName[16];
+#define MI_SET_PFN_PROCESS_NAME(pfn, x) memcpy(pfn->ProcessName, x, min(sizeof(x), sizeof(pfn->ProcessName)))
 #endif
 
     // HACK until WS lists are supported

ntoskrnl/mm/ARM3/mminit.c

@@ -801,7 +801,7 @@ MiBuildPfnDatabaseFromPages(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
                 Pfn1->u3.e1.CacheAttribute = MiNonCached;
 #if MI_TRACE_PFNS
                 Pfn1->PfnUsage = MI_USAGE_INIT_MEMORY;
-                memcpy(Pfn1->ProcessName, "Initial PDE", 16);
+                MI_SET_PFN_PROCESS_NAME(Pfn1, "Initial PDE");
 #endif
             }
             else
@@ -848,7 +848,7 @@ MiBuildPfnDatabaseFromPages(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
                                 Pfn2->u3.e1.CacheAttribute = MiNonCached;
 #if MI_TRACE_PFNS
                                 Pfn2->PfnUsage = MI_USAGE_INIT_MEMORY;
-                                memcpy(Pfn1->ProcessName, "Initial PTE", 16);
+                                MI_SET_PFN_PROCESS_NAME(Pfn2, "Initial PTE");
 #endif
                             }
                         }

ntoskrnl/mm/ARM3/pagfault.c

@@ -899,6 +899,9 @@ MiResolvePageFileFault(_In_ BOOLEAN StoreInstruction,
     ASSERT(CurrentProcess > HYDRA_PROCESS);
     ASSERT(*OldIrql != MM_NOIRQL);
 
+    MI_SET_USAGE(MI_USAGE_PAGE_FILE);
+    MI_SET_PROCESS(CurrentProcess);
+
     /* We must hold the PFN lock */
     MI_ASSERT_PFN_LOCK_HELD();
 
@@ -1210,6 +1213,9 @@ MiResolveProtoPteFault(IN BOOLEAN StoreInstruction,
         ASSERT(TempPte.u.Hard.Valid == 1);
         ProtoPageFrameIndex = PFN_FROM_PTE(&TempPte);
 
+        MI_SET_USAGE(MI_USAGE_COW);
+        MI_SET_PROCESS(Process);
+
         /* Get a new page for the private copy */
         if (Process > HYDRA_PROCESS)
             Color = MI_GET_NEXT_PROCESS_COLOR(Process);
@@ -2210,6 +2216,9 @@ UserFault:
 
                 ASSERT(MmAvailablePages > 0);
 
+                MI_SET_USAGE(MI_USAGE_COW);
+                MI_SET_PROCESS(CurrentProcess);
+
                 /* Allocate a new page and copy it */
                 PageFrameIndex = MiRemoveAnyPage(MI_GET_NEXT_PROCESS_COLOR(CurrentProcess));
                 OldPageFrameIndex = PFN_FROM_PTE(&TempPte);

ntoskrnl/mm/ARM3/pfnlist.c

@@ -254,8 +254,8 @@ MiUnlinkFreeOrZeroedPage(IN PMMPFN Entry)
     ASSERT(MI_PFN_CURRENT_USAGE != MI_USAGE_NOT_SET);
     Entry->PfnUsage = MI_PFN_CURRENT_USAGE;
     memcpy(Entry->ProcessName, MI_PFN_CURRENT_PROCESS_NAME, 16);
-//    MI_PFN_CURRENT_USAGE = MI_USAGE_NOT_SET;
-//    memcpy(MI_PFN_CURRENT_PROCESS_NAME, "Not Set", 16);
+    MI_PFN_CURRENT_USAGE = MI_USAGE_NOT_SET;
+    MI_SET_PROCESS2("Not Set");
 #endif
 }
 
@@ -459,11 +459,11 @@ MiRemovePageByColor(IN PFN_NUMBER PageIndex,
     MiDecrementAvailablePages();
 
 #if MI_TRACE_PFNS
-    //ASSERT(MI_PFN_CURRENT_USAGE != MI_USAGE_NOT_SET);
+    ASSERT(MI_PFN_CURRENT_USAGE != MI_USAGE_NOT_SET);
     Pfn1->PfnUsage = MI_PFN_CURRENT_USAGE;
     memcpy(Pfn1->ProcessName, MI_PFN_CURRENT_PROCESS_NAME, 16);
-    //MI_PFN_CURRENT_USAGE = MI_USAGE_NOT_SET;
-    //memcpy(MI_PFN_CURRENT_PROCESS_NAME, "Not Set", 16);
+    MI_PFN_CURRENT_USAGE = MI_USAGE_NOT_SET;
+    MI_SET_PROCESS2("Not Set");
 #endif
 
     /* Return the page */
@@ -937,9 +937,8 @@ MiInsertPageInList(IN PMMPFNLIST ListHead,
         ColorHead->Count++;
 
 #if MI_TRACE_PFNS
-            //ASSERT(MI_PFN_CURRENT_USAGE == MI_USAGE_NOT_SET);
+            ASSERT(MI_PFN_CURRENT_USAGE == MI_USAGE_NOT_SET);
             Pfn1->PfnUsage = MI_USAGE_FREE_PAGE;
-            MI_PFN_CURRENT_USAGE = MI_USAGE_NOT_SET;
             RtlZeroMemory(Pfn1->ProcessName, 16);
 #endif
     }

ntoskrnl/mm/ARM3/sysldr.c

@@ -188,14 +188,7 @@ MiLoadImageSection(_Inout_ PSECTION *SectionPtr,
         /* Some debug stuff */
         MI_SET_USAGE(MI_USAGE_DRIVER_PAGE);
 #if MI_TRACE_PFNS
-        if (FileName->Buffer)
-        {
-            PWCHAR pos = NULL;
-            ULONG len = 0;
-            pos = wcsrchr(FileName->Buffer, '\\');
-            len = wcslen(pos) * sizeof(WCHAR);
-            if (pos) snprintf(MI_PFN_CURRENT_PROCESS_NAME, min(16, len), "%S", pos);
-        }
+        MI_SET_PROCESS_USTR(FileName);
 #endif
 
         /* Grab a page */

ntoskrnl/mm/freelist.c

@@ -610,6 +610,21 @@ MmAllocPage(ULONG Type)
 
     OldIrql = MiAcquirePfnLock();
 
+#if MI_TRACE_PFNS
+    switch(Type)
+    {
+    case MC_CACHE:
+    case MC_SYSTEM:
+        MI_SET_USAGE(MI_USAGE_CACHE);
+        break;
+    case MC_USER:
+        MI_SET_USAGE(MI_USAGE_SECTION);
+        break;
+    default:
+        ASSERT(FALSE);
+    }
+#endif
+
     PfnOffset = MiRemoveZeroPage(MI_GET_NEXT_COLOR());
     if (!PfnOffset)
     {

ntoskrnl/mm/section.c

@@ -1803,9 +1803,6 @@ MmNotPresentFaultSectionView(PMMSUPPORT AddressSpace,
         MmUnlockSectionSegment(Segment);
 
         MmUnlockAddressSpace(AddressSpace);
-        MI_SET_USAGE(MI_USAGE_SECTION);
-        if (Process) MI_SET_PROCESS2(Process->ImageFileName);
-        if (!Process) MI_SET_PROCESS2("Kernel Section");
         Status = MmRequestPageMemoryConsumer(MC_USER, TRUE, &Page);
         if (!NT_SUCCESS(Status))
         {
@@ -1975,9 +1972,6 @@ MmAccessFaultSectionView(PMMSUPPORT AddressSpace,
     /*
      * Allocate a page
      */
-    MI_SET_USAGE(MI_USAGE_SECTION);
-    if (Process) MI_SET_PROCESS2(Process->ImageFileName);
-    if (!Process) MI_SET_PROCESS2("Kernel Section");
     Status = MmRequestPageMemoryConsumer(MC_USER, TRUE, &NewPage);
     if (!NT_SUCCESS(Status))
     {