3 files changed, 13 insertions, 8 deletions
diff --git a/IntelFrameworkModulePkg/Bus/Pci/PciBusDxe/PciBus.h b/IntelFrameworkModulePkg/Bus/Pci/PciBusDxe/PciBus.h
index 74a4abf4df..2c09061695 100644
--- a/IntelFrameworkModulePkg/Bus/Pci/PciBusDxe/PciBus.h
+++ b/IntelFrameworkModulePkg/Bus/Pci/PciBusDxe/PciBus.h
@@ -45,6 +45,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Library/DevicePathLib.h>
 #include <Library/PcdLib.h>
 #include <Library/PciIncompatibleDeviceSupportLib.h>
+#include <Library/PeCoffLib.h>
 
 #include <IndustryStandard/Pci.h>
 #include <IndustryStandard/PeImage.h>
diff --git a/IntelFrameworkModulePkg/Bus/Pci/PciBusDxe/PciBusDxe.inf b/IntelFrameworkModulePkg/Bus/Pci/PciBusDxe/PciBusDxe.inf
index 3c8b63cd78..8fa69ec729 100644
--- a/IntelFrameworkModulePkg/Bus/Pci/PciBusDxe/PciBusDxe.inf
+++ b/IntelFrameworkModulePkg/Bus/Pci/PciBusDxe/PciBusDxe.inf
@@ -86,7 +86,7 @@
   BaseLib
   UefiDriverEntryPoint
   DebugLib
-
+  PeCoffLib
 
 [Guids]
   gEfiPciOptionRomTableGuid                     # SOMETIMES_CONSUMED  System Table
diff --git a/IntelFrameworkModulePkg/Bus/Pci/PciBusDxe/PciDriverOverride.c b/IntelFrameworkModulePkg/Bus/Pci/PciBusDxe/PciDriverOverride.c
index 6c9d6f8079..6ec208453e 100644
--- a/IntelFrameworkModulePkg/Bus/Pci/PciBusDxe/PciDriverOverride.c
+++ b/IntelFrameworkModulePkg/Bus/Pci/PciBusDxe/PciDriverOverride.c
@@ -101,9 +101,8 @@ AddDriver (
   )
 {
   EFI_STATUS                    Status;
-  EFI_IMAGE_DOS_HEADER          *DosHdr;
-  EFI_IMAGE_NT_HEADERS          *PeHdr;
   EFI_LOADED_IMAGE_PROTOCOL     *LoadedImage;
+  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;
   PCI_DRIVER_OVERRIDE_LIST      *Node;
 
   Status = gBS->HandleProtocol (DriverImageHandle, &gEfiLoadedImageProtocolGuid, (VOID **) &LoadedImage);
@@ -123,16 +122,21 @@ AddDriver (
 
   PciIoDevice->BusOverride  = TRUE;
 
-  DosHdr                    = (EFI_IMAGE_DOS_HEADER *) LoadedImage->ImageBase;
-  if (DosHdr->e_magic != EFI_IMAGE_DOS_SIGNATURE) {
+  ImageContext.Handle    = LoadedImage->ImageBase;
+  ImageContext.ImageRead = PeCoffLoaderImageReadFromMemory;
+
+  //
+  // Get information about the image 
+  //
+  Status = PeCoffLoaderGetImageInfo (&ImageContext);
+  if (EFI_ERROR (Status)) {
     return EFI_SUCCESS;
   }
 
-  PeHdr = (EFI_IMAGE_NT_HEADERS *) ((UINTN) LoadedImage->ImageBase + DosHdr->e_lfanew);
-
-  if (PeHdr->FileHeader.Machine != EFI_IMAGE_MACHINE_EBC) {
+  if (ImageContext.Machine != EFI_IMAGE_MACHINE_EBC) {
     return EFI_SUCCESS;
   }
+
   return EFI_SUCCESS;
 }