[NTOSKRNL]
[reactos.git] / reactos / ntoskrnl / mm / ARM3 / sysldr.c
index 51d5cec..b77bc90 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 *******************************************************************/
 
@@ -95,18 +95,20 @@ 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;
+    KIRQL OldIrql;
+    PFN_NUMBER PageFrameIndex;
     PAGED_CODE();
 
     /* Detect session load */
     if (SessionLoad)
     {
         /* Fail */
-        DPRINT1("Session loading not yet supported!\n");
-        while (TRUE);
+        UNIMPLEMENTED_DBGBREAK("Session loading not yet supported!\n");
+        return STATUS_NOT_IMPLEMENTED;
     }
 
     /* Not session load, shouldn't have an entry */
@@ -150,14 +152,20 @@ 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;
+    PteCount = ROUND_TO_PAGES(Section->ImageSection->ImageInformation.ImageFileSize) >> 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;
@@ -167,31 +175,46 @@ MiLoadImageSection(IN OUT PVOID *SectionPtr,
     *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);
+    /* Lock the PFN database */
+    OldIrql = KeAcquireQueuedSpinLock(LockQueuePfnLock);
+
+    /* Some debug stuff */
+    MI_SET_USAGE(MI_USAGE_DRIVER_PAGE);
 #if MI_TRACE_PFNS
+    if (FileName->Buffer)
+    {
         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);
-        }
+        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 */
+    /* Loop the new driver PTEs */
+    TempPte = ValidKernelPte;
+    while (PointerPte < LastPte)
+    {
+        /* Make sure the PTE is not valid for whatever reason */
+        ASSERT(PointerPte->u.Hard.Valid == 0);
+
+        /* Grab a page */
+        PageFrameIndex = MiRemoveAnyPage(MI_GET_NEXT_COLOR());
+
+        /* Initialize its PFN entry */
+        MiInitializePfn(PageFrameIndex, PointerPte, TRUE);
+
+        /* Write the PTE */
+        TempPte.u.Hard.PageFrameNumber = PageFrameIndex;
         MI_WRITE_VALID_PTE(PointerPte, TempPte);
 
         /* Move on */
         PointerPte++;
     }
 
+    /* Release the PFN lock */
+    KeReleaseQueuedSpinLock(LockQueuePfnLock, OldIrql);
+
     /* Copy the image */
     RtlCopyMemory(DriverBase, Base, PteCount << PAGE_SHIFT);
 
@@ -324,7 +347,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);
     }
 
@@ -551,8 +574,10 @@ MiProcessLoaderEntry(IN PLDR_DATA_TABLE_ENTRY LdrEntry,
     OldIrql = KeAcquireSpinLockRaiseToSynch(&PsLoadedModuleSpinLock);
 
     /* Insert or remove from the list */
-    Insert ? InsertTailList(&PsLoadedModuleList, &LdrEntry->InLoadOrderLinks) :
-             RemoveEntryList(&LdrEntry->InLoadOrderLinks);
+    if (Insert)
+        InsertTailList(&PsLoadedModuleList, &LdrEntry->InLoadOrderLinks);
+    else
+        RemoveEntryList(&LdrEntry->InLoadOrderLinks);
 
     /* Release locks */
     KeReleaseSpinLock(&PsLoadedModuleSpinLock, OldIrql);
@@ -690,7 +715,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();
 
@@ -712,7 +737,7 @@ MiSnapThunk(IN PVOID DllBase,
         /* Copy the procedure name */
         RtlStringCbCopyA(*MissingApi,
                          MAXIMUM_FILENAME_LENGTH,
-                                                (PCHAR)&NameImport->Name[0]);
+                         (PCHAR)&NameImport->Name[0]);
 
         /* Setup name tables */
         DPRINT("Import name: %s\n", NameImport->Name);
@@ -797,9 +822,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 */
@@ -821,9 +846,9 @@ MiSnapThunk(IN PVOID DllBase,
                                              InLoadOrderLinks);
 
                 /* Check if it matches */
-                if (RtlPrefixString((PSTRING)&ForwarderName,
-                                    (PSTRING)&LdrEntry->BaseDllName,
-                                    TRUE))
+                if (RtlPrefixUnicodeString(&ForwarderName,
+                                           &LdrEntry->BaseDllName,
+                                           TRUE))
                 {
                     /* Get the forwarder export directory */
                     ForwardExportDirectory =
@@ -920,7 +945,7 @@ MmUnloadSystemImage(IN PVOID ImageHandle)
             /* Unload the symbols */
             DbgUnLoadImageSymbols(&TempName,
                                   BaseAddress,
-                                  (ULONG_PTR)ZwCurrentProcess());
+                                  (ULONG_PTR)PsGetCurrentProcessId());
             RtlFreeAnsiString(&TempName);
         }
     }
@@ -978,6 +1003,7 @@ MiResolveImageReferences(IN PVOID ImageBase,
                          OUT PWCHAR *MissingDriver,
                          OUT PLOAD_IMPORTS *LoadImports)
 {
+    static UNICODE_STRING DriversFolderName = RTL_CONSTANT_STRING(L"drivers\\");
     PCHAR MissingApiBuffer = *MissingApi, ImportName;
     PIMAGE_IMPORT_DESCRIPTOR ImportDescriptor, CurrentImport;
     ULONG ImportSize, ImportCount = 0, LoadedImportsSize, ExportSize;
@@ -996,6 +1022,9 @@ MiResolveImageReferences(IN PVOID ImageBase,
     DPRINT("%s - ImageBase: %p. ImageFileDirectory: %wZ\n",
            __FUNCTION__, ImageBase, ImageFileDirectory);
 
+    /* No name string buffer yet */
+    NameString.Buffer = NULL;
+
     /* Assume no imports */
     *LoadImports = MM_SYSLDR_NO_IMPORTS;
 
@@ -1058,9 +1087,8 @@ MiResolveImageReferences(IN PVOID ImageBase,
         if ((GdiLink) && (NormalLink))
         {
             /* It's not, it's importing stuff it shouldn't be! */
-            MiDereferenceImports(LoadedImports);
-            if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
-            return STATUS_PROCEDURE_NOT_FOUND;
+            Status = STATUS_PROCEDURE_NOT_FOUND;
+            goto Failure;
         }
 
         /* Check for user-mode printer or video card drivers, which don't belong */
@@ -1072,9 +1100,8 @@ MiResolveImageReferences(IN PVOID ImageBase,
             !(_strnicmp(ImportName, "gdi32", sizeof("gdi32") - 1)))
         {
             /* This is not kernel code */
-            MiDereferenceImports(LoadedImports);
-            if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
-            return STATUS_PROCEDURE_NOT_FOUND;
+            Status = STATUS_PROCEDURE_NOT_FOUND;
+            goto Failure;
         }
 
         /* Check if this is a "core" import, which doesn't get referenced */
@@ -1097,9 +1124,7 @@ MiResolveImageReferences(IN PVOID ImageBase,
         if (!NT_SUCCESS(Status))
         {
             /* Failed */
-            MiDereferenceImports(LoadedImports);
-            if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
-            return Status;
+            goto Failure;
         }
 
         /* We don't support name prefixes yet */
@@ -1154,73 +1179,97 @@ CheckDllState:
             DllName.Buffer = ExAllocatePoolWithTag(NonPagedPool,
                                                    DllName.MaximumLength,
                                                    TAG_LDR_WSTR);
-            if (DllName.Buffer)
-            {
-                /* Setup the base length and copy it */
-                DllName.Length = ImageFileDirectory->Length;
-                RtlCopyMemory(DllName.Buffer,
-                              ImageFileDirectory->Buffer,
-                              ImageFileDirectory->Length);
-
-                /* Now add the import name and null-terminate it */
-                RtlAppendUnicodeStringToString(&DllName,
-                                               &NameString);
-                DllName.Buffer[DllName.Length / sizeof(WCHAR)] = UNICODE_NULL;
-
-                /* Load the image */
-                Status = MmLoadSystemImage(&DllName,
-                                           NamePrefix,
-                                           NULL,
-                                           FALSE,
-                                           (PVOID)&DllEntry,
-                                           &DllBase);
-                if (NT_SUCCESS(Status))
-                {
-                    /* We can free the DLL Name */
-                    ExFreePoolWithTag(DllName.Buffer, TAG_LDR_WSTR);
-                }
-                else
-                {
-                    /* Fill out the information for the error */
-                    *MissingDriver = DllName.Buffer;
-                    *(PULONG)MissingDriver |= 1;
-                    *MissingApi = NULL;
-                }
-            }
-            else
+            if (!DllName.Buffer)
             {
                 /* We're out of resources */
                 Status = STATUS_INSUFFICIENT_RESOURCES;
+                goto Failure;
             }
 
-            /* Check if we're OK until now */
-            if (NT_SUCCESS(Status))
+            /* Add the import name to the base directory */
+            RtlCopyUnicodeString(&DllName, ImageFileDirectory);
+            RtlAppendUnicodeStringToString(&DllName,
+                                           &NameString);
+
+            /* Load the image */
+            Status = MmLoadSystemImage(&DllName,
+                                       NamePrefix,
+                                       NULL,
+                                       FALSE,
+                                       (PVOID *)&DllEntry,
+                                       &DllBase);
+
+            /* win32k / GDI drivers can also import from system32 folder */
+            if ((Status == STATUS_OBJECT_NAME_NOT_FOUND) &&
+                (MI_IS_SESSION_ADDRESS(ImageBase) || 1)) // HACK
             {
-                /* We're now loaded */
-                Loaded = TRUE;
+                /* Free the old name buffer */
+                ExFreePoolWithTag(DllName.Buffer, TAG_LDR_WSTR);
 
-                /* Sanity check */
-                ASSERT(DllBase = DllEntry->DllBase);
+                /* Calculate size for a string the adds 'drivers\' */
+                DllName.MaximumLength += DriversFolderName.Length;
 
-                /* Call the initialization routines */
-                Status = MmCallDllInitialize(DllEntry, &PsLoadedModuleList);
-                if (!NT_SUCCESS(Status))
+                /* Allocate the new buffer */
+                DllName.Buffer = ExAllocatePoolWithTag(NonPagedPool,
+                                                       DllName.MaximumLength,
+                                                       TAG_LDR_WSTR);
+                if (!DllName.Buffer)
                 {
-                    /* We failed, unload the image */
-                    MmUnloadSystemImage(DllEntry);
-                    while (TRUE);
-                    Loaded = FALSE;
+                    /* We're out of resources */
+                    Status = STATUS_INSUFFICIENT_RESOURCES;
+                    goto Failure;
                 }
+
+                /* Copy image directory and append 'drivers\' folder name */
+                RtlCopyUnicodeString(&DllName, ImageFileDirectory);
+                RtlAppendUnicodeStringToString(&DllName, &DriversFolderName);
+
+                /* Now add the import name */
+                RtlAppendUnicodeStringToString(&DllName, &NameString);
+
+                /* Try once again to load the image */
+                Status = MmLoadSystemImage(&DllName,
+                                           NamePrefix,
+                                           NULL,
+                                           FALSE,
+                                           (PVOID *)&DllEntry,
+                                           &DllBase);
             }
 
-            /* Check if we failed by here */
             if (!NT_SUCCESS(Status))
             {
+                /* Fill out the information for the error */
+                *MissingDriver = DllName.Buffer;
+                *(PULONG)MissingDriver |= 1;
+                *MissingApi = NULL;
+
+                DPRINT1("Failed to load dependency: %wZ\n", &DllName);
+
+                /* Don't free the name */
+                DllName.Buffer = NULL;
+
                 /* Cleanup and return */
-                RtlFreeUnicodeString(&NameString);
-                MiDereferenceImports(LoadedImports);
-                if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
-                return Status;
+                goto Failure;
+            }
+
+            /* We can free the DLL Name */
+            ExFreePoolWithTag(DllName.Buffer, TAG_LDR_WSTR);
+            DllName.Buffer = NULL;
+
+            /* We're now loaded */
+            Loaded = TRUE;
+
+            /* Sanity check */
+            ASSERT(DllBase == DllEntry->DllBase);
+
+            /* Call the initialization routines */
+            Status = MmCallDllInitialize(DllEntry, &PsLoadedModuleList);
+            if (!NT_SUCCESS(Status))
+            {
+                /* We failed, unload the image */
+                MmUnloadSystemImage(DllEntry);
+                ERROR_DBGBREAK("MmCallDllInitialize failed with status 0x%x\n", Status);
+                Loaded = FALSE;
             }
 
             /* Loop again to make sure that everything is OK */
@@ -1251,10 +1300,9 @@ CheckDllState:
         if (!ExportDirectory)
         {
             /* Cleanup and return */
-            MiDereferenceImports(LoadedImports);
-            if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
             DPRINT1("Warning: Driver failed to load, %S not found\n", *MissingDriver);
-            return STATUS_DRIVER_ENTRYPOINT_NOT_FOUND;
+            Status = STATUS_DRIVER_ENTRYPOINT_NOT_FOUND;
+            goto Failure;
         }
 
         /* Make sure we have an IAT */
@@ -1281,9 +1329,7 @@ CheckDllState:
                 if (!NT_SUCCESS(Status))
                 {
                     /* Cleanup and return */
-                    MiDereferenceImports(LoadedImports);
-                    if (LoadedImports) ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
-                    return Status;
+                    goto Failure;
                 }
 
                 /* Reset the buffer */
@@ -1360,6 +1406,264 @@ CheckDllState:
 
     /* Return success */
     return STATUS_SUCCESS;
+
+Failure:
+
+    /* Cleanup and return */
+    RtlFreeUnicodeString(&NameString);
+
+    if (LoadedImports)
+    {
+        MiDereferenceImports(LoadedImports);
+        ExFreePoolWithTag(LoadedImports, TAG_LDR_IMPORTS);
+    }
+
+    return Status;
+}
+
+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;
+
+        /* Only process boot loaded images. Other drivers are processed by
+           MmFreeDriverInitialization */
+        if (LdrEntry->Flags & LDRP_MM_LOADED)
+        {
+            /* Keep going */
+            NextEntry = NextEntry->Flink;
+            continue;
+        }
+
+        /* 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;
+
+    /* 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 */
+    MiDeleteSystemPageableVm(StartPte, PageCount, 0, NULL);
 }
 
 VOID
@@ -1377,7 +1681,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;
 
@@ -1473,8 +1777,8 @@ MiReloadBootLoadedDrivers(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
         if (!PointerPte)
         {
             /* Shouldn't happen */
-            DPRINT1("[Mm0]: Couldn't allocate driver section!\n");
-            while (TRUE);
+            ERROR_FATAL("[Mm0]: Couldn't allocate driver section!\n");
+            return;
         }
 
         /* This is the new virtual address for the module */
@@ -1526,8 +1830,8 @@ MiReloadBootLoadedDrivers(IN PLOADER_PARAMETER_BLOCK LoaderBlock)
             if (!NT_SUCCESS(Status))
             {
                 /* This shouldn't happen */
-                DPRINT1("Relocations failed!\n");
-                while (TRUE);
+                ERROR_FATAL("Relocations failed!\n");
+                return;
             }
         }
 
@@ -1609,8 +1913,8 @@ MiBuildImportsForBootDrivers(VOID)
         }
         else
         {
-            /* No referencing needed */
-            LdrEntry->LoadCount = 0;
+            /* Add a reference for all other modules as well */
+            LdrEntry->LoadCount = 1;
         }
 
         /* Remember this came from the loader */
@@ -1666,7 +1970,7 @@ MiBuildImportsForBootDrivers(VOID)
         /* 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))
         {
@@ -1724,9 +2028,8 @@ MiBuildImportsForBootDrivers(VOID)
                 if (*ImageThunk)
                 {
                     /* Should not be happening */
-                    DPRINT1("Broken IAT entry for %p at %p (%lx)\n",
-                            LdrEntry, ImageThunk, *ImageThunk);
-                    ASSERT(FALSE);
+                    ERROR_FATAL("Broken IAT entry for %p at %p (%lx)\n",
+                                LdrEntry, ImageThunk, *ImageThunk);
                 }
 
                 /* Reset if we hit this */
@@ -1933,9 +2236,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);
@@ -2110,8 +2417,7 @@ MiWriteProtectSystemImage(IN PVOID ImageBase)
     else
     {
         /* Not supported */
-        DPRINT1("Session drivers not supported\n");
-        ASSERT(FALSE);
+        UNIMPLEMENTED_DBGBREAK("Session drivers not supported\n");
     }
 
     /* These are the only protection masks we care about */
@@ -2277,10 +2583,15 @@ MiSetPagingOfDriver(IN PMMPTE PointerPte,
 {
     PVOID ImageBase;
     PETHREAD CurrentThread = PsGetCurrentThread();
-    PFN_NUMBER PageCount = 0, PageFrameIndex;
+    PFN_COUNT PageCount = 0;
+    PFN_NUMBER PageFrameIndex;
     PMMPFN Pfn1;
     PAGED_CODE();
 
+    /* The page fault handler is broken and doesn't page back in! */
+    DPRINT1("WARNING: MiSetPagingOfDriver() called, but paging is broken! ignoring!\n");
+    return;
+
     /* Get the driver's base address */
     ImageBase = MiPteToAddress(PointerPte);
     ASSERT(MI_IS_SESSION_IMAGE_ADDRESS(ImageBase) == FALSE);
@@ -2442,7 +2753,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);
@@ -2461,6 +2776,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;
@@ -2665,8 +2981,8 @@ LoaderScan:
         else
         {
             /* We don't support session loading yet */
-            DPRINT1("Unsupported Session-Load!\n");
-            while (TRUE);
+            UNIMPLEMENTED_DBGBREAK("Unsupported Session-Load!\n");
+            Status = STATUS_NOT_IMPLEMENTED;
         }
 
         /* Do cleanup */
@@ -2702,7 +3018,12 @@ LoaderScan:
                             &IoStatusBlock,
                             FILE_SHARE_READ | FILE_SHARE_DELETE,
                             0);
-        if (!NT_SUCCESS(Status)) goto Quickie;
+        if (!NT_SUCCESS(Status))
+        {
+            DPRINT1("ZwOpenFile failed for '%wZ' with status 0x%x\n",
+                    FileName, Status);
+            goto Quickie;
+        }
 
         /* Validate it */
         Status = MmCheckSystemImage(FileHandle, FALSE);
@@ -2741,7 +3062,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,
@@ -2757,8 +3082,8 @@ LoaderScan:
         if (Flags)
         {
             /* We don't support session loading yet */
-            DPRINT1("Unsupported Session-Load!\n");
-            while (TRUE);
+            UNIMPLEMENTED_DBGBREAK("Unsupported Session-Load!\n");
+            goto Quickie;
         }
 
         /* Check the loader list again, we should end up in the path below */
@@ -2779,7 +3104,7 @@ LoaderScan:
     ASSERT(Status != STATUS_ALREADY_COMMITTED);
 
     /* Get the size of the driver */
-    DriverSize = ((PROS_SECTION_OBJECT)Section)->ImageSection->ImageSize;
+    DriverSize = ((PROS_SECTION_OBJECT)Section)->ImageSection->ImageInformation.ImageFileSize;
 
     /* Make sure we're not being loaded into session space */
     if (!Flags)
@@ -2800,7 +3125,11 @@ LoaderScan:
     }
 
     /* 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,
@@ -2809,8 +3138,11 @@ 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);
@@ -2891,6 +3223,7 @@ LoaderScan:
 
     /* Resolve imports */
     MissingApiName = Buffer;
+    MissingDriverName = NULL;
     Status = MiResolveImageReferences(ModuleLoadBase,
                                       &BaseDirectory,
                                       NULL,
@@ -2899,6 +3232,24 @@ LoaderScan:
                                       &LoadedImports);
     if (!NT_SUCCESS(Status))
     {
+        BOOLEAN NeedToFreeString = FALSE;
+
+        /* If the lowest bit is set to 1, this is a hint that we need to free */
+        if (*(ULONG_PTR*)&MissingDriverName & 1)
+        {
+            NeedToFreeString = TRUE;
+            *(ULONG_PTR*)&MissingDriverName &= ~1;
+        }
+
+        DPRINT1("MiResolveImageReferences failed with status 0x%x\n", Status);
+        DPRINT1(" Missing driver '%ls', missing API '%s'\n",
+                MissingDriverName, MissingApiName);
+
+        if (NeedToFreeString)
+        {
+            ExFreePoolWithTag(MissingDriverName, TAG_LDR_WSTR);
+        }
+
         /* Fail */
         MiProcessLoaderEntry(LdrEntry, FALSE);
 
@@ -2975,7 +3326,7 @@ LoaderScan:
         /* Notify the debugger */
         DbgLoadImageSymbols(&AnsiTemp,
                             LdrEntry->DllBase,
-                            (ULONG_PTR)ZwCurrentProcess());
+                            (ULONG_PTR)PsGetCurrentProcessId());
         LdrEntry->Flags |= LDRP_DEBUG_SYMBOLS_LOADED;
     }
 
@@ -3167,3 +3518,4 @@ MmGetSystemRoutineAddress(IN PUNICODE_STRING SystemRoutineName)
     return ProcAddress;
 }
 
+/* EOF */