[NTOSKRNL]
[reactos.git] / reactos / ntoskrnl / mm / ARM3 / sysldr.c
index d8f7157..3f42163 100644 (file)
@@ -1,11 +1,11 @@
 /*
-* PROJECT:         ReactOS Kernel
-* LICENSE:         BSD - See COPYING.ARM in the top level directory
-* FILE:            ntoskrnl/mm/ARM3/sysldr.c
-* PURPOSE:         Contains the Kernel Loader (SYSLDR) for loading PE files.
-* PROGRAMMERS:     Alex Ionescu (alex.ionescu@reactos.org)
-*                  ReactOS Portable Systems Group
-*/
+ * PROJECT:         ReactOS Kernel
+ * LICENSE:         BSD - See COPYING.ARM in the top level directory
+ * FILE:            ntoskrnl/mm/ARM3/sysldr.c
+ * PURPOSE:         Contains the Kernel Loader (SYSLDR) for loading PE files.
+ * PROGRAMMERS:     Alex Ionescu (alex.ionescu@reactos.org)
+ *                  ReactOS Portable Systems Group
+ */
 
 /* INCLUDES *******************************************************************/
 
@@ -13,7 +13,6 @@
 #define NDEBUG
 #include <debug.h>
 
-#line 16 "ARMĀ³::LOADER"
 #define MODULE_INVOLVED_IN_ARM3
 #include "../ARM3/miarm.h"
 
@@ -96,7 +95,7 @@ MiLoadImageSection(IN OUT PVOID *SectionPtr,
     KAPC_STATE ApcState;
     LARGE_INTEGER SectionOffset = {{0, 0}};
     BOOLEAN LoadSymbols = FALSE;
-    PFN_NUMBER PteCount;
+    PFN_COUNT PteCount;
     PMMPTE PointerPte, LastPte;
     PVOID DriverBase;
     MMPTE TempPte;
@@ -112,7 +111,7 @@ MiLoadImageSection(IN OUT PVOID *SectionPtr,
 
     /* Not session load, shouldn't have an entry */
     ASSERT(LdrEntry == NULL);
-    
+
     /* Attach to the system process */
     KeStackAttachProcess(&PsInitialSystemProcess->Pcb, &ApcState);
 
@@ -151,27 +150,45 @@ MiLoadImageSection(IN OUT PVOID *SectionPtr,
     if (!NT_SUCCESS(Status))
     {
         /* Detach and return */
+        DPRINT1("MmMapViewOfSection failed with status 0x%x\n", Status);
         KeUnstackDetachProcess(&ApcState);
         return Status;
     }
-    
+
     /* Reserve system PTEs needed */
     PteCount = ROUND_TO_PAGES(Section->ImageSection->ImageSize) >> PAGE_SHIFT;
     PointerPte = MiReserveSystemPtes(PteCount, SystemPteSpace);
-    if (!PointerPte) return STATUS_INSUFFICIENT_RESOURCES;
-    
+    if (!PointerPte)
+    {
+        DPRINT1("MiReserveSystemPtes failed\n");
+        KeUnstackDetachProcess(&ApcState);
+        return STATUS_INSUFFICIENT_RESOURCES;
+    }
+
     /* New driver base */
     LastPte = PointerPte + PteCount;
     DriverBase = MiPteToAddress(PointerPte);
 
     /* The driver is here */
     *ImageBase = DriverBase;
+    DPRINT1("Loading: %wZ at %p with %lx pages\n", FileName, DriverBase, PteCount);
 
     /* Loop the new driver PTEs */
     TempPte = ValidKernelPte;
     while (PointerPte < LastPte)
     {
         /* Allocate a page */
+        MI_SET_USAGE(MI_USAGE_DRIVER_PAGE);
+#if MI_TRACE_PFNS
+        PWCHAR pos = NULL;
+        ULONG len = 0;
+        if (FileName->Buffer)
+        {
+            pos = wcsrchr(FileName->Buffer, '\\');
+            len = wcslen(pos) * sizeof(WCHAR);
+            if (pos) snprintf(MI_PFN_CURRENT_PROCESS_NAME, min(16, len), "%S", pos);
+        }
+#endif
         TempPte.u.Hard.PageFrameNumber = MiAllocatePfn(PointerPte, MM_EXECUTE);
 
         /* Write it */
@@ -180,7 +197,7 @@ MiLoadImageSection(IN OUT PVOID *SectionPtr,
         /* Move on */
         PointerPte++;
     }
-        
+
     /* Copy the image */
     RtlCopyMemory(DriverBase, Base, PteCount << PAGE_SHIFT);
 
@@ -313,7 +330,7 @@ MmCallDllInitialize(IN PLDR_DATA_TABLE_ENTRY LdrEntry,
     if (wcschr(ImportName.Buffer, L'.'))
     {
         /* Remove the extension */
-        ImportName.Length = (wcschr(ImportName.Buffer, L'.') -
+        ImportName.Length = (USHORT)(wcschr(ImportName.Buffer, L'.') -
             ImportName.Buffer) * sizeof(WCHAR);
     }
 
@@ -325,7 +342,7 @@ MmCallDllInitialize(IN PLDR_DATA_TABLE_ENTRY LdrEntry,
     Status = DllInit(&RegPath);
 
     /* Clean up */
-    ExFreePool(RegPath.Buffer);
+    ExFreePoolWithTag(RegPath.Buffer, TAG_LDR_WSTR);
 
     /* Return status value which DllInitialize returned */
     return Status;
@@ -374,14 +391,14 @@ MiDereferenceImports(IN PLOAD_IMPORTS ImportList)
         /* Then there's nothing to do */
         return STATUS_SUCCESS;
     }
-    
+
     /* Check for single-entry */
     if ((ULONG_PTR)ImportList & MM_SYSLDR_SINGLE_ENTRY)
     {
         /* Set it up */
         SingleEntry.Count = 1;
         SingleEntry.Entry[0] = (PVOID)((ULONG_PTR)ImportList &~ MM_SYSLDR_SINGLE_ENTRY);
-        
+
         /* Use this as the import list */
         ImportList = &SingleEntry;
     }
@@ -392,31 +409,31 @@ MiDereferenceImports(IN PLOAD_IMPORTS ImportList)
         /* Get the entry */
         LdrEntry = ImportList->Entry[i];
         DPRINT1("%wZ <%wZ>\n", &LdrEntry->FullDllName, &LdrEntry->BaseDllName);
-        
+
         /* Skip boot loaded images */
         if (LdrEntry->LoadedImports == MM_SYSLDR_BOOT_LOADED) continue;
-        
+
         /* Dereference the entry */
         ASSERT(LdrEntry->LoadCount >= 1);
         if (!--LdrEntry->LoadCount)
         {
             /* Save the import data in case unload fails */
             CurrentImports = LdrEntry->LoadedImports;
-            
+
             /* This is the last entry */
             LdrEntry->LoadedImports = MM_SYSLDR_NO_IMPORTS;
             if (MiCallDllUnloadAndUnloadDll(LdrEntry))
             {
                 /* Unloading worked, parse this DLL's imports too */
                 MiDereferenceImports(CurrentImports);
-                
+
                 /* Check if we had valid imports */
                 if ((CurrentImports != MM_SYSLDR_BOOT_LOADED) ||
                     (CurrentImports != MM_SYSLDR_NO_IMPORTS) ||
                     !((ULONG_PTR)LdrEntry->LoadedImports & MM_SYSLDR_SINGLE_ENTRY))
                 {
                     /* Free them */
-                    ExFreePool(CurrentImports);
+                    ExFreePoolWithTag(CurrentImports, TAG_LDR_IMPORTS);
                 }
             }
             else
@@ -426,7 +443,7 @@ MiDereferenceImports(IN PLOAD_IMPORTS ImportList)
             }
         }
     }
-    
+
     /* Done */
     return STATUS_SUCCESS;
 }
@@ -447,7 +464,7 @@ MiClearImports(IN PLDR_DATA_TABLE_ENTRY LdrEntry)
     }
 
     /* Otherwise, free the import list */
-    ExFreePool(LdrEntry->LoadedImports);
+    ExFreePoolWithTag(LdrEntry->LoadedImports, TAG_LDR_IMPORTS);
     LdrEntry->LoadedImports = MM_SYSLDR_BOOT_LOADED;
 }
 
@@ -551,6 +568,7 @@ MiProcessLoaderEntry(IN PLDR_DATA_TABLE_ENTRY LdrEntry,
 
 VOID
 NTAPI
+INIT_FUNCTION
 MiUpdateThunks(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
                IN PVOID OldBase,
                IN PVOID NewBase,
@@ -559,7 +577,20 @@ MiUpdateThunks(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
     ULONG_PTR OldBaseTop, Delta;
     PLDR_DATA_TABLE_ENTRY LdrEntry;
     PLIST_ENTRY NextEntry;
-    ULONG ImportSize, i;
+    ULONG ImportSize;
+    //
+    // FIXME: MINGW-W64 must fix LD to generate drivers that Windows can load,
+    // since a real version of Windows would fail at this point, but they seem
+    // busy implementing features such as "HotPatch" support in GCC 4.6 instead,
+    // a feature which isn't even used by Windows. Priorities, priorities...
+    // Please note that Microsoft WDK EULA and license prohibits using
+    // the information contained within it for the generation of "non-Windows"
+    // drivers, which is precisely what LD will generate, since an LD driver
+    // will not load on Windows.
+    //
+#ifdef _WORKING_LINKER_
+    ULONG i;
+#endif
     PULONG_PTR ImageThunk;
     PIMAGE_IMPORT_DESCRIPTOR ImportDescriptor;
 
@@ -599,7 +630,6 @@ MiUpdateThunks(IN PLOADER_PARAMETER_BLOCK LoaderBlock,
         }
 #else
         /* Get the import table */
-        i = ImportSize;
         ImportDescriptor = RtlImageDirectoryEntryToData(LdrEntry->DllBase,
                                                         TRUE,
                                                         IMAGE_DIRECTORY_ENTRY_IMPORT,
@@ -666,7 +696,7 @@ MiSnapThunk(IN PVOID DllBase,
     ULONG ForwardExportSize;
     PIMAGE_EXPORT_DIRECTORY ForwardExportDirectory;
     PIMAGE_IMPORT_BY_NAME ForwardName;
-    ULONG ForwardLength;
+    SIZE_T ForwardLength;
     IMAGE_THUNK_DATA ForwardThunk;
     PAGED_CODE();
 
@@ -686,9 +716,9 @@ MiSnapThunk(IN PVOID DllBase,
         NameImport = (PIMAGE_IMPORT_BY_NAME)Name->u1.AddressOfData;
 
         /* Copy the procedure name */
-        strncpy(*MissingApi,
-                (PCHAR)&NameImport->Name[0],
-                MAXIMUM_FILENAME_LENGTH - 1);
+        RtlStringCbCopyA(*MissingApi,
+                         MAXIMUM_FILENAME_LENGTH,
+                         (PCHAR)&NameImport->Name[0]);
 
         /* Setup name tables */
         DPRINT("Import name: %s\n", NameImport->Name);
@@ -734,7 +764,11 @@ MiSnapThunk(IN PVOID DllBase,
             }
 
             /* Check if we couldn't find it */
-            if (High < Low) return STATUS_DRIVER_ENTRYPOINT_NOT_FOUND;
+            if (High < Low)
+            {
+                DPRINT1("Warning: Driver failed to load, %s not found\n", NameImport->Name);
+                return STATUS_DRIVER_ENTRYPOINT_NOT_FOUND;
+            }
 
             /* Otherwise, this is the ordinal */
             Ordinal = OrdinalTable[Mid];
@@ -769,9 +803,9 @@ MiSnapThunk(IN PVOID DllBase,
 
             /* Build the forwarder name */
             DllName.Buffer = (PCHAR)Address->u1.Function;
-            DllName.Length = strchr(DllName.Buffer, '.') -
-                             DllName.Buffer +
-                             sizeof(ANSI_NULL);
+            DllName.Length = (USHORT)(strchr(DllName.Buffer, '.') -
+                                      DllName.Buffer) +
+                                      sizeof(ANSI_NULL);
             DllName.MaximumLength = DllName.Length;
 
             /* Convert it */
@@ -920,7 +954,7 @@ MmUnloadSystemImage(IN PVOID ImageHandle)
         if (LdrEntry->FullDllName.Buffer)
         {
             /* Free it */
-            ExFreePool(LdrEntry->FullDllName.Buffer);
+            ExFreePoolWithTag(LdrEntry->FullDllName.Buffer, TAG_LDR_WSTR);
         }
 
         /* Check if we had a section */
@@ -931,7 +965,7 @@ MmUnloadSystemImage(IN PVOID ImageHandle)
         }
 
         /* Free the entry */
-        ExFreePool(LdrEntry);
+        ExFreePoolWithTag(LdrEntry, TAG_MODULE_OBJECT);
     }
 
     /* Release the system lock and return */
@@ -994,7 +1028,7 @@ MiResolveImageReferences(IN PVOID ImageBase,
         LoadedImportsSize = ImportCount * sizeof(PVOID) + sizeof(SIZE_T);
         LoadedImports = ExAllocatePoolWithTag(PagedPool,
                                               LoadedImportsSize,
-                                              'TDmM');
+                                              TAG_LDR_IMPORTS);
         if (LoadedImports)
         {
             /* Zero it */
@@ -1031,7 +1065,7 @@ MiResolveImageReferences(IN PVOID ImageBase,
         {
             /* It's not, it's importing stuff it shouldn't be! */
             MiDereferenceImports(LoadedImports);
-            if (LoadedImports) ExFreePool(LoadedImports);
+            if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
             return STATUS_PROCEDURE_NOT_FOUND;
         }
 
@@ -1045,7 +1079,7 @@ MiResolveImageReferences(IN PVOID ImageBase,
         {
             /* This is not kernel code */
             MiDereferenceImports(LoadedImports);
-            if (LoadedImports) ExFreePool(LoadedImports);
+            if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
             return STATUS_PROCEDURE_NOT_FOUND;
         }
 
@@ -1070,7 +1104,7 @@ MiResolveImageReferences(IN PVOID ImageBase,
         {
             /* Failed */
             MiDereferenceImports(LoadedImports);
-            if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_WSTR);
+            if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
             return Status;
         }
 
@@ -1125,7 +1159,7 @@ CheckDllState:
                                     sizeof(UNICODE_NULL);
             DllName.Buffer = ExAllocatePoolWithTag(NonPagedPool,
                                                    DllName.MaximumLength,
-                                                   'TDmM');
+                                                   TAG_LDR_WSTR);
             if (DllName.Buffer)
             {
                 /* Setup the base length and copy it */
@@ -1135,9 +1169,9 @@ CheckDllState:
                               ImageFileDirectory->Length);
 
                 /* Now add the import name and null-terminate it */
-                RtlAppendStringToString((PSTRING)&DllName,
-                                        (PSTRING)&NameString);
-                DllName.Buffer[(DllName.MaximumLength - 1) / sizeof(WCHAR)] = UNICODE_NULL;
+                RtlAppendUnicodeStringToString(&DllName,
+                                               &NameString);
+                DllName.Buffer[DllName.Length / sizeof(WCHAR)] = UNICODE_NULL;
 
                 /* Load the image */
                 Status = MmLoadSystemImage(&DllName,
@@ -1149,7 +1183,7 @@ CheckDllState:
                 if (NT_SUCCESS(Status))
                 {
                     /* We can free the DLL Name */
-                    ExFreePool(DllName.Buffer);
+                    ExFreePoolWithTag(DllName.Buffer, TAG_LDR_WSTR);
                 }
                 else
                 {
@@ -1157,6 +1191,8 @@ CheckDllState:
                     *MissingDriver = DllName.Buffer;
                     *(PULONG)MissingDriver |= 1;
                     *MissingApi = NULL;
+
+                    DPRINT1("Failed to load dependency: %wZ\n", &DllName);
                 }
             }
             else
@@ -1172,7 +1208,7 @@ CheckDllState:
                 Loaded = TRUE;
 
                 /* Sanity check */
-                ASSERT(DllBase = DllEntry->DllBase);
+                ASSERT(DllBase == DllEntry->DllBase);
 
                 /* Call the initialization routines */
                 Status = MmCallDllInitialize(DllEntry, &PsLoadedModuleList);
@@ -1180,6 +1216,7 @@ CheckDllState:
                 {
                     /* We failed, unload the image */
                     MmUnloadSystemImage(DllEntry);
+                    DPRINT1("MmCallDllInitialize failed with status 0x%x\n", Status);
                     while (TRUE);
                     Loaded = FALSE;
                 }
@@ -1191,7 +1228,7 @@ CheckDllState:
                 /* Cleanup and return */
                 RtlFreeUnicodeString(&NameString);
                 MiDereferenceImports(LoadedImports);
-                if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_WSTR);
+                if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
                 return Status;
             }
 
@@ -1224,7 +1261,8 @@ CheckDllState:
         {
             /* Cleanup and return */
             MiDereferenceImports(LoadedImports);
-            if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_WSTR);
+            if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
+            DPRINT1("Warning: Driver failed to load, %S not found\n", *MissingDriver);
             return STATUS_DRIVER_ENTRYPOINT_NOT_FOUND;
         }
 
@@ -1253,7 +1291,7 @@ CheckDllState:
                 {
                     /* Cleanup and return */
                     MiDereferenceImports(LoadedImports);
-                    if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_WSTR);
+                    if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
                     return Status;
                 }
 
@@ -1269,7 +1307,7 @@ CheckDllState:
     /* Check if we have an import list */
     if (LoadedImports)
     {
-        /* Reset the count again, and loop entries*/
+        /* Reset the count again, and loop entries */
         ImportCount = 0;
         for (i = 0; i < LoadedImports->Count; i++)
         {
@@ -1286,13 +1324,13 @@ CheckDllState:
         if (!ImportCount)
         {
             /* Free the list and set it to no imports */
-            ExFreePoolWithTag(LoadedImports, TAG_LDR_WSTR);
+            ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
             LoadedImports = MM_SYSLDR_NO_IMPORTS;
         }
         else if (ImportCount == 1)
         {
             /* Just one entry, we can free the table and only use our entry */
-            ExFreePoolWithTag(LoadedImports, TAG_LDR_WSTR);
+            ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
             LoadedImports = (PLOAD_IMPORTS)ImportEntry;
         }
         else if (ImportCount != LoadedImports->Count)
@@ -1301,7 +1339,7 @@ CheckDllState:
             LoadedImportsSize = ImportCount * sizeof(PVOID) + sizeof(SIZE_T);
             NewImports = ExAllocatePoolWithTag(PagedPool,
                                                LoadedImportsSize,
-                                               'TDmM');
+                                               TAG_LDR_IMPORTS);
             if (NewImports)
             {
                 /* Set count */
@@ -1320,7 +1358,7 @@ CheckDllState:
                 }
 
                 /* Free the old copy */
-                ExFreePoolWithTag(LoadedImports, TAG_LDR_WSTR);
+                ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
                 LoadedImports = NewImports;
             }
         }
@@ -1335,6 +1373,244 @@ CheckDllState:
 
 VOID
 NTAPI
+MiFreeInitializationCode(IN PVOID InitStart,
+                         IN PVOID InitEnd)
+{
+    PMMPTE PointerPte;
+    PFN_NUMBER PagesFreed;
+
+    /* Get the start PTE */
+    PointerPte = MiAddressToPte(InitStart);
+    ASSERT(MI_IS_PHYSICAL_ADDRESS(InitStart) == FALSE);
+
+    /*  Compute the number of pages we expect to free */
+    PagesFreed = (PFN_NUMBER)(MiAddressToPte(InitEnd) - PointerPte + 1);
+    
+    /* Try to actually free them */
+    PagesFreed = MiDeleteSystemPageableVm(PointerPte,
+                                          PagesFreed,
+                                          0,
+                                          NULL);
+}
+
+VOID
+NTAPI
+INIT_FUNCTION
+MiFindInitializationCode(OUT PVOID *StartVa,
+                         OUT PVOID *EndVa)
+{
+    ULONG Size, SectionCount, Alignment;
+    PLDR_DATA_TABLE_ENTRY LdrEntry;
+    ULONG_PTR DllBase, InitStart, InitEnd, ImageEnd, InitCode;
+    PLIST_ENTRY NextEntry;
+    PIMAGE_NT_HEADERS NtHeader;
+    PIMAGE_SECTION_HEADER Section, LastSection;
+    BOOLEAN InitFound;
+    
+    /* So we don't free our own code yet */
+    InitCode = (ULONG_PTR)&MiFindInitializationCode;
+
+    /* Assume failure */
+    *StartVa = NULL;
+
+    /* Enter a critical region while we loop the list */
+    KeEnterCriticalRegion();
+
+    /* Loop all loaded modules */
+    NextEntry = PsLoadedModuleList.Flink;
+    while (NextEntry != &PsLoadedModuleList)
+    {
+        /* Get the loader entry and its DLL base */
+        LdrEntry = CONTAINING_RECORD(NextEntry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
+        DllBase = (ULONG_PTR)LdrEntry->DllBase;
+        
+        /* Get the NT header */
+        NtHeader = RtlImageNtHeader((PVOID)DllBase);
+        if (!NtHeader)
+        {
+            /* Keep going */
+            NextEntry = NextEntry->Flink;
+            continue;
+        }
+
+        /* Get the first section, the section count, and scan them all */
+        Section = IMAGE_FIRST_SECTION(NtHeader);
+        SectionCount = NtHeader->FileHeader.NumberOfSections;
+        InitStart = 0;
+        while (SectionCount > 0)
+        {
+            /* Assume failure */
+            InitFound = FALSE;
+
+            /* Is this the INIT section or a discardable section? */
+            if ((*(PULONG)Section->Name == 'TINI') ||
+                ((Section->Characteristics & IMAGE_SCN_MEM_DISCARDABLE)))
+            {
+                /* Remember this */
+                InitFound = TRUE;
+            }
+
+            if (InitFound)
+            {
+                /* Pick the biggest size -- either raw or virtual */
+                Size = max(Section->SizeOfRawData, Section->Misc.VirtualSize);
+                
+                /* Read the section alignment */
+                Alignment = NtHeader->OptionalHeader.SectionAlignment;
+
+                /* Align the start and end addresses appropriately */
+                InitStart = DllBase + Section->VirtualAddress;
+                InitEnd = ((Alignment + InitStart + Size - 2) & 0xFFFFF000) - 1;                        
+                InitStart = (InitStart + (PAGE_SIZE - 1)) & 0xFFFFF000;
+
+                /* Have we reached the last section? */
+                if (SectionCount == 1)
+                {
+                    /* Remember this */
+                    LastSection = Section;
+                }
+                else
+                {
+                    /* We have not, loop all the sections */
+                    LastSection = NULL;
+                    do
+                    {
+                        /* Keep going until we find a non-discardable section range */
+                        SectionCount--;
+                        Section++;
+                        if (Section->Characteristics & IMAGE_SCN_MEM_DISCARDABLE)
+                        {
+                            /* Discardable, so record it, then keep going */
+                            LastSection = Section;
+                        }
+                        else
+                        {
+                            /* Non-contigous discard flag, or no flag, break out */
+                            break;
+                        }
+                    }
+                    while (SectionCount > 1);
+                }
+
+                /* Have we found a discardable or init section? */
+                if (LastSection)
+                {
+                    /* Pick the biggest size -- either raw or virtual */
+                    Size = max(LastSection->SizeOfRawData, LastSection->Misc.VirtualSize);
+
+                    /* Use this as the end of the section address */
+                    InitEnd = DllBase + LastSection->VirtualAddress + Size - 1;
+
+                    /* Have we reached the last section yet? */
+                    if (SectionCount != 1)
+                    {
+                        /* Then align this accross the session boundary */
+                        InitEnd = ((Alignment + InitEnd - 1) & 0XFFFFF000) - 1;
+                    }
+                }
+
+                /* Make sure we don't let the init section go past the image */
+                ImageEnd = DllBase + LdrEntry->SizeOfImage;
+                if (InitEnd > ImageEnd) InitEnd = (ImageEnd - 1) | (PAGE_SIZE - 1);
+
+                /* Make sure we have a valid, non-zero init section */
+                if (InitStart <= InitEnd)
+                {
+                    /* Make sure we are not within this code itself */
+                    if ((InitCode >= InitStart) && (InitCode <= InitEnd))
+                    {
+                        /* Return it, we can't free ourselves now */
+                        ASSERT(*StartVa == 0);
+                        *StartVa = (PVOID)InitStart;
+                        *EndVa = (PVOID)InitEnd;
+                    }
+                    else
+                    {
+                        /* This isn't us -- go ahead and free it */
+                        ASSERT(MI_IS_PHYSICAL_ADDRESS((PVOID)InitStart) == FALSE);
+                        MiFreeInitializationCode((PVOID)InitStart, (PVOID)InitEnd);
+                    }
+                }
+            }
+            
+            /* Move to the next section */
+            SectionCount--;
+            Section++;
+        }
+        
+        /* Move to the next module */
+        NextEntry = NextEntry->Flink;
+    }
+
+    /* Leave the critical region and return */
+    KeLeaveCriticalRegion();
+}
+
+/* 
+ * Note: This function assumes that all discardable sections are at the end of
+ * the PE file. It searches backwards until it finds the non-discardable section
+ */
+VOID
+NTAPI
+MmFreeDriverInitialization(IN PLDR_DATA_TABLE_ENTRY LdrEntry)
+{
+    PMMPTE StartPte, EndPte;
+    PFN_NUMBER PageCount;
+    PVOID DllBase;
+    ULONG i;
+    PIMAGE_NT_HEADERS NtHeader;
+    PIMAGE_SECTION_HEADER Section, DiscardSection;
+    ULONG PagesDeleted;
+
+    /* Get the base address and the page count */
+    DllBase = LdrEntry->DllBase;
+    PageCount = LdrEntry->SizeOfImage >> PAGE_SHIFT;
+
+    /* Get the last PTE in this image */
+    EndPte = MiAddressToPte(DllBase) + PageCount;
+
+    /* Get the NT header */
+    NtHeader = RtlImageNtHeader(DllBase);
+    if (!NtHeader) return;
+
+    /* Get the last section and loop each section backwards */
+    Section = IMAGE_FIRST_SECTION(NtHeader) + NtHeader->FileHeader.NumberOfSections;
+    DiscardSection = NULL;
+    for (i = 0; i < NtHeader->FileHeader.NumberOfSections; i++)
+    {
+        /* Go back a section and check if it's discardable */
+        Section--;
+        if (Section->Characteristics & IMAGE_SCN_MEM_DISCARDABLE)
+        {
+            /* It is, select it for freeing */
+            DiscardSection = Section;
+        }
+        else
+        {
+            /* No more discardable sections exist, bail out */
+            break;
+        }
+    }
+
+    /* Bail out if there's nothing to free */
+    if (!DiscardSection) return;
+
+    /* Push the DLL base to the first disacrable section, and get its PTE */
+    DllBase = (PVOID)ROUND_TO_PAGES((ULONG_PTR)DllBase + DiscardSection->VirtualAddress);
+    ASSERT(MI_IS_PHYSICAL_ADDRESS(DllBase) == FALSE);
+    StartPte = MiAddressToPte(DllBase);
+
+    /* Check how many pages to free total */
+    PageCount = (PFN_NUMBER)(EndPte - StartPte);
+    if (!PageCount) return;
+
+    /* Delete this many PTEs */
+    PagesDeleted = MiDeleteSystemPageableVm(StartPte, PageCount, 0, NULL);
+}
+
+VOID
+NTAPI
+INIT_FUNCTION
 MiReloadBootLoadedDrivers(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
 {
     PLIST_ENTRY NextEntry;
@@ -1347,7 +1623,7 @@ MiReloadBootLoadedDrivers(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
     PVOID DllBase, NewImageAddress;
     NTSTATUS Status;
     PMMPTE PointerPte, StartPte, LastPte;
-    PFN_NUMBER PteCount;
+    PFN_COUNT PteCount;
     PMMPFN Pfn1;
     MMPTE TempPte, OldPte;
 
@@ -1368,6 +1644,23 @@ MiReloadBootLoadedDrivers(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
                 (ULONG_PTR)LdrEntry->DllBase + LdrEntry->SizeOfImage,
                 &LdrEntry->FullDllName);
 
+        /* Get the first PTE and the number of PTEs we'll need */
+        PointerPte = StartPte = MiAddressToPte(LdrEntry->DllBase);
+        PteCount = ROUND_TO_PAGES(LdrEntry->SizeOfImage) >> PAGE_SHIFT;
+        LastPte = StartPte + PteCount;
+
+#if MI_TRACE_PFNS
+        /* Loop the PTEs */
+        while (PointerPte < LastPte)
+        {
+            ULONG len;
+            ASSERT(PointerPte->u.Hard.Valid == 1);
+            Pfn1 = MiGetPfnEntry(PFN_FROM_PTE(PointerPte));
+            len = wcslen(LdrEntry->BaseDllName.Buffer) * sizeof(WCHAR);
+            snprintf(Pfn1->ProcessName, min(16, len), "%S", LdrEntry->BaseDllName.Buffer);
+            PointerPte++;
+        }
+#endif
         /* Skip kernel and HAL */
         /* ROS HACK: Skip BOOTVID/KDCOM too */
         i++;
@@ -1406,13 +1699,9 @@ MiReloadBootLoadedDrivers(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
 
         /* Remember the original address */
         DllBase = LdrEntry->DllBase;
-        
-        /* Get the first PTE and the number of PTEs we'll need */
-        PointerPte = StartPte = MiAddressToPte(LdrEntry->DllBase);
-        PteCount = ROUND_TO_PAGES(LdrEntry->SizeOfImage) >> PAGE_SHIFT;
-        LastPte = StartPte + PteCount;
-        
+
         /* Loop the PTEs */
+        PointerPte = StartPte;
         while (PointerPte < LastPte)
         {
             /* Mark the page modified in the PFN database */
@@ -1420,11 +1709,11 @@ MiReloadBootLoadedDrivers(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
             Pfn1 = MiGetPfnEntry(PFN_FROM_PTE(PointerPte));
             ASSERT(Pfn1->u3.e1.Rom == 0);
             Pfn1->u3.e1.Modified = TRUE;
-            
+
             /* Next */
             PointerPte++;
         }
-        
+
         /* Now reserve system PTEs for the image */
         PointerPte = MiReserveSystemPtes(PteCount, SystemPteSpace);
         if (!PointerPte)
@@ -1433,7 +1722,7 @@ MiReloadBootLoadedDrivers(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
             DPRINT1("[Mm0]: Couldn't allocate driver section!\n");
             while (TRUE);
         }
-        
+
         /* This is the new virtual address for the module */
         LastPte = PointerPte + PteCount;
         NewImageAddress = MiPteToAddress(PointerPte);
@@ -1441,7 +1730,7 @@ MiReloadBootLoadedDrivers(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
         /* Sanity check */
         DPRINT("[Mm0]: Copying from: %p to: %p\n", DllBase, NewImageAddress);
         ASSERT(ExpInitializationPhase == 0);
-        
+
         /* Loop the new driver PTEs */
         TempPte = ValidKernelPte;
         while (PointerPte < LastPte)
@@ -1460,7 +1749,7 @@ MiReloadBootLoadedDrivers(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
             PointerPte++;
             StartPte++;
         }
-        
+
         /* Update position */
         PointerPte -= PteCount;
 
@@ -1503,13 +1792,14 @@ MiReloadBootLoadedDrivers(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
         LdrEntry->EntryPoint = (PVOID)((ULONG_PTR)NewImageAddress +
                                 NtHeader->OptionalHeader.AddressOfEntryPoint);
         LdrEntry->SizeOfImage = PteCount << PAGE_SHIFT;
-        
+
         /* FIXME: We'll need to fixup the PFN linkage when switching to ARM3 */
     }
 }
 
 NTSTATUS
 NTAPI
+INIT_FUNCTION
 MiBuildImportsForBootDrivers(VOID)
 {
     PLIST_ENTRY NextEntry, NextEntry2;
@@ -1523,7 +1813,7 @@ MiBuildImportsForBootDrivers(VOID)
     ULONG_PTR DllBase, DllEnd;
     ULONG Modules = 0, i, j = 0;
     PIMAGE_IMPORT_DESCRIPTOR ImportDescriptor;
-    
+
     /* Initialize variables */
     KernelEntry = HalEntry = LastEntry = NULL;
 
@@ -1547,7 +1837,7 @@ MiBuildImportsForBootDrivers(VOID)
             /* Found it */
             HalEntry = LdrEntry;
         }
-        
+
         /* Check if this is a driver DLL */
         if (LdrEntry->Flags & LDRP_DRIVER_DEPENDENT_DLL)
         {
@@ -1566,9 +1856,9 @@ MiBuildImportsForBootDrivers(VOID)
         else
         {
             /* No referencing needed */
-            LdrEntry->LoadCount = 0;    
+            LdrEntry->LoadCount = 0;
         }
-        
+
         /* Remember this came from the loader */
         LdrEntry->LoadedImports = MM_SYSLDR_BOOT_LOADED;
 
@@ -1576,12 +1866,12 @@ MiBuildImportsForBootDrivers(VOID)
         NextEntry = NextEntry->Flink;
         Modules++;
     }
-    
+
     /* We must have at least found the kernel and HAL */
     if (!(HalEntry) || (!KernelEntry)) return STATUS_NOT_FOUND;
-    
+
     /* Allocate the list */
-    EntryArray = ExAllocatePoolWithTag(PagedPool, Modules * sizeof(PVOID), 'TDmM');
+    EntryArray = ExAllocatePoolWithTag(PagedPool, Modules * sizeof(PVOID), TAG_LDR_IMPORTS);
     if (!EntryArray) return STATUS_INSUFFICIENT_RESOURCES;
 
     /* Loop the loaded module list again */
@@ -1598,7 +1888,7 @@ MiBuildImportsForBootDrivers(VOID)
                                                   TRUE,
                                                   IMAGE_DIRECTORY_ENTRY_IAT,
                                                   &ImportSize);
-        if (!ImageThunk)                                                
+        if (!ImageThunk)
 #else
         /* Get its imports */
         ImportDescriptor = RtlImageDirectoryEntryToData(LdrEntry->DllBase,
@@ -1613,16 +1903,16 @@ MiBuildImportsForBootDrivers(VOID)
             NextEntry = NextEntry->Flink;
             continue;
         }
-        
+
         /* Clear the list and count the number of IAT thunks */
         RtlZeroMemory(EntryArray, Modules * sizeof(PVOID));
 #ifdef _WORKING_LOADER_
         ImportSize /= sizeof(ULONG_PTR);
-        
+
         /* Scan the thunks */
         for (i = 0, DllBase = 0, DllEnd = 0; i < ImportSize; i++, ImageThunk++)
 #else
-        i = DllBase = DllEnd = 0;
+        DllBase = DllEnd = i = 0;
         while ((ImportDescriptor->Name) &&
                (ImportDescriptor->OriginalFirstThunk))
         {
@@ -1644,7 +1934,7 @@ MiBuildImportsForBootDrivers(VOID)
                     continue;
                 }
             }
-            
+
             /* Loop the loaded module list to locate this address owner */
             j = 0;
             NextEntry2 = PsLoadedModuleList.Flink;
@@ -1654,11 +1944,11 @@ MiBuildImportsForBootDrivers(VOID)
                 LdrEntry2 = CONTAINING_RECORD(NextEntry2,
                                               LDR_DATA_TABLE_ENTRY,
                                               InLoadOrderLinks);
-                                              
+
                 /* Get the address range for this module */
                 DllBase = (ULONG_PTR)LdrEntry2->DllBase;
                 DllEnd = DllBase + LdrEntry2->SizeOfImage;
-                
+
                 /* Check if this IAT entry matches it */
                 if ((*ImageThunk >= DllBase) && (*ImageThunk < DllEnd))
                 {
@@ -1667,12 +1957,12 @@ MiBuildImportsForBootDrivers(VOID)
                     EntryArray[j] = LdrEntry2;
                     break;
                 }
-                
+
                 /* Keep searching */
                 NextEntry2 = NextEntry2->Flink;
                 j++;
             }
-            
+
             /* Do we have a thunk outside the range? */
             if ((*ImageThunk < DllBase) || (*ImageThunk >= DllEnd))
             {
@@ -1684,19 +1974,19 @@ MiBuildImportsForBootDrivers(VOID)
                             LdrEntry, ImageThunk, *ImageThunk);
                     ASSERT(FALSE);
                 }
-                
+
                 /* Reset if we hit this */
                 DllBase = 0;
             }
 #ifndef _WORKING_LOADER_
             ImageThunk++;
             }
-        
+
             i++;
             ImportDescriptor++;
 #endif
         }
-        
+
         /* Now scan how many imports we really have */
         for (i = 0, ImportSize = 0; i < Modules; i++)
         {
@@ -1710,7 +2000,7 @@ MiBuildImportsForBootDrivers(VOID)
                 ImportSize++;
             }
         }
-        
+
         /* Do we have any imports after all? */
         if (!ImportSize)
         {
@@ -1729,12 +2019,12 @@ MiBuildImportsForBootDrivers(VOID)
             LoadedImportsSize = ImportSize * sizeof(PVOID) + sizeof(SIZE_T);
             LoadedImports = ExAllocatePoolWithTag(PagedPool,
                                                   LoadedImportsSize,
-                                                  'TDmM');
+                                                  TAG_LDR_IMPORTS);
             ASSERT(LoadedImports);
-            
+
             /* Save the count */
             LoadedImports->Count = ImportSize;
-            
+
             /* Now copy all imports */
             for (i = 0, j = 0; i < Modules; i++)
             {
@@ -1750,38 +2040,39 @@ MiBuildImportsForBootDrivers(VOID)
                     j++;
                 }
             }
-            
+
             /* Should had as many entries as we expected */
             ASSERT(j == ImportSize);
             LdrEntry->LoadedImports = LoadedImports;
         }
-        
+
         /* Next */
         NextEntry = NextEntry->Flink;
     }
-    
+
     /* Free the initial array */
-    ExFreePool(EntryArray);
-    
+    ExFreePoolWithTag(EntryArray, TAG_LDR_IMPORTS);
+
     /* FIXME: Might not need to keep the HAL/Kernel imports around */
-    
+
     /* Kernel and HAL are loaded at boot */
     KernelEntry->LoadedImports = MM_SYSLDR_BOOT_LOADED;
     HalEntry->LoadedImports = MM_SYSLDR_BOOT_LOADED;
-    
+
     /* All worked well */
     return STATUS_SUCCESS;
 }
 
 VOID
 NTAPI
+INIT_FUNCTION
 MiLocateKernelSections(IN PLDR_DATA_TABLE_ENTRY LdrEntry)
 {
     ULONG_PTR DllBase;
     PIMAGE_NT_HEADERS NtHeaders;
     PIMAGE_SECTION_HEADER SectionHeader;
     ULONG Sections, Size;
-    
+
     /* Get the kernel section header */
     DllBase = (ULONG_PTR)LdrEntry->DllBase;
     NtHeaders = RtlImageNtHeader((PVOID)DllBase);
@@ -1793,7 +2084,7 @@ MiLocateKernelSections(IN PLDR_DATA_TABLE_ENTRY LdrEntry)
     {
         /* Grab the size of the section */
         Size = max(SectionHeader->SizeOfRawData, SectionHeader->Misc.VirtualSize);
-        
+
         /* Check for .RSRC section */
         if (*(PULONG)SectionHeader->Name == 'rsr.')
         {
@@ -1816,7 +2107,7 @@ MiLocateKernelSections(IN PLDR_DATA_TABLE_ENTRY LdrEntry)
             {
                 /* Found Mm* Pool code */
                 MmPoolCodeStart = DllBase + SectionHeader->VirtualAddress;
-                MmPoolCodeEnd = ExPoolCodeStart + Size;                
+                MmPoolCodeEnd = ExPoolCodeStart + Size;
             }
         }
         else if ((*(PULONG)SectionHeader->Name == 'YSIM') &&
@@ -1826,7 +2117,7 @@ MiLocateKernelSections(IN PLDR_DATA_TABLE_ENTRY LdrEntry)
             MmPteCodeStart = DllBase + SectionHeader->VirtualAddress;
             MmPteCodeEnd = ExPoolCodeStart + Size;
         }
-        
+
         /* Keep going */
         Sections--;
         SectionHeader++;
@@ -1835,6 +2126,7 @@ MiLocateKernelSections(IN PLDR_DATA_TABLE_ENTRY LdrEntry)
 
 BOOLEAN
 NTAPI
+INIT_FUNCTION
 MiInitializeLoadedModuleList(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
 {
     PLDR_DATA_TABLE_ENTRY LdrEntry, NewEntry;
@@ -1853,7 +2145,7 @@ MiInitializeLoadedModuleList(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
                                  LDR_DATA_TABLE_ENTRY,
                                  InLoadOrderLinks);
     PsNtosImageBase = (ULONG_PTR)LdrEntry->DllBase;
-    
+
     /* Locate resource section, pool code, and system pte code */
     MiLocateKernelSections(LdrEntry);
 
@@ -1877,7 +2169,7 @@ MiInitializeLoadedModuleList(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
         EntrySize = sizeof(LDR_DATA_TABLE_ENTRY) +
                     LdrEntry->BaseDllName.MaximumLength +
                     sizeof(UNICODE_NULL);
-        NewEntry = ExAllocatePoolWithTag(NonPagedPool, EntrySize, TAG_LDR_WSTR);
+        NewEntry = ExAllocatePoolWithTag(NonPagedPool, EntrySize, TAG_MODULE_OBJECT);
         if (!NewEntry) return FALSE;
 
         /* Copy the entry over */
@@ -1887,9 +2179,13 @@ MiInitializeLoadedModuleList(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
         NewEntry->FullDllName.Buffer =
             ExAllocatePoolWithTag(PagedPool,
                                   LdrEntry->FullDllName.MaximumLength +
-                                  sizeof(UNICODE_NULL),
+                                      sizeof(UNICODE_NULL),
                                   TAG_LDR_WSTR);
-        if (!NewEntry->FullDllName.Buffer) return FALSE;
+        if (!NewEntry->FullDllName.Buffer)
+        {
+            ExFreePoolWithTag(NewEntry, TAG_MODULE_OBJECT);
+            return FALSE;
+        }
 
         /* Set the base name */
         NewEntry->BaseDllName.Buffer = (PVOID)(NewEntry + 1);
@@ -1965,11 +2261,11 @@ MiUseLargeDriverPage(IN ULONG NumberOfPtes,
             /* Keep trying */
             NextEntry = NextEntry->Flink;
         }
-        
+
         /* If we didn't find the driver, it doesn't need large pages */
         if (DriverFound == FALSE) return FALSE;
     }
+
     /* Nothing to do yet */
     DPRINT1("Large pages not supported!\n");
     return FALSE;
@@ -1981,13 +2277,13 @@ MiComputeDriverProtection(IN BOOLEAN SessionSpace,
                           IN ULONG SectionProtection)
 {
     ULONG Protection = MM_ZERO_ACCESS;
-    
+
     /* Check if the caller gave anything */
     if (SectionProtection)
     {
         /* Always turn on execute access */
         SectionProtection |= IMAGE_SCN_MEM_EXECUTE;
-        
+
         /* Check if the registry setting is on or not */
         if (!MmEnforceWriteProtection)
         {
@@ -1995,11 +2291,11 @@ MiComputeDriverProtection(IN BOOLEAN SessionSpace,
             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)
     {
@@ -2015,10 +2311,10 @@ MiComputeDriverProtection(IN BOOLEAN SessionSpace,
             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;
 }
@@ -2040,20 +2336,20 @@ MiWriteProtectSystemImage(IN PVOID ImageBase)
     PIMAGE_NT_HEADERS NtHeaders;
     PIMAGE_SECTION_HEADER Section;
     PFN_NUMBER DriverPages;
-    ULONG CurrentProtection, SectionProtection, CombinedProtection, ProtectionMask;
+    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) */
     if (MI_IS_PHYSICAL_ADDRESS(ImageBase)) return;
-    
+
     /* Get the image headers */
     NtHeaders = RtlImageNtHeader(ImageBase);
     if (!NtHeaders) return;
-    
+
     /* Check if this is a session driver or not */
     if (!MI_IS_SESSION_ADDRESS(ImageBase))
     {
@@ -2067,13 +2363,13 @@ MiWriteProtectSystemImage(IN PVOID ImageBase)
         DPRINT1("Session drivers not supported\n");
         ASSERT(FALSE);
     }
-    
+
     /* 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);
@@ -2085,7 +2381,7 @@ MiWriteProtectSystemImage(IN PVOID ImageBase)
     {
         /* Get the section size */
         Size = max(Section->SizeOfRawData, Section->Misc.VirtualSize);
-        
+
         /* Get its virtual address */
         BaseAddress = (ULONG_PTR)ImageBase + Section->VirtualAddress;
         if (BaseAddress < CurrentAddress)
@@ -2094,24 +2390,24 @@ MiWriteProtectSystemImage(IN PVOID ImageBase)
             DPRINT1("Badly linked image!\n");
             return;
         }
-        
+
         /* Remember the current address */
         CurrentAddress = BaseAddress + Size - 1;
-        
+
         /* Next */
         Sections--;
         Section++;
     }
-    
+
     /* Get the number of sections and the first section header */
     Sections = NtHeaders->FileHeader.NumberOfSections;
     ASSERT(Sections != 0);
     Section = IMAGE_FIRST_SECTION(NtHeaders);
-    
+
     /* Set the address at the end to initialize the loop */
     CurrentAddress = (ULONG_PTR)Section + Sections - 1;
     CurrentProtection = IMAGE_SCN_MEM_WRITE | IMAGE_SCN_MEM_READ;
-    
+
     /* Set the PTE points for the image, and loop its sections */
     StartPte = MiAddressToPte(ImageBase);
     LastPte = StartPte + DriverPages;
@@ -2119,105 +2415,105 @@ MiWriteProtectSystemImage(IN PVOID ImageBase)
     {
         /* Get the section size */
         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 */
         if (PointerPte >= LastPte) break;
-        
+
         /* Get the requested protection from the image header */
         SectionProtection = Section->Characteristics & ProtectionMask;
         if (SectionProtection == CurrentProtection)
         {
             /* Same protection, so merge the request */
             CurrentAddress = BaseAddress + Size - 1;
-        
+
             /* Next */
             Sections--;
             Section++;
             continue;
         }
-        
+
         /* This is now a new section, so close up the old one */
         CurrentPte = MiAddressToPte(CurrentAddress);
-        
+
         /* Check for overlap */
         if (CurrentPte == PointerPte)
         {
             /* 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 */
         if (CurrentPte >= StartPte)
         {
             /* 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 = MiPteToAddress(CurrentAddress);
+    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);
@@ -2230,25 +2526,22 @@ MiSetPagingOfDriver(IN PMMPTE PointerPte,
                     IN PMMPTE LastPte)
 {
     PVOID ImageBase;
-    PETHREAD CurrentThread;
-    PFN_NUMBER PageCount = 0, PageFrameIndex;
+    PETHREAD CurrentThread = PsGetCurrentThread();
+    PFN_COUNT PageCount = 0;
+    PFN_NUMBER PageFrameIndex;
     PMMPFN Pfn1;
     PAGED_CODE();
-    
+
     /* Get the driver's base address */
     ImageBase = MiPteToAddress(PointerPte);
     ASSERT(MI_IS_SESSION_IMAGE_ADDRESS(ImageBase) == FALSE);
-    
+
     /* If this is a large page, it's stuck in physical memory */
     if (MI_IS_PHYSICAL_ADDRESS(ImageBase)) return;
 
-    /* We should lock the system working set -- we don't have one yet, so just be consistent */
-    CurrentThread = PsGetCurrentThread();
-    KeEnterGuardedRegion();
-    ASSERT((CurrentThread->OwnsSystemWorkingSetExclusive == 0) &&
-           (CurrentThread->OwnsSystemWorkingSetShared == 0));
-    CurrentThread->OwnsSystemWorkingSetExclusive = 1;
-    
+    /* Lock the working set */
+    MiLockWorkingSet(CurrentThread, &MmSystemCacheWs);
+
     /* Loop the PTEs */
     while (PointerPte <= LastPte)
     {
@@ -2258,20 +2551,18 @@ MiSetPagingOfDriver(IN PMMPTE PointerPte,
             PageFrameIndex = PFN_FROM_PTE(PointerPte);
             Pfn1 = MiGetPfnEntry(PageFrameIndex);
             ASSERT(Pfn1->u2.ShareCount == 1);
-            
+
             /* No working sets in ReactOS yet */
             PageCount++;
         }
-        
+
         ImageBase = (PVOID)((ULONG_PTR)ImageBase + PAGE_SIZE);
         PointerPte++;
     }
-    
-    /* Release the working set "lock" */
-    ASSERT(KeAreAllApcsDisabled() == TRUE);
-    CurrentThread->OwnsSystemWorkingSetExclusive = 0;
-    KeLeaveGuardedRegion();
-    
+
+    /* Release the working set */
+    MiUnlockWorkingSet(CurrentThread, &MmSystemCacheWs);
+
     /* Do we have any driver pages? */
     if (PageCount)
     {
@@ -2290,16 +2581,16 @@ MiEnablePagingOfDriver(IN PLDR_DATA_TABLE_ENTRY LdrEntry)
     PIMAGE_SECTION_HEADER Section;
     PMMPTE PointerPte = NULL, LastPte = NULL;
     if (MmDisablePagingExecutive) return;
-    
+
     /* Get the driver base address and its NT header */
     ImageBase = (ULONG_PTR)LdrEntry->DllBase;
     NtHeaders = RtlImageNtHeader((PVOID)ImageBase);
     if (!NtHeaders) return;
-    
+
     /* Get the sections and their alignment */
     Sections = NtHeaders->FileHeader.NumberOfSections;
     Alignment = NtHeaders->OptionalHeader.SectionAlignment - 1;
-    
+
     /* Loop each section */
     Section = IMAGE_FIRST_SECTION(NtHeaders);
     while (Sections)
@@ -2316,10 +2607,10 @@ MiEnablePagingOfDriver(IN PLDR_DATA_TABLE_ENTRY LdrEntry)
                                                            Section->
                                                            VirtualAddress));
             }
-            
+
             /* Compute the size */
             Size = max(Section->SizeOfRawData, Section->Misc.VirtualSize);
-            
+
             /* Find the last PTE that maps this section */
             LastPte = MiAddressToPte(ImageBase +
                                      Section->VirtualAddress +
@@ -2337,12 +2628,12 @@ MiEnablePagingOfDriver(IN PLDR_DATA_TABLE_ENTRY LdrEntry)
                 PointerPte = NULL;
             }
         }
-        
+
         /* Keep searching */
         Sections--;
         Section++;
     }
-    
+
     /* Handle the straggler */
     if (PointerPte) MiSetPagingOfDriver(PointerPte, LastPte);
 }
@@ -2386,7 +2677,7 @@ MmCheckSystemImage(IN HANDLE ImageHandle,
     PIMAGE_NT_HEADERS NtHeaders;
     OBJECT_ATTRIBUTES ObjectAttributes;
     PAGED_CODE();
-    
+
     /* Setup the object attributes */
     InitializeObjectAttributes(&ObjectAttributes,
                                NULL,
@@ -2402,7 +2693,11 @@ MmCheckSystemImage(IN HANDLE ImageHandle,
                              PAGE_EXECUTE,
                              SEC_IMAGE,
                              ImageHandle);
-    if (!NT_SUCCESS(Status)) return Status;
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("ZwCreateSection failed with status 0x%x\n", Status);
+        return Status;
+    }
 
     /* Make sure we're in the system process */
     KeStackAttachProcess(&PsInitialSystemProcess->Pcb, &ApcState);
@@ -2421,6 +2716,7 @@ MmCheckSystemImage(IN HANDLE ImageHandle,
     if (!NT_SUCCESS(Status))
     {
         /* We failed, close the handle and return */
+        DPRINT1("ZwMapViewOfSection failed with status 0x%x\n", Status);
         KeUnstackDetachProcess(&ApcState);
         ZwClose(SectionHandle);
         return Status;
@@ -2444,7 +2740,7 @@ MmCheckSystemImage(IN HANDLE ImageHandle,
             Status = STATUS_IMAGE_CHECKSUM_MISMATCH;
             goto Fail;
         }
-        
+
         /* Make sure it's a real image */
         NtHeaders = RtlImageNtHeader(ViewBase);
         if (!NtHeaders)
@@ -2453,7 +2749,7 @@ MmCheckSystemImage(IN HANDLE ImageHandle,
             Status = STATUS_IMAGE_CHECKSUM_MISMATCH;
             goto Fail;
         }
-        
+
         /* Make sure it's for the correct architecture */
         if ((NtHeaders->FileHeader.Machine != IMAGE_FILE_MACHINE_NATIVE) ||
             (NtHeaders->OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR_MAGIC))
@@ -2521,7 +2817,7 @@ MmLoadSystemImage(IN PUNICODE_STRING FileName,
     }
 
     /* Allocate a buffer we'll use for names */
-    Buffer = ExAllocatePoolWithTag(NonPagedPool, MAX_PATH, 'nLmM');
+    Buffer = ExAllocatePoolWithTag(NonPagedPool, MAX_PATH, TAG_LDR_WSTR);
     if (!Buffer) return STATUS_INSUFFICIENT_RESOURCES;
 
     /* Check for a separator */
@@ -2565,7 +2861,7 @@ MmLoadSystemImage(IN PUNICODE_STRING FileName,
 
     /* Check if we already have a name, use it instead */
     if (LoadedName) BaseName = *LoadedName;
-    
+
     /* Check for loader snap debugging */
     if (NtGlobalFlag & FLG_SHOW_LDR_SNAPS)
     {
@@ -2662,7 +2958,11 @@ LoaderScan:
                             &IoStatusBlock,
                             FILE_SHARE_READ | FILE_SHARE_DELETE,
                             0);
-        if (!NT_SUCCESS(Status)) goto Quickie;
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("ZwOpenFile failed with status 0x%x\n", Status);
+            goto Quickie;
+        }
 
         /* Validate it */
         Status = MmCheckSystemImage(FileHandle, FALSE);
@@ -2701,7 +3001,11 @@ LoaderScan:
                                  PAGE_EXECUTE,
                                  SEC_IMAGE,
                                  FileHandle);
-        if (!NT_SUCCESS(Status)) goto Quickie;
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("ZwCreateSection failed with status 0x%x\n", Status);
+            goto Quickie;
+        }
 
         /* Now get the section pointer */
         Status = ObReferenceObjectByHandle(SectionHandle,
@@ -2758,9 +3062,13 @@ LoaderScan:
         ObDereferenceObject(Section);
         Section = NULL;
     }
-    
+
     /* Check for failure of the load earlier */
-    if (!NT_SUCCESS(Status)) goto Quickie;
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("MiLoadImageSection failed with status 0x%x\n", Status);
+        goto Quickie;
+    }
 
     /* Relocate the driver */
     Status = LdrRelocateImageWithBias(ModuleLoadBase,
@@ -2769,9 +3077,12 @@ LoaderScan:
                                       STATUS_SUCCESS,
                                       STATUS_CONFLICTING_ADDRESSES,
                                       STATUS_INVALID_IMAGE_FORMAT);
-    if (!NT_SUCCESS(Status)) goto Quickie;
+    if (!NT_SUCCESS(Status))
+    {
+        DPRINT1("LdrRelocateImageWithBias failed with status 0x%x\n", Status);
+        goto Quickie;
+    }
 
-    
     /* Get the NT Header */
     NtHeader = RtlImageNtHeader(ModuleLoadBase);
 
@@ -2859,6 +3170,8 @@ LoaderScan:
                                       &LoadedImports);
     if (!NT_SUCCESS(Status))
     {
+        DPRINT1("MiResolveImageReferences failed with status 0x%x\n", Status);
+
         /* Fail */
         MiProcessLoaderEntry(LdrEntry, FALSE);
 
@@ -2866,7 +3179,7 @@ LoaderScan:
         if (LdrEntry->FullDllName.Buffer)
         {
             /* Free it */
-            ExFreePool(LdrEntry->FullDllName.Buffer);
+            ExFreePoolWithTag(LdrEntry->FullDllName.Buffer, TAG_LDR_WSTR);
         }
 
         /* Free the entry itself */
@@ -2960,8 +3273,8 @@ Quickie:
     /* If we have a file handle, close it */
     if (FileHandle) ZwClose(FileHandle);
 
-    /* Check if we had a prefix */
-    if (NamePrefix) ExFreePool(PrefixName.Buffer);
+    /* Check if we had a prefix (not supported yet - PrefixName == *FileName now) */
+    /* if (NamePrefix) ExFreePool(PrefixName.Buffer); */
 
     /* Free the name buffer and return status */
     ExFreePoolWithTag(Buffer, TAG_LDR_WSTR);
@@ -2975,7 +3288,7 @@ MiLookupDataTableEntry(IN PVOID Address)
     PLDR_DATA_TABLE_ENTRY LdrEntry, FoundEntry = NULL;
     PLIST_ENTRY NextEntry;
     PAGED_CODE();
-    
+
     /* Loop entries */
     NextEntry = PsLoadedModuleList.Flink;
     do
@@ -2984,7 +3297,7 @@ MiLookupDataTableEntry(IN PVOID Address)
         LdrEntry =  CONTAINING_RECORD(NextEntry,
                                       LDR_DATA_TABLE_ENTRY,
                                       InLoadOrderLinks);
-        
+
         /* Check if the address matches */
         if ((Address >= LdrEntry->DllBase) &&
             (Address < (PVOID)((ULONG_PTR)LdrEntry->DllBase +
@@ -2994,11 +3307,11 @@ MiLookupDataTableEntry(IN PVOID Address)
             FoundEntry = LdrEntry;
             break;
         }
-        
+
         /* Move on */
         NextEntry = NextEntry->Flink;
     } while(NextEntry != &PsLoadedModuleList);
-    
+
     /* Return the entry */
     return FoundEntry;
 }
@@ -3127,3 +3440,4 @@ MmGetSystemRoutineAddress(IN PUNICODE_STRING SystemRoutineName)
     return ProcAddress;
 }
 
+/* EOF */