1, Use #if defined() to judge the switching macro such as compiler macro.

[mirror_edk2.git] / DuetPkg / PciBusNoEnumerationDxe / PciDriverOverride.c

diff --git a/DuetPkg/PciBusNoEnumerationDxe/PciDriverOverride.c b/DuetPkg/PciBusNoEnumerationDxe/PciDriverOverride.c
index 3556d240ffd30f3df47e4b5260d807f187ce90b7..17bf212a194f431b2cbd665699b8e0a50947e7fc 100644 (file)
--- a/DuetPkg/PciBusNoEnumerationDxe/PciDriverOverride.c
+++ b/DuetPkg/PciBusNoEnumerationDxe/PciDriverOverride.c
@@ -26,7 +26,7 @@ Revision History
 EFI_STATUS\r
 EFIAPI\r
 GetDriver(\r
-  IN     struct _EFI_BUS_SPECIFIC_DRIVER_OVERRIDE_PROTOCOL  *This,\r
+  IN     EFI_BUS_SPECIFIC_DRIVER_OVERRIDE_PROTOCOL          *This,\r
   IN OUT EFI_HANDLE                                         *DriverImageHandle\r
   );\r
 \r
@@ -135,9 +135,8 @@ Returns:
 \r
 {\r
   EFI_STATUS                    Status;\r
-  EFI_IMAGE_DOS_HEADER          *DosHdr;\r
-  EFI_IMAGE_NT_HEADERS          *PeHdr;\r
   EFI_LOADED_IMAGE_PROTOCOL     *LoadedImage;\r
+  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;\r
   PCI_DRIVER_OVERRIDE_LIST      *Node;\r
 #if (EFI_SPECIFICATION_VERSION < 0x00020000)\r
   EFI_DRIVER_OS_HANDOFF_HEADER  *DriverOsHandoffHeader;\r
@@ -167,14 +166,19 @@ Returns:
 \r
   PciIoDevice->BusOverride  = TRUE;\r
 \r
-  DosHdr                    = (EFI_IMAGE_DOS_HEADER *) LoadedImage->ImageBase;\r
-  if (DosHdr->e_magic != EFI_IMAGE_DOS_SIGNATURE) {\r
+\r
+  ImageContext.Handle    = LoadedImage->ImageBase;\r
+  ImageContext.ImageRead = PeCoffLoaderImageReadFromMemory;\r
+\r
+  //\r
+  // Get information about the image \r
+  //\r
+  Status = PeCoffLoaderGetImageInfo (&ImageContext);\r
+  if (EFI_ERROR (Status)) {\r
     return EFI_SUCCESS;\r
   }\r
 \r
-  PeHdr = (EFI_IMAGE_NT_HEADERS *) ((UINTN) LoadedImage->ImageBase + DosHdr->e_lfanew);\r
-\r
-  if (PeHdr->FileHeader.Machine != EFI_IMAGE_MACHINE_EBC) {\r
+  if (ImageContext.Machine != EFI_IMAGE_MACHINE_EBC) {\r
     return EFI_SUCCESS;\r
   }\r
 \r