[NTOS:MM][NTOS:CC] Rewrite some cache memory management functions (#7510)

Use section object pointer with byte offset instead of using base address. This simplifies the Mm functions themselves and also the code in Cc that calls them.
Also add minor fixes for MmFlushSegment and MmPurgeSegment.

ntoskrnl/cc/copy.c

@@ -653,7 +653,9 @@ CcCopyWrite (
             CurrentOffset += VacbLength;
 
             /* Tell Mm */
-            Status = MmMakePagesDirty(NULL, Add2Ptr(Vacb->BaseAddress, VacbOffset), VacbLength);
+            Status = MmMakeSegmentDirty(FileObject->SectionObjectPointer,
+                                        Vacb->FileOffset.QuadPart + VacbOffset,
+                                        VacbLength);
             if (!NT_SUCCESS(Status))
                 ExRaiseStatus(Status);
         }
@@ -913,7 +915,9 @@ CcZeroData (
             Length -= VacbLength;
 
             /* Tell Mm */
-            Status = MmMakePagesDirty(NULL, Add2Ptr(Vacb->BaseAddress, VacbOffset), VacbLength);
+            Status = MmMakeSegmentDirty(FileObject->SectionObjectPointer,
+                                        Vacb->FileOffset.QuadPart + VacbOffset,
+                                        VacbLength);
             if (!NT_SUCCESS(Status))
                 ExRaiseStatus(Status);
         }

ntoskrnl/cc/pin.c

@@ -551,17 +551,18 @@ CcSetDirtyPinnedData (
     IN PLARGE_INTEGER Lsn)
 {
     PINTERNAL_BCB iBcb = CONTAINING_RECORD(Bcb, INTERNAL_BCB, PFCB);
+    PROS_VACB Vacb = iBcb->Vacb;
 
     CCTRACE(CC_API_DEBUG, "Bcb=%p Lsn=%p\n", Bcb, Lsn);
 
     /* Tell Mm */
-    MmMakePagesDirty(NULL,
+    MmMakeSegmentDirty(Vacb->SharedCacheMap->FileObject->SectionObjectPointer,
-                     Add2Ptr(iBcb->Vacb->BaseAddress, iBcb->PFCB.MappedFileOffset.QuadPart - iBcb->Vacb->FileOffset.QuadPart),
+                       iBcb->PFCB.MappedFileOffset.QuadPart,
-                     iBcb->PFCB.MappedLength);
+                       iBcb->PFCB.MappedLength);
 
-    if (!iBcb->Vacb->Dirty)
+    if (!Vacb->Dirty)
     {
-        CcRosMarkDirtyVacb(iBcb->Vacb);
+        CcRosMarkDirtyVacb(Vacb);
     }
 }
 

ntoskrnl/cc/view.c

@@ -936,22 +936,22 @@ CcRosEnsureVacbResident(
     _In_ ULONG Length
 )
 {
-    PVOID BaseAddress;
+    PROS_SHARED_CACHE_MAP SharedCacheMap = Vacb->SharedCacheMap;
 
     ASSERT((Offset + Length) <= VACB_MAPPING_GRANULARITY);
 
 #if 0
-    if ((Vacb->FileOffset.QuadPart + Offset) > Vacb->SharedCacheMap->SectionSize.QuadPart)
+    if ((Vacb->FileOffset.QuadPart + Offset) > SharedCacheMap->SectionSize.QuadPart)
     {
         DPRINT1("Vacb read beyond the file size!\n");
         return FALSE;
     }
 #endif
 
-    BaseAddress = (PVOID)((ULONG_PTR)Vacb->BaseAddress + Offset);
-
     /* Check if the pages are resident */
-    if (!MmArePagesResident(NULL, BaseAddress, Length))
+    if (!MmIsDataSectionResident(SharedCacheMap->FileObject->SectionObjectPointer,
+                                 Vacb->FileOffset.QuadPart + Offset,
+                                 Length))
     {
         if (!Wait)
         {
@@ -960,7 +960,6 @@ CcRosEnsureVacbResident(
 
         if (!NoRead)
         {
-            PROS_SHARED_CACHE_MAP SharedCacheMap = Vacb->SharedCacheMap;
             NTSTATUS Status = MmMakeDataSectionResident(SharedCacheMap->FileObject->SectionObjectPointer,
                                                         Vacb->FileOffset.QuadPart + Offset,
                                                         Length,

ntoskrnl/include/internal/mm.h

@@ -1502,16 +1502,16 @@ MmMapViewInSystemSpaceEx(
 
 BOOLEAN
 NTAPI
-MmArePagesResident(
+MmIsDataSectionResident(
-    _In_ PEPROCESS Process,
+    _In_ PSECTION_OBJECT_POINTERS SectionObjectPointer,
-    _In_ PVOID BaseAddress,
+    _In_ LONGLONG Offset,
     _In_ ULONG Length);
 
 NTSTATUS
 NTAPI
-MmMakePagesDirty(
+MmMakeSegmentDirty(
-    _In_ PEPROCESS Process,
+    _In_ PSECTION_OBJECT_POINTERS SectionObjectPointer,
-    _In_ PVOID Address,
+    _In_ LONGLONG Offset,
     _In_ ULONG Length);
 
 NTSTATUS

ntoskrnl/mm/section.c

@@ -4775,10 +4775,13 @@ MmCreateSection (OUT PVOID  * Section,
     return Status;
 }
 
+/* This function is not used. It is left for future use, when per-process
+ * address space is considered. */
+#if 0
 BOOLEAN
 NTAPI
 MmArePagesResident(
-    _In_ PEPROCESS Process,
+    _In_opt_ PEPROCESS Process,
     _In_ PVOID Address,
     _In_ ULONG Length)
 {
@@ -4826,6 +4829,7 @@ MmArePagesResident(
     MmUnlockAddressSpace(AddressSpace);
     return Ret;
 }
+#endif
 
 /* Like CcPurgeCache but for the in-memory segment */
 BOOLEAN
@@ -4859,9 +4863,9 @@ MmPurgeSegment(
         /* We must calculate the length for ourselves */
         /* FIXME: All of this is suboptimal */
         ULONG ElemCount = RtlNumberGenericTableElements(&Segment->PageTable);
-        /* No page. Nothing to purge */
         if (!ElemCount)
         {
+            /* No page. Nothing to purge */
             MmUnlockSectionSegment(Segment);
             MmDereferenceSegment(Segment);
             return TRUE;
@@ -4871,6 +4875,9 @@ MmPurgeSegment(
         PurgeEnd.QuadPart = PageTable->FileOffset.QuadPart + _countof(PageTable->PageEntries) * PAGE_SIZE;
     }
 
+    /* Find byte offset of the page to start */
+    PurgeStart.QuadPart = PAGE_ROUND_DOWN(PurgeStart.QuadPart);
+
     while (PurgeStart.QuadPart < PurgeEnd.QuadPart)
     {
         ULONG_PTR Entry = MmGetPageEntrySectionSegment(Segment, &PurgeStart);
@@ -4920,6 +4927,48 @@ MmPurgeSegment(
     return TRUE;
 }
 
+BOOLEAN
+NTAPI
+MmIsDataSectionResident(
+    _In_ PSECTION_OBJECT_POINTERS SectionObjectPointer,
+    _In_ LONGLONG Offset,
+    _In_ ULONG Length)
+{
+    PMM_SECTION_SEGMENT Segment;
+    LARGE_INTEGER RangeStart, RangeEnd;
+    BOOLEAN Ret = TRUE;
+
+    RangeStart.QuadPart = Offset;
+    if (!NT_SUCCESS(RtlLongLongAdd(RangeStart.QuadPart, Length, &RangeEnd.QuadPart)))
+        return FALSE;
+
+    Segment = MiGrabDataSection(SectionObjectPointer);
+    if (!Segment)
+        return FALSE;
+
+    /* Find byte offset of the page to start */
+    RangeStart.QuadPart = PAGE_ROUND_DOWN(RangeStart.QuadPart);
+
+    MmLockSectionSegment(Segment);
+
+    while (RangeStart.QuadPart < RangeEnd.QuadPart)
+    {
+        ULONG_PTR Entry = MmGetPageEntrySectionSegment(Segment, &RangeStart);
+        if ((Entry == 0) || IS_SWAP_FROM_SSE(Entry))
+        {
+            Ret = FALSE;
+            break;
+        }
+
+        RangeStart.QuadPart += PAGE_SIZE;
+    }
+
+    MmUnlockSectionSegment(Segment);
+    MmDereferenceSegment(Segment);
+
+    return Ret;
+}
+
 NTSTATUS
 NTAPI
 MmMakeDataSectionResident(
@@ -4940,6 +4989,63 @@ MmMakeDataSectionResident(
     return Status;
 }
 
+NTSTATUS
+NTAPI
+MmMakeSegmentDirty(
+    _In_ PSECTION_OBJECT_POINTERS SectionObjectPointer,
+    _In_ LONGLONG Offset,
+    _In_ ULONG Length)
+{
+    PMM_SECTION_SEGMENT Segment;
+    LARGE_INTEGER RangeStart, RangeEnd;
+    NTSTATUS Status;
+
+    RangeStart.QuadPart = Offset;
+    Status = RtlLongLongAdd(RangeStart.QuadPart, Length, &RangeEnd.QuadPart);
+    if (!NT_SUCCESS(Status))
+        return Status;
+
+    Segment = MiGrabDataSection(SectionObjectPointer);
+    if (!Segment)
+        return STATUS_NOT_MAPPED_VIEW;
+
+    /* Find byte offset of the page to start */
+    RangeStart.QuadPart = PAGE_ROUND_DOWN(RangeStart.QuadPart);
+
+    MmLockSectionSegment(Segment);
+
+    while (RangeStart.QuadPart < RangeEnd.QuadPart)
+    {
+        ULONG_PTR Entry = MmGetPageEntrySectionSegment(Segment, &RangeStart);
+
+        /* Let any pending read proceed */
+        while (MM_IS_WAIT_PTE(Entry))
+        {
+            MmUnlockSectionSegment(Segment);
+            KeDelayExecutionThread(KernelMode, FALSE, &TinyTime);
+            MmLockSectionSegment(Segment);
+            Entry = MmGetPageEntrySectionSegment(Segment, &RangeStart);
+        }
+
+        /* We are called from Cc, this can't be backed by the page files */
+        ASSERT(!IS_SWAP_FROM_SSE(Entry));
+
+        /* If there is no page there, there is nothing to make dirty */
+        if (Entry != 0)
+        {
+            /* Dirtify the entry */
+            MmSetPageEntrySectionSegment(Segment, &RangeStart, DIRTY_SSE(Entry));
+        }
+
+        RangeStart.QuadPart += PAGE_SIZE;
+    }
+
+    MmUnlockSectionSegment(Segment);
+    MmDereferenceSegment(Segment);
+
+    return STATUS_SUCCESS;
+}
+
 NTSTATUS
 NTAPI
 MmFlushSegment(
@@ -4991,8 +5097,8 @@ MmFlushSegment(
         FlushEnd.QuadPart = PageTable->FileOffset.QuadPart + _countof(PageTable->PageEntries) * PAGE_SIZE;
     }
 
-    FlushStart.QuadPart >>= PAGE_SHIFT;
+    /* Find byte offset of the page to start */
-    FlushStart.QuadPart <<= PAGE_SHIFT;
+    FlushStart.QuadPart = PAGE_ROUND_DOWN(FlushStart.QuadPart);
 
     while (FlushStart.QuadPart < FlushEnd.QuadPart)
     {
@@ -5197,10 +5303,13 @@ MmCheckDirtySegment(
     return FALSE;
 }
 
+/* This function is not used. It is left for future use, when per-process
+ * address space is considered. */
+#if 0
 NTSTATUS
 NTAPI
 MmMakePagesDirty(
-    _In_ PEPROCESS Process,
+    _In_opt_ PEPROCESS Process,
     _In_ PVOID Address,
     _In_ ULONG Length)
 {
@@ -5267,6 +5376,7 @@ MmMakePagesDirty(
     MmUnlockAddressSpace(AddressSpace);
     return STATUS_SUCCESS;
 }
+#endif
 
 NTSTATUS
 NTAPI