2 files changed, 43 insertions, 19 deletions

diff --git a/Private/NewBoot/Source/HEL/AMD64/AMD64-BootKit.cxx b/Private/NewBoot/Source/HEL/AMD64/AMD64-BootKit.cxx
index b45e661f..1533194b 100644
--- a/Private/NewBoot/Source/HEL/AMD64/AMD64-BootKit.cxx
+++ b/Private/NewBoot/Source/HEL/AMD64/AMD64-BootKit.cxx
@@ -122,39 +122,58 @@ HCore::VoidPtr BFileReader::ReadAll(SizeT &size) {
         .WriteString(mPath)
         .WriteString(L"\r\n");
 
-    UInt32 bufSz = KIB(350);
+    UInt32 *bufSz = nullptr;
     VoidPtr buf = nullptr;
 
-    BS->AllocatePool(EfiLoaderCode, bufSz, &buf);
+    BS->AllocatePool(EfiLoaderCode, sizeof(UInt32), (VoidPtr *)&bufSz);
+    *bufSz = 0;
+
+    BS->AllocatePool(EfiLoaderCode, *bufSz, &buf);
 
     if (!buf) return nullptr;
 
     EfiFileDevicePathProtocol filePath{0};
-    filePath.Proto.Length[0] = (sizeof(EfiDevicePathProtocol));
-    filePath.Proto.Length[1] = (sizeof(EfiDevicePathProtocol) + kPathLen) >> 8;
 
     filePath.Proto.Type = kEFIMediaDevicePath;
     filePath.Proto.SubType = kEFIMediaDevicePath;  // from all drives.
 
     BCopyMem(filePath.Path, mPath, kPathLen);
 
-    auto err = loadFile->LoadFile(loadFile, &filePath, false, (UInt32 *)&bufSz,
-                                  (VoidPtr *)&buf);
+    auto err = loadFile->LoadFile(loadFile, &filePath, true, bufSz, buf);
 
-    size = bufSz;
+    size = *bufSz;
 
-    if (buf) {
+    if (err == kEfiOk) {
       writer.WriteString(L"HCoreLdr: Loaded: ")
           .WriteString(mPath)
           .WriteString(L"\r\n");
     } else {
-      writer.WriteString(L"HCoreLdr: Error: ")
-          .WriteString(mPath)
-          .WriteString(L" , EFI-Code: ")
-          .WriteCharacter(err + 48)
-          .WriteString(L"\r\n");
+      BS->FreePool(buf);
+      buf = nullptr;
+
+      switch (err) {
+        case 2: {
+          writer.WriteString(L"HCoreLdr: Error: ")
+              .WriteString(mPath)
+              .WriteString(L", Code: Invalid-Parameter")
+              .WriteString(L"\r\n");
+
+          break;
+        }
+        case 14: {
+          writer.WriteString(L"HCoreLdr: Error: ")
+              .WriteString(mPath)
+              .WriteString(L" , EFI-Code: Not-Found")
+              .WriteString(L"\r\n");
+
+          break;
+        }
+      }
     }
 
+    BS->FreePool(bufSz);
+    bufSz = nullptr;
+
     return buf;
   }
 
diff --git a/Private/NewBoot/Source/HEL/AMD64/AMD64-Main.cxx b/Private/NewBoot/Source/HEL/AMD64/AMD64-Main.cxx
index 7d179e42..d7583873 100644
--- a/Private/NewBoot/Source/HEL/AMD64/AMD64-Main.cxx
+++ b/Private/NewBoot/Source/HEL/AMD64/AMD64-Main.cxx
@@ -7,11 +7,11 @@
  *      ========================================================
  */
 
-#include "EFIKit/EFI.hxx"
 #define __BOOTLOADER__ 1
 
 #include <BootKit/BootKit.hxx>
 #include <EFIKit/Api.hxx>
+#include <KernelKit/PE.hpp>
 
 // don't remove EfiGUID, it will call initializer_list!
 
@@ -27,17 +27,22 @@ EFI_EXTERN_C int EfiMain(EfiHandlePtr ImageHandle,
 
   UInt64 mapKey = 0;
 
-  BFileReader reader(L"EFI\\BOOT\\HCoreKrnl.exe");
+  BFileReader reader(L"HCoreKrnl.exe");
 
   SizeT sz = 0UL;
 
-  auto blob = reader.ReadAll(sz);
+  PEImagePtr blob = (PEImagePtr)reader.ReadAll(sz);
 
   if (!blob || sz < 1)
-    KeRuntimeStop(L"HCoreLdr_NoSuchKernel",
-                  L"Couldn't find HCoreKrnl.exe! Aborting...");
+    KeRuntimeStop(L"HCoreLdr_NoSuchKernel", L"Couldn't find HCoreKrnl.exe!");
 
-  writer.WriteString(L"HCoreLdr: Running HCoreKrnl.exe...\r\n");
+  ExecHeaderPtr headerPtr = (ExecHeaderPtr)blob;
+
+  if (blob[0] == kMagMz0 && blob[1] == kMagMz1) {
+    writer.WriteString(L"HCoreLdr: Running HCoreKrnl.exe...\r\n");
+  } else {
+    KeRuntimeStop(L"HCoreLdr_NotPE", L"Not a PE file! Aborting...");
+  }
 
   EFI::ExitBootServices(SystemTable, mapKey, ImageHandle);
   EFI::Stop();