[NTOS:Mm] Rewrite MiWriteProtectSystemImage (#749)

* The previous version was overcomplicated and broken and therefore disabled.
* The new version also enforces NX protection on x64.
* Now that protecting works, also protect the boot loaded images.

ntoskrnl/include/internal/amd64/mm.h

@@ -4,6 +4,7 @@
 #pragma once
 
 #define _MI_PAGING_LEVELS 4
+#define _MI_HAS_NO_EXECUTE 1
 
 /* Memory layout base addresses (This is based on Vista!) */
 #define MI_USER_PROBE_ADDRESS           (PVOID)0x000007FFFFFF0000ULL

ntoskrnl/include/internal/arm/mm.h

@@ -4,6 +4,7 @@
 #pragma once
 
 #define _MI_PAGING_LEVELS 2
+#define _MI_HAS_NO_EXECUTE 1
 
 /* Memory layout base addresses */
 #define MI_USER_PROBE_ADDRESS                   (PVOID)0x7FFF0000

ntoskrnl/include/internal/i386/mm.h

@@ -5,8 +5,10 @@
 
 #ifdef _PAE_
 #define _MI_PAGING_LEVELS 3
+#define _MI_HAS_NO_EXECUTE 1
 #else
 #define _MI_PAGING_LEVELS 2
+#define _MI_NO_EXECUTE 0
 #endif
 
 /* Memory layout base addresses */

ntoskrnl/include/ntoskrnl.h

@@ -94,6 +94,14 @@
 
 #define ExRaiseStatus RtlRaiseStatus
 
+/* Also defined in fltkernel.h, but we don't want the entire header */
+#ifndef Add2Ptr
+#define Add2Ptr(P,I) ((PVOID)((PUCHAR)(P) + (I)))
+#endif
+#ifndef PtrOffset
+#define PtrOffset(B,O) ((ULONG)((ULONG_PTR)(O) - (ULONG_PTR)(B)))
+#endif
+
 //
 // Switch for enabling global page support
 //

ntoskrnl/mm/ARM3/miarm.h

@@ -83,7 +83,7 @@ C_ASSERT(SYSTEM_PD_SIZE == PAGE_SIZE);
 // while on certain architectures such as ARM, it is enabling the cache which
 // requires a flag.
 //
-#if defined(_M_IX86) || defined(_M_AMD64)
+#if defined(_M_IX86)
 //
 // Access Flags
 //
@@ -109,6 +109,34 @@ C_ASSERT(SYSTEM_PD_SIZE == PAGE_SIZE);
 #define PTE_ENABLE_CACHE        0
 #define PTE_DISABLE_CACHE       0x10
 #define PTE_WRITECOMBINED_CACHE 0x10
+#define PTE_PROTECT_MASK        0x612
+#elif defined(_M_AMD64)
+//
+// Access Flags
+//
+#define PTE_READONLY            0x8000000000000000ULL
+#define PTE_EXECUTE             0x0000000000000000ULL
+#define PTE_EXECUTE_READ        PTE_EXECUTE /* EXECUTE implies READ on x64 */
+#define PTE_READWRITE           0x8000000000000002ULL
+#define PTE_WRITECOPY           0x8000000000000200ULL
+#define PTE_EXECUTE_READWRITE   0x0000000000000002ULL
+#define PTE_EXECUTE_WRITECOPY   0x0000000000000200ULL
+#define PTE_PROTOTYPE           0x0000000000000400ULL
+
+//
+// State Flags
+//
+#define PTE_VALID               0x0000000000000001ULL
+#define PTE_ACCESSED            0x0000000000000020ULL
+#define PTE_DIRTY               0x0000000000000040ULL
+
+//
+// Cache flags
+//
+#define PTE_ENABLE_CACHE        0x0000000000000000ULL
+#define PTE_DISABLE_CACHE       0x0000000000000010ULL
+#define PTE_WRITECOMBINED_CACHE 0x0000000000000010ULL
+#define PTE_PROTECT_MASK        0x8000000000000612ULL
 #elif defined(_M_ARM)
 #define PTE_READONLY            0x200
 #define PTE_EXECUTE             0 // Not worrying about NX yet
@@ -118,16 +146,23 @@ C_ASSERT(SYSTEM_PD_SIZE == PAGE_SIZE);
 #define PTE_EXECUTE_READWRITE   0 // Not worrying about NX yet
 #define PTE_EXECUTE_WRITECOPY   0 // Not worrying about NX yet
 #define PTE_PROTOTYPE           0x400 // Using the Shared bit
+
 //
 // Cache flags
 //
 #define PTE_ENABLE_CACHE        0
 #define PTE_DISABLE_CACHE       0x10
 #define PTE_WRITECOMBINED_CACHE 0x10
+#define PTE_PROTECT_MASK        0x610
 #else
 #error Define these please!
 #endif
 
+//
+// Mask for image section page protection
+//
+#define IMAGE_SCN_PROTECTION_MASK (IMAGE_SCN_MEM_WRITE | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_EXECUTE)
+
 extern const ULONG_PTR MmProtectToPteMask[32];
 extern const ULONG MmProtectToValue[32];
 
@@ -2263,6 +2298,11 @@ MiMakePdeExistAndMakeValid(
     IN KIRQL OldIrql
 );
 
+VOID
+NTAPI
+MiWriteProtectSystemImage(
+    _In_ PVOID ImageBase);
+
 //
 // MiRemoveZeroPage will use inline code to zero out the page manually if only
 // free pages are available. In some scenarios, we don't/can't run that piece of

ntoskrnl/mm/ARM3/sysldr.c

@@ -2341,252 +2341,146 @@ MiUseLargeDriverPage(IN ULONG NumberOfPtes,
     return FALSE;
 }
 
-ULONG
+VOID
 NTAPI
-MiComputeDriverProtection(IN BOOLEAN SessionSpace,
+MiSetSystemCodeProtection(
-                          IN ULONG SectionProtection)
+    _In_ PMMPTE FirstPte,
+    _In_ PMMPTE LastPte,
+    _In_ ULONG Protection)
 {
-    ULONG Protection = MM_ZERO_ACCESS;
+    PMMPTE PointerPte;
+    MMPTE TempPte;
 
-    /* Check if the caller gave anything */
+    /* Loop the PTEs */
-    if (SectionProtection)
+    for (PointerPte = FirstPte; PointerPte <= LastPte; PointerPte++)
     {
-        /* Always turn on execute access */
+        /* Read the PTE */
-        SectionProtection |= IMAGE_SCN_MEM_EXECUTE;
+        TempPte = *PointerPte;
+
+        /* Make sure it's valid */
+        ASSERT(TempPte.u.Hard.Valid == 1);
+
+        /* Update the protection */
+        TempPte.u.Hard.Write = BooleanFlagOn(Protection, IMAGE_SCN_MEM_WRITE);
+#if _MI_HAS_NO_EXECUTE
+        TempPte.u.Hard.NoExecute = !BooleanFlagOn(Protection, IMAGE_SCN_MEM_EXECUTE);
+#endif
+
+        MI_UPDATE_VALID_PTE(PointerPte, TempPte);
+    }
+
+    /* Flush it all */
+    KeFlushEntireTb(TRUE, TRUE);
+
+    return;
+}
+
+VOID
+NTAPI
+MiWriteProtectSystemImage(
+    _In_ PVOID ImageBase)
+{
+    PIMAGE_NT_HEADERS NtHeaders;
+    PIMAGE_SECTION_HEADER SectionHeaders, Section;
+    ULONG i;
+    PVOID SectionBase, SectionEnd;
+    ULONG SectionSize;
+    ULONG Protection;
+    PMMPTE FirstPte, LastPte;
 
     /* Check if the registry setting is on or not */
     if (!MmEnforceWriteProtection)
     {
-            /* Turn on write access too */
+        /* Ignore section protection */
-            SectionProtection |= (IMAGE_SCN_MEM_WRITE | IMAGE_SCN_MEM_EXECUTE);
-        }
-    }
-
-    /* Convert to internal PTE flags */
-    if (SectionProtection & IMAGE_SCN_MEM_EXECUTE) Protection |= MM_EXECUTE;
-    if (SectionProtection & IMAGE_SCN_MEM_READ) Protection |= MM_READONLY;
-
-    /* Check for write access */
-    if (SectionProtection & IMAGE_SCN_MEM_WRITE)
-    {
-        /* Session space is not supported */
-        if (SessionSpace)
-        {
-            DPRINT1("Session drivers not supported\n");
-            ASSERT(SessionSpace == FALSE);
-        }
-        else
-        {
-            /* Convert to internal PTE flag */
-            Protection = (Protection & MM_EXECUTE) ? MM_EXECUTE_READWRITE : MM_READWRITE;
-        }
-    }
-
-    /* If there's no access at all by now, convert to internal no access flag */
-    if (Protection == MM_ZERO_ACCESS) Protection = MM_NOACCESS;
-
-    /* Return the computed PTE protection */
-    return Protection;
-}
-
-VOID
-NTAPI
-MiSetSystemCodeProtection(IN PMMPTE FirstPte,
-                          IN PMMPTE LastPte,
-                          IN ULONG ProtectionMask)
-{
-    /* I'm afraid to introduce regressions at the moment... */
         return;
-}
+    }
 
-VOID
+    /* Large page mapped images are not supported */
-NTAPI
+    NT_ASSERT(!MI_IS_PHYSICAL_ADDRESS(ImageBase));
-MiWriteProtectSystemImage(IN PVOID ImageBase)
-{
-    PIMAGE_NT_HEADERS NtHeaders;
-    PIMAGE_SECTION_HEADER Section;
-    PFN_NUMBER DriverPages;
-    ULONG CurrentProtection, SectionProtection, CombinedProtection = 0, ProtectionMask;
-    ULONG Sections, Size;
-    ULONG_PTR BaseAddress, CurrentAddress;
-    PMMPTE PointerPte, StartPte, LastPte, CurrentPte, ComboPte = NULL;
-    ULONG CurrentMask, CombinedMask = 0;
-    PAGED_CODE();
 
-    /* No need to write protect physical memory-backed drivers (large pages) */
+    /* Session images are not yet supported */
-    if (MI_IS_PHYSICAL_ADDRESS(ImageBase)) return;
+    NT_ASSERT(!MI_IS_SESSION_ADDRESS(ImageBase));
 
-    /* Get the image headers */
+    /* Get the NT headers */
     NtHeaders = RtlImageNtHeader(ImageBase);
-    if (!NtHeaders) return;
+    if (NtHeaders == NULL)
-
-    /* Check if this is a session driver or not */
-    if (!MI_IS_SESSION_ADDRESS(ImageBase))
     {
-        /* Don't touch NT4 drivers */
+        DPRINT1("Failed to get NT headers for image @ %p\n", ImageBase);
-        if (NtHeaders->OptionalHeader.MajorOperatingSystemVersion < 5) return;
-        if (NtHeaders->OptionalHeader.MajorImageVersion < 5) return;
-    }
-    else
-    {
-        /* Not supported */
-        UNIMPLEMENTED_DBGBREAK("Session drivers not supported\n");
-    }
-
-    /* These are the only protection masks we care about */
-    ProtectionMask = IMAGE_SCN_MEM_WRITE | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_EXECUTE;
-
-    /* Calculate the number of pages this driver is occupying */
-    DriverPages = BYTES_TO_PAGES(NtHeaders->OptionalHeader.SizeOfImage);
-
-    /* Get the number of sections and the first section header */
-    Sections = NtHeaders->FileHeader.NumberOfSections;
-    ASSERT(Sections != 0);
-    Section = IMAGE_FIRST_SECTION(NtHeaders);
-
-    /* Loop all the sections */
-    CurrentAddress = (ULONG_PTR)ImageBase;
-    while (Sections)
-    {
-        /* Get the section size */
-        Size = max(Section->SizeOfRawData, Section->Misc.VirtualSize);
-
-        /* Get its virtual address */
-        BaseAddress = (ULONG_PTR)ImageBase + Section->VirtualAddress;
-        if (BaseAddress < CurrentAddress)
-        {
-            /* Windows doesn't like these */
-            DPRINT1("Badly linked image!\n");
         return;
     }
 
-        /* Remember the current address */
+    /* Don't touch NT4 drivers */
-        CurrentAddress = BaseAddress + Size - 1;
+    if ((NtHeaders->OptionalHeader.MajorOperatingSystemVersion < 5) ||
-
+        (NtHeaders->OptionalHeader.MajorSubsystemVersion < 5))
-        /* Next */
+    {
-        Sections--;
+        DPRINT1("Skipping NT 4 driver @ %p\n", ImageBase);
-        Section++;
+        return;
     }
 
-    /* Get the number of sections and the first section header */
+    /* Get the section headers */
-    Sections = NtHeaders->FileHeader.NumberOfSections;
+    SectionHeaders = IMAGE_FIRST_SECTION(NtHeaders);
-    ASSERT(Sections != 0);
-    Section = IMAGE_FIRST_SECTION(NtHeaders);
 
-    /* Set the address at the end to initialize the loop */
+    /* Get the base address of the first section */
-    CurrentAddress = (ULONG_PTR)Section + Sections - 1;
+    SectionBase = Add2Ptr(ImageBase, SectionHeaders[0].VirtualAddress);
-    CurrentProtection = IMAGE_SCN_MEM_WRITE | IMAGE_SCN_MEM_READ;
 
-    /* Set the PTE points for the image, and loop its sections */
+    /* Start protecting the image header as R/O */
-    StartPte = MiAddressToPte(ImageBase);
+    FirstPte = MiAddressToPte(ImageBase);
-    LastPte = StartPte + DriverPages;
+    LastPte = MiAddressToPte(SectionBase) - 1;
-    while (Sections)
+    Protection = IMAGE_SCN_MEM_READ;
+    if (LastPte >= FirstPte)
     {
-        /* Get the section size */
+        MiSetSystemCodeProtection(FirstPte, LastPte, IMAGE_SCN_MEM_READ);
-        Size = max(Section->SizeOfRawData, Section->Misc.VirtualSize);
-
-        /* Get its virtual address and PTE */
-        BaseAddress = (ULONG_PTR)ImageBase + Section->VirtualAddress;
-        PointerPte = MiAddressToPte(BaseAddress);
-
-        /* Check if we were already protecting a run, and found a new run */
-        if ((ComboPte) && (PointerPte > ComboPte))
-        {
-            /* Compute protection */
-            CombinedMask = MiComputeDriverProtection(FALSE, CombinedProtection);
-
-            /* Set it */
-            MiSetSystemCodeProtection(ComboPte, ComboPte, CombinedMask);
-
-            /* Check for overlap */
-            if (ComboPte == StartPte) StartPte++;
-
-            /* One done, reset variables */
-            ComboPte = NULL;
-            CombinedProtection = 0;
     }
 
-        /* Break out when needed */
+    /* Loop the sections */
-        if (PointerPte >= LastPte) break;
+    for (i = 0; i < NtHeaders->FileHeader.NumberOfSections; i++)
-
-        /* Get the requested protection from the image header */
-        SectionProtection = Section->Characteristics & ProtectionMask;
-        if (SectionProtection == CurrentProtection)
     {
-            /* Same protection, so merge the request */
+        /* Get the section base address and size */
-            CurrentAddress = BaseAddress + Size - 1;
+        Section = &SectionHeaders[i];
+        SectionBase = Add2Ptr(ImageBase, Section->VirtualAddress);
+        SectionSize = max(Section->SizeOfRawData, Section->Misc.VirtualSize);
 
-            /* Next */
+        /* Get the first PTE of this section */
-            Sections--;
+        FirstPte = MiAddressToPte(SectionBase);
-            Section++;
+
+        /* Check for overlap with the previous range */
+        if (FirstPte == LastPte)
+        {
+            /* Combine the old and new protection by ORing them */
+            Protection |= (Section->Characteristics & IMAGE_SCN_PROTECTION_MASK);
+
+            /* Update the protection for this PTE */
+            MiSetSystemCodeProtection(FirstPte, FirstPte, Protection);
+
+            /* Skip this PTE */
+            FirstPte++;
+        }
+
+        /* There can not be gaps! */
+        NT_ASSERT(FirstPte == (LastPte + 1));
+
+        /* Get the end of the section and the last PTE */
+        SectionEnd = Add2Ptr(SectionBase, SectionSize - 1);
+        NT_ASSERT(SectionEnd < Add2Ptr(ImageBase, NtHeaders->OptionalHeader.SizeOfImage));
+        LastPte = MiAddressToPte(SectionEnd);
+
+        /* If there are no more pages (after an overlap), skip this section */
+        if (LastPte < FirstPte)
+        {
+            NT_ASSERT(FirstPte == (LastPte + 1));
             continue;
         }
 
-        /* This is now a new section, so close up the old one */
+        /* Get the section protection */
-        CurrentPte = MiAddressToPte(CurrentAddress);
+        Protection = (Section->Characteristics & IMAGE_SCN_PROTECTION_MASK);
 
-        /* Check for overlap */
+        /* Update the protection for this section */
-        if (CurrentPte == PointerPte)
+        MiSetSystemCodeProtection(FirstPte, LastPte, Protection);
-        {
-            /* Skip the last PTE, since it overlaps with us */
-            CurrentPte--;
-
-            /* And set the PTE we will merge with */
-            ASSERT((ComboPte == NULL) || (ComboPte == PointerPte));
-            ComboPte = PointerPte;
-
-            /* Get the most flexible protection by merging both */
-            CombinedMask |= (SectionProtection | CurrentProtection);
     }
 
-        /* Loop any PTEs left */
+    /* Image should end with the last section */
-        if (CurrentPte >= StartPte)
+    NT_ASSERT(ALIGN_UP_POINTER_BY(SectionEnd, PAGE_SIZE) == 
-        {
+              Add2Ptr(ImageBase, NtHeaders->OptionalHeader.SizeOfImage));
-            /* Sanity check */
-            ASSERT(StartPte < LastPte);
-
-            /* Make sure we don't overflow past the last PTE in the driver */
-            if (CurrentPte >= LastPte) CurrentPte = LastPte - 1;
-            ASSERT(CurrentPte >= StartPte);
-
-            /* Compute the protection and set it */
-            CurrentMask = MiComputeDriverProtection(FALSE, CurrentProtection);
-            MiSetSystemCodeProtection(StartPte, CurrentPte, CurrentMask);
-        }
-
-        /* Set new state */
-        StartPte = PointerPte;
-        CurrentAddress = BaseAddress + Size - 1;
-        CurrentProtection = SectionProtection;
-
-        /* Next */
-        Sections--;
-        Section++;
-    }
-
-    /* Is there a leftover section to merge? */
-    if (ComboPte)
-    {
-        /* Compute and set the protection */
-        CombinedMask = MiComputeDriverProtection(FALSE, CombinedProtection);
-        MiSetSystemCodeProtection(ComboPte, ComboPte, CombinedMask);
-
-        /* Handle overlap */
-        if (ComboPte == StartPte) StartPte++;
-    }
-
-    /* Finally, handle the last section */
-    CurrentPte = MiAddressToPte(CurrentAddress);
-    if ((StartPte < LastPte) && (CurrentPte >= StartPte))
-    {
-        /* Handle overlap */
-        if (CurrentPte >= LastPte) CurrentPte = LastPte - 1;
-        ASSERT(CurrentPte >= StartPte);
-
-        /* Compute and set the protection */
-        CurrentMask = MiComputeDriverProtection(FALSE, CurrentProtection);
-        MiSetSystemCodeProtection(StartPte, CurrentPte, CurrentMask);
-    }
 }
 
 VOID

ntoskrnl/mm/mminit.c

@@ -203,6 +203,8 @@ MmInitSystem(IN ULONG Phase,
     PMMPTE PointerPte;
     MMPTE TempPte = ValidKernelPte;
     PFN_NUMBER PageFrameNumber;
+    PLIST_ENTRY ListEntry;
+    PLDR_DATA_TABLE_ENTRY DataTableEntry;
 
     /* Initialize the kernel address space */
     ASSERT(Phase == 1);
@@ -271,6 +273,18 @@ MmInitSystem(IN ULONG Phase,
     /* Initialize the balance set manager */
     MmInitBsmThread();
 
+    /* Loop the boot loaded images */
+    for (ListEntry = PsLoadedModuleList.Flink;
+         ListEntry != &PsLoadedModuleList;
+         ListEntry = ListEntry->Flink)
+    {
+        /* Get the data table entry */
+        DataTableEntry = CONTAINING_RECORD(ListEntry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
+
+        /* Set up the image protection */
+        MiWriteProtectSystemImage(DataTableEntry->DllBase);
+    }
+
     return TRUE;
 }