[FREELDR] Add cleanup in PE-image-handling-related failure code paths.

boot/freeldr/freeldr/disk/scsiport.c

@@ -1653,7 +1653,7 @@ LoadBootDeviceDriver(VOID)
     Success = PeLdrLoadImage(NtBootDdPath, LoaderBootDriver, &ImageBase);
     if (!Success)
     {
-        /* That's OK. File simply doesn't exist */
+        /* That's OK, file simply doesn't exist */
         return ESUCCESS;
     }
 
@@ -1661,7 +1661,11 @@ LoadBootDeviceDriver(VOID)
     Success = PeLdrAllocateDataTableEntry(&ModuleListHead, "ntbootdd.sys",
                                           "NTBOOTDD.SYS", ImageBase, &BootDdDTE);
     if (!Success)
+    {
+        /* Cleanup and bail out */
+        MmFreeMemory(ImageBase);
         return EIO;
+    }
 
     /* Add the PE part of freeldr.sys to the list of loaded executables, it
        contains ScsiPort* exports, imported by ntbootdd.sys */
@@ -1669,20 +1673,27 @@ LoadBootDeviceDriver(VOID)
                                           "FREELDR.SYS", &__ImageBase, &FreeldrDTE);
     if (!Success)
     {
-        RemoveEntryList(&BootDdDTE->InLoadOrderLinks);
+        /* Cleanup and bail out */
+        PeLdrFreeDataTableEntry(BootDdDTE);
+        MmFreeMemory(ImageBase);
         return EIO;
     }
 
     /* Fix imports */
     Success = PeLdrScanImportDescriptorTable(&ModuleListHead, "", BootDdDTE);
+    if (!Success)
+    {
+        /* Cleanup and bail out */
+        PeLdrFreeDataTableEntry(FreeldrDTE);
+        PeLdrFreeDataTableEntry(BootDdDTE);
+        MmFreeMemory(ImageBase);
+        return EIO;
+    }
 
-    /* Now unlinkt the DTEs, they won't be valid later */
+    /* Now unlink the DTEs, they won't be valid later */
     RemoveEntryList(&BootDdDTE->InLoadOrderLinks);
     RemoveEntryList(&FreeldrDTE->InLoadOrderLinks);
 
-    if (!Success)
-        return EIO;
-
     /* Change imports to PA */
     ImportTable = (PIMAGE_IMPORT_DESCRIPTOR)RtlImageDirectoryEntryToData(VaToPa(BootDdDTE->DllBase),
         TRUE, IMAGE_DIRECTORY_ENTRY_IMPORT, &ImportTableSize);
@@ -1705,7 +1716,7 @@ LoadBootDeviceDriver(VOID)
                                                 NtHeaders->OptionalHeader.ImageBase - (ULONG_PTR)BootDdDTE->DllBase,
                                                 "FreeLdr",
                                                 TRUE,
-                                                TRUE, /* in case of conflict still return success */
+                                                TRUE, /* In case of conflict still return success */
                                                 FALSE);
     if (!Success)
         return EIO;

boot/freeldr/freeldr/include/peloader.h

@@ -40,6 +40,11 @@ PeLdrAllocateDataTableEntry(
     IN PVOID BasePA,
     OUT PLDR_DATA_TABLE_ENTRY *NewEntry);
 
+VOID
+PeLdrFreeDataTableEntry(
+    // _In_ PLIST_ENTRY ModuleListHead,
+    _In_ PLDR_DATA_TABLE_ENTRY Entry);
+
 BOOLEAN
 PeLdrScanImportDescriptorTable(
     IN OUT PLIST_ENTRY ModuleListHead,

boot/freeldr/freeldr/lib/peloader.c

@@ -115,7 +115,7 @@ PeLdrpBindImportName(
     //      DllBase, ImageBase, ThunkData, ExportDirectory, ExportSize, ProcessForwards);
 
     /* Check passed DllBase param */
-    if(DllBase == NULL)
+    if (DllBase == NULL)
     {
         WARN("DllBase == NULL!\n");
         return FALSE;
@@ -368,7 +368,7 @@ PeLdrpLoadAndScanReferencedDll(
     Success = PeLdrLoadImage(FullDllName, LoaderBootDriver, &BasePA);
     if (!Success)
     {
-        ERR("PeLdrLoadImage() failed\n");
+        ERR("PeLdrLoadImage('%s') failed\n", FullDllName);
         return Success;
     }
 
@@ -380,7 +380,9 @@ PeLdrpLoadAndScanReferencedDll(
                                           DataTableEntry);
     if (!Success)
     {
-        ERR("PeLdrAllocateDataTableEntry() failed\n");
+        /* Cleanup and bail out */
+        ERR("PeLdrAllocateDataTableEntry('%s') failed\n", FullDllName);
+        MmFreeMemory(BasePA);
         return Success;
     }
 
@@ -392,7 +394,10 @@ PeLdrpLoadAndScanReferencedDll(
     Success = PeLdrScanImportDescriptorTable(ModuleListHead, DirectoryPath, *DataTableEntry);
     if (!Success)
     {
+        /* Cleanup and bail out */
         ERR("PeLdrScanImportDescriptorTable() failed\n");
+        PeLdrFreeDataTableEntry(*DataTableEntry);
+        MmFreeMemory(BasePA);
         return Success;
     }
 
@@ -603,7 +608,7 @@ PeLdrAllocateDataTableEntry(
     PIMAGE_NT_HEADERS NtHeaders;
     USHORT Length;
 
-    TRACE("PeLdrAllocateDataTableEntry(, '%s', '%s', %p)\n",
+    TRACE("PeLdrAllocateDataTableEntry('%s', '%s', %p)\n",
           BaseDllName, FullDllName, BasePA);
 
     /* Allocate memory for a data table entry, zero-initialize it */
@@ -706,6 +711,20 @@ PeLdrAllocateDataTableEntry(
     return TRUE;
 }
 
+VOID
+PeLdrFreeDataTableEntry(
+    // _In_ PLIST_ENTRY ModuleListHead,
+    _In_ PLDR_DATA_TABLE_ENTRY Entry)
+{
+    // ASSERT(ModuleListHead);
+    ASSERT(Entry);
+
+    RemoveEntryList(&Entry->InLoadOrderLinks);
+    FrLdrHeapFree(VaToPa(Entry->FullDllName.Buffer), TAG_WLDR_NAME);
+    FrLdrHeapFree(VaToPa(Entry->BaseDllName.Buffer), TAG_WLDR_NAME);
+    FrLdrHeapFree(Entry, TAG_WLDR_DTE);
+}
+
 /*
  * PeLdrLoadImage loads the specified image from the file (it doesn't
  * perform any additional operations on the filename, just directly
@@ -730,13 +749,13 @@ PeLdrLoadImage(
     LARGE_INTEGER Position;
     ULONG i, BytesRead;
 
-    TRACE("PeLdrLoadImage(%s, %ld)\n", FileName, MemoryType);
+    TRACE("PeLdrLoadImage('%s', %ld)\n", FileName, MemoryType);
 
     /* Open the image file */
     Status = ArcOpen((PSTR)FileName, OpenReadOnly, &FileId);
     if (Status != ESUCCESS)
     {
-        WARN("ArcOpen(FileName: '%s') failed. Status: %u\n", FileName, Status);
+        WARN("ArcOpen('%s') failed. Status: %u\n", FileName, Status);
         return FALSE;
     }
 
@@ -744,8 +763,7 @@ PeLdrLoadImage(
     Status = ArcRead(FileId, HeadersBuffer, SECTOR_SIZE * 2, &BytesRead);
     if (Status != ESUCCESS)
     {
-        ERR("ArcRead(File: '%s') failed. Status: %u\n", FileName, Status);
+        ERR("ArcRead('%s') failed. Status: %u\n", FileName, Status);
-        UiMessageBox("Error reading from file.");
         ArcClose(FileId);
         return FALSE;
     }
@@ -755,7 +773,6 @@ PeLdrLoadImage(
     if (!NtHeaders)
     {
         ERR("No NT header found in \"%s\"\n", FileName);
-        UiMessageBox("Error: No NT header found.");
         ArcClose(FileId);
         return FALSE;
     }
@@ -764,7 +781,6 @@ PeLdrLoadImage(
     if (((NtHeaders->FileHeader.Characteristics & IMAGE_FILE_EXECUTABLE_IMAGE) == 0))
     {
         ERR("Not an executable image \"%s\"\n", FileName);
-        UiMessageBox("Not an executable image.");
         ArcClose(FileId);
         return FALSE;
     }
@@ -773,26 +789,25 @@ PeLdrLoadImage(
     NumberOfSections = NtHeaders->FileHeader.NumberOfSections;
     SectionHeader = IMAGE_FIRST_SECTION(NtHeaders);
 
-    /* Try to allocate this memory, if fails - allocate somewhere else */
+    /* Try to allocate this memory; if it fails, allocate somewhere else */
     PhysicalBase = MmAllocateMemoryAtAddress(NtHeaders->OptionalHeader.SizeOfImage,
                        (PVOID)((ULONG)NtHeaders->OptionalHeader.ImageBase & (KSEG0_BASE - 1)),
                        MemoryType);
 
     if (PhysicalBase == NULL)
     {
-        /* It's ok, we don't panic - let's allocate again at any other "low" place */
+        /* Don't fail, allocate again at any other "low" place */
         PhysicalBase = MmAllocateMemoryWithType(NtHeaders->OptionalHeader.SizeOfImage, MemoryType);
 
         if (PhysicalBase == NULL)
         {
             ERR("Failed to alloc %lu bytes for image %s\n", NtHeaders->OptionalHeader.SizeOfImage, FileName);
-            UiMessageBox("Failed to alloc pages for image.");
             ArcClose(FileId);
             return FALSE;
         }
     }
 
-    /* This is the real image base - in form of a virtual address */
+    /* This is the real image base, in form of a virtual address */
     VirtualBase = PaToVa(PhysicalBase);
 
     TRACE("Base PA: 0x%X, VA: 0x%X\n", PhysicalBase, VirtualBase);
@@ -805,10 +820,10 @@ PeLdrLoadImage(
         Status = ArcRead(FileId, (PUCHAR)PhysicalBase + sizeof(HeadersBuffer), NtHeaders->OptionalHeader.SizeOfHeaders - sizeof(HeadersBuffer), &BytesRead);
         if (Status != ESUCCESS)
         {
-            ERR("ArcRead(File: '%s') failed. Status: %u\n", FileName, Status);
+            ERR("ArcRead('%s') failed. Status: %u\n", FileName, Status);
-            UiMessageBox("Error reading headers.");
+            // UiMessageBox("Error reading headers.");
             ArcClose(FileId);
-            return FALSE;
+            goto Failure;
         }
     }
 
@@ -824,9 +839,6 @@ PeLdrLoadImage(
     /* Load the first section */
     SectionHeader = IMAGE_FIRST_SECTION(NtHeaders);
 
-    /* Fill output parameters */
-    *ImageBasePA = PhysicalBase;
-
     /* Walk through each section and read it (check/fix any possible
        bad situations, if they arise) */
     for (i = 0; i < NumberOfSections; i++)
@@ -868,7 +880,7 @@ PeLdrLoadImage(
             }
         }
 
-        /* Size of data is less than the virtual size - fill up the remainder with zeroes */
+        /* Size of data is less than the virtual size: fill up the remainder with zeroes */
         if (SizeOfRawData < VirtualSize)
         {
             TRACE("PeLdrLoadImage(): SORD %d < VS %d\n", SizeOfRawData, VirtualSize);
@@ -878,25 +890,35 @@ PeLdrLoadImage(
         SectionHeader++;
     }
 
-    /* We are done with the file - close it */
+    /* We are done with the file, close it */
     ArcClose(FileId);
 
-    /* If loading failed - return right now */
+    /* If loading failed, return right now */
     if (Status != ESUCCESS)
-        return FALSE;
+        goto Failure;
 
     /* Relocate the image, if it needs it */
     if (NtHeaders->OptionalHeader.ImageBase != (ULONG_PTR)VirtualBase)
     {
         WARN("Relocating %p -> %p\n", NtHeaders->OptionalHeader.ImageBase, VirtualBase);
-        return (BOOLEAN)LdrRelocateImageWithBias(PhysicalBase,
+        Status = LdrRelocateImageWithBias(PhysicalBase,
-                                                 (ULONG_PTR)VirtualBase - (ULONG_PTR)PhysicalBase,
+                                          (ULONG_PTR)VirtualBase - (ULONG_PTR)PhysicalBase,
-                                                 "FreeLdr",
+                                          "FreeLdr",
-                                                 TRUE,
+                                          ESUCCESS,
-                                                 TRUE, /* in case of conflict still return success */
+                                          ESUCCESS, /* In case of conflict still return success */
-                                                 FALSE);
+                                          ENOEXEC);
+        if (Status != ESUCCESS)
+            goto Failure;
     }
 
+    /* Fill output parameters */
+    *ImageBasePA = PhysicalBase;
+
     TRACE("PeLdrLoadImage() done, PA = %p\n", *ImageBasePA);
     return TRUE;
+
+Failure:
+    /* Cleanup and bail out */
+    MmFreeMemory(PhysicalBase);
+    return FALSE;
 }

boot/freeldr/freeldr/ntldr/winldr.c

@@ -333,7 +333,9 @@ WinLdrLoadDeviceDriver(PLIST_ENTRY LoadOrderListHead,
     Success = PeLdrAllocateDataTableEntry(LoadOrderListHead, DllName, DllName, DriverBase, DriverDTE);
     if (!Success)
     {
+        /* Cleanup and bail out */
         ERR("PeLdrAllocateDataTableEntry('%s') failed\n", DllName);
+        MmFreeMemory(DriverBase);
         return FALSE;
     }
 
@@ -345,7 +347,10 @@ WinLdrLoadDeviceDriver(PLIST_ENTRY LoadOrderListHead,
     Success = PeLdrScanImportDescriptorTable(LoadOrderListHead, FullPath, *DriverDTE);
     if (!Success)
     {
+        /* Cleanup and bail out */
         ERR("PeLdrScanImportDescriptorTable('%s') failed\n", FullPath);
+        PeLdrFreeDataTableEntry(*DriverDTE);
+        MmFreeMemory(DriverBase);
         return FALSE;
     }
 
@@ -482,7 +487,7 @@ WinLdrDetectVersion(VOID)
 }
 
 static
-BOOLEAN
+PVOID
 LoadModule(
     IN OUT PLOADER_PARAMETER_BLOCK LoaderBlock,
     IN PCCH Path,
@@ -495,7 +500,7 @@ LoadModule(
     BOOLEAN Success;
     CHAR FullFileName[MAX_PATH];
     CHAR ProgressString[256];
-    PVOID BaseAddress = NULL;
+    PVOID BaseAddress;
 
     RtlStringCbPrintfA(ProgressString, sizeof(ProgressString), "Loading %s...", File);
     if (!SosEnabled) UiDrawProgressBarCenter(Percentage, 100, ProgressString);
@@ -508,15 +513,10 @@ LoadModule(
     if (!Success)
     {
         ERR("PeLdrLoadImage('%s') failed\n", File);
-        return FALSE;
+        return NULL;
     }
     TRACE("%s loaded successfully at %p\n", File, BaseAddress);
 
-    /*
-     * Cheat about the base DLL name if we are loading
-     * the Kernel Debugger Transport DLL, to make the
-     * PE loader happy.
-     */
     Success = PeLdrAllocateDataTableEntry(&LoaderBlock->LoadOrderListHead,
                                           ImportName,
                                           FullFileName,
@@ -524,10 +524,13 @@ LoadModule(
                                           Dte);
     if (!Success)
     {
+        /* Cleanup and bail out */
         ERR("PeLdrAllocateDataTableEntry('%s') failed\n", FullFileName);
+        MmFreeMemory(BaseAddress);
+        BaseAddress = NULL;
     }
 
-    return Success;
+    return BaseAddress;
 }
 
 static
@@ -541,7 +544,8 @@ LoadWindowsCore(IN USHORT OperatingSystemVersion,
     BOOLEAN Success;
     PCSTR Option;
     ULONG OptionLength;
-    PLDR_DATA_TABLE_ENTRY HalDTE, KdComDTE = NULL;
+    PVOID KernelBase, HalBase, KdDllBase = NULL;
+    PLDR_DATA_TABLE_ENTRY HalDTE, KdDllDTE = NULL;
     CHAR DirPath[MAX_PATH];
     CHAR HalFileName[MAX_PATH];
     CHAR KernelFileName[MAX_PATH];
@@ -585,17 +589,33 @@ LoadWindowsCore(IN USHORT OperatingSystemVersion,
 
     TRACE("HAL file = '%s' ; Kernel file = '%s'\n", HalFileName, KernelFileName);
 
+    /*
+     * Load the core NT files: Kernel, HAL and KD transport DLL.
+     * Cheat about their base DLL name so as to satisfy the imports/exports,
+     * even if the corresponding underlying files do not have the same names
+     * -- this happens e.g. with UP vs. MP kernel, standard vs. ACPI hal, or
+     * different KD transport DLLs.
+     */
+
     /* Load the Kernel */
-    if (!LoadModule(LoaderBlock, DirPath, KernelFileName, "ntoskrnl.exe", LoaderSystemCode, KernelDTE, 30))
+    KernelBase = LoadModule(LoaderBlock, DirPath, KernelFileName,
+                            "ntoskrnl.exe", LoaderSystemCode, KernelDTE, 30);
+    if (!KernelBase)
     {
-        ERR("LoadModule() failed for %s\n", KernelFileName);
+        ERR("LoadModule('%s') failed\n", KernelFileName);
+        UiMessageBox("Could not load %s", KernelFileName);
         return FALSE;
     }
 
     /* Load the HAL */
-    if (!LoadModule(LoaderBlock, DirPath, HalFileName, "hal.dll", LoaderHalCode, &HalDTE, 45))
+    HalBase = LoadModule(LoaderBlock, DirPath, HalFileName,
+                         "hal.dll", LoaderHalCode, &HalDTE, 45);
+    if (!HalBase)
     {
-        ERR("LoadModule() failed for %s\n", HalFileName);
+        ERR("LoadModule('%s') failed\n", HalFileName);
+        UiMessageBox("Could not load %s", HalFileName);
+        PeLdrFreeDataTableEntry(*KernelDTE);
+        MmFreeMemory(KernelBase);
         return FALSE;
     }
 
@@ -651,10 +671,12 @@ LoadWindowsCore(IN USHORT OperatingSystemVersion,
              * Load the transport DLL. Override the base DLL name of the
              * loaded transport DLL to the default "KDCOM.DLL" name.
              */
-            if (!LoadModule(LoaderBlock, DirPath, KdDllName, "kdcom.dll", LoaderSystemCode, &KdComDTE, 60))
+            KdDllBase = LoadModule(LoaderBlock, DirPath, KdDllName,
+                                   "kdcom.dll", LoaderSystemCode, &KdDllDTE, 60);
+            if (!KdDllBase)
             {
                 /* The transport DLL being optional, just ignore the failure */
-                WARN("LoadModule() failed for %s\n", KdDllName);
+                WARN("LoadModule('%s') failed\n", KdDllName);
             }
         }
     }
@@ -730,11 +752,42 @@ LoadWindowsCore(IN USHORT OperatingSystemVersion,
     }
 
     /* Load all referenced DLLs for Kernel, HAL and Kernel Debugger Transport DLL */
-    Success  = PeLdrScanImportDescriptorTable(&LoaderBlock->LoadOrderListHead, DirPath, *KernelDTE);
+    Success = PeLdrScanImportDescriptorTable(&LoaderBlock->LoadOrderListHead, DirPath, *KernelDTE);
-    Success &= PeLdrScanImportDescriptorTable(&LoaderBlock->LoadOrderListHead, DirPath, HalDTE);
+    if (!Success)
-    if (KdComDTE)
     {
-        Success &= PeLdrScanImportDescriptorTable(&LoaderBlock->LoadOrderListHead, DirPath, KdComDTE);
+        UiMessageBox("Could not load %s", KernelFileName);
+        goto Quit;
+    }
+    Success = PeLdrScanImportDescriptorTable(&LoaderBlock->LoadOrderListHead, DirPath, HalDTE);
+    if (!Success)
+    {
+        UiMessageBox("Could not load %s", HalFileName);
+        goto Quit;
+    }
+    if (KdDllDTE)
+    {
+        Success = PeLdrScanImportDescriptorTable(&LoaderBlock->LoadOrderListHead, DirPath, KdDllDTE);
+        if (!Success)
+        {
+            UiMessageBox("Could not load %s", KdDllName);
+            goto Quit;
+        }
+    }
+
+Quit:
+    if (!Success)
+    {
+        /* Cleanup and bail out */
+        if (KdDllDTE)
+            PeLdrFreeDataTableEntry(KdDllDTE);
+        if (KdDllBase) // Optional
+            MmFreeMemory(KdDllBase);
+
+        PeLdrFreeDataTableEntry(HalDTE);
+        MmFreeMemory(HalBase);
+
+        PeLdrFreeDataTableEntry(*KernelDTE);
+        MmFreeMemory(KernelBase);
     }
 
     return Success;