SecurePE: Replaced old PE loader with Secure one.

ArmPkg/ArmPkg.dsc

@@ -54,9 +54,9 @@
   UefiLib|MdePkg/Library/UefiLib/UefiLib.inf
   DevicePathLib|MdePkg/Library/UefiDevicePathLib/UefiDevicePathLib.inf
   UefiRuntimeServicesTableLib|MdePkg/Library/UefiRuntimeServicesTableLib/UefiRuntimeServicesTableLib.inf
-  PeCoffLib|MdePkg/Library/BasePeCoffLib/BasePeCoffLib.inf
+  PeCoffLib2|MdePkg/Library/BasePeCoffLib2/BasePeCoffLib2.inf
-  PeCoffGetEntryPointLib|MdePkg/Library/BasePeCoffGetEntryPointLib/BasePeCoffGetEntryPointLib.inf
+  UefiImageLib|MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibPeCoff.inf
-  PeCoffExtraActionLib|MdePkg/Library/BasePeCoffExtraActionLibNull/BasePeCoffExtraActionLibNull.inf
+  UefiImageExtraActionLib|MdePkg/Library/BaseUefiImageExtraActionLibNull/BaseUefiImageExtraActionLibNull.inf
 
   UefiHiiServicesLib|MdeModulePkg/Library/UefiHiiServicesLib/UefiHiiServicesLib.inf
   HiiLib|MdeModulePkg/Library/UefiHiiLib/UefiHiiLib.inf
@@ -112,7 +112,7 @@
   ArmPkg/Library/ArmCacheMaintenanceLib/ArmCacheMaintenanceLib.inf
   ArmPkg/Library/ArmPsciResetSystemLib/ArmPsciResetSystemLib.inf
   ArmPkg/Library/DebugAgentSymbolsBaseLib/DebugAgentSymbolsBaseLib.inf
-  ArmPkg/Library/DebugPeCoffExtraActionLib/DebugPeCoffExtraActionLib.inf
+  ArmPkg/Library/DebugUefiImageExtraActionLib/DebugUefiImageExtraActionLib.inf
   ArmPkg/Library/DefaultExceptionHandlerLib/DefaultExceptionHandlerLib.inf
   ArmPkg/Library/SemiHostingDebugLib/SemiHostingDebugLib.inf
   ArmPkg/Library/SemiHostingSerialPortLib/SemiHostingSerialPortLib.inf

ArmPkg/Drivers/ArmCrashDumpDxe/ArmCrashDumpDxe.dsc

@@ -34,7 +34,6 @@
   DevicePathLib|MdePkg/Library/UefiDevicePathLib/UefiDevicePathLib.inf
   MemoryAllocationLib|MdePkg/Library/UefiMemoryAllocationLib/UefiMemoryAllocationLib.inf
   PcdLib|MdePkg/Library/BasePcdLibNull/BasePcdLibNull.inf
-  PeCoffGetEntryPointLib|MdePkg/Library/BasePeCoffGetEntryPointLib/BasePeCoffGetEntryPointLib.inf
   PrintLib|MdePkg/Library/BasePrintLib/BasePrintLib.inf
   SerialPortLib|MdePkg/Library/BaseSerialPortLibNull/BaseSerialPortLibNull.inf
   UefiBootServicesTableLib|MdePkg/Library/UefiBootServicesTableLib/UefiBootServicesTableLib.inf

ArmPkg/Drivers/CpuDxe/CpuDxe.h

@@ -20,7 +20,6 @@
 #include <Library/UefiBootServicesTableLib.h>
 #include <Library/DxeServicesTableLib.h>
 #include <Library/CacheMaintenanceLib.h>
-#include <Library/PeCoffGetEntryPointLib.h>
 #include <Library/UefiLib.h>
 #include <Library/CpuLib.h>
 #include <Library/DefaultExceptionHandlerLib.h>

ArmPkg/Drivers/CpuDxe/CpuDxe.inf

@@ -49,7 +49,6 @@
   DxeServicesTableLib
   HobLib
   MemoryAllocationLib
-  PeCoffGetEntryPointLib
   UefiDriverEntryPoint
   UefiLib
 

ArmPkg/Include/Library/ArmStandaloneMmCoreEntryPoint.h

@@ -25,7 +25,7 @@
 
 #include <Library/ArmSvcLib.h>
 #include <Library/ArmFfaLib.h>
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
 #include <Library/FvLib.h>
 
 #define CPU_INFO_FLAG_PRIMARY_CPU  0x00000001
@@ -157,56 +157,6 @@ typedef RETURN_STATUS (*REGION_PERMISSION_UPDATE_FUNC) (
   IN  UINT64                Length
   );
 
-/**
-  Privileged firmware assigns RO & Executable attributes to all memory occupied
-  by the Boot Firmware Volume. This function sets the correct permissions of
-  sections in the Standalone MM Core module to be able to access RO and RW data
-  and make further progress in the boot process.
-
-  @param  [in] ImageContext           Pointer to PE/COFF image context
-  @param  [in] ImageBase              Base of image in memory
-  @param  [in] SectionHeaderOffset    Offset of PE/COFF image section header
-  @param  [in] NumberOfSections       Number of Sections
-  @param  [in] TextUpdater            Function to change code permissions
-  @param  [in] ReadOnlyUpdater        Function to change RO permissions
-  @param  [in] ReadWriteUpdater       Function to change RW permissions
-
-**/
-EFI_STATUS
-EFIAPI
-UpdateMmFoundationPeCoffPermissions (
-  IN  CONST PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext,
-  IN  EFI_PHYSICAL_ADDRESS                ImageBase,
-  IN  UINT32                              SectionHeaderOffset,
-  IN  CONST  UINT16                       NumberOfSections,
-  IN  REGION_PERMISSION_UPDATE_FUNC       TextUpdater,
-  IN  REGION_PERMISSION_UPDATE_FUNC       ReadOnlyUpdater,
-  IN  REGION_PERMISSION_UPDATE_FUNC       ReadWriteUpdater
-  );
-
-/**
-  Privileged firmware assigns RO & Executable attributes to all memory occupied
-  by the Boot Firmware Volume. This function locates the section information of
-  the Standalone MM Core module to be able to change permissions of the
-  individual sections later in the boot process.
-
-  @param  [in]      TeData                Pointer to PE/COFF image data
-  @param  [in, out] ImageContext          Pointer to PE/COFF image context
-  @param  [out]     ImageBase             Pointer to ImageBase variable
-  @param  [in, out] SectionHeaderOffset   Offset of PE/COFF image section header
-  @param  [in, out] NumberOfSections      Number of Sections
-
-**/
-EFI_STATUS
-EFIAPI
-GetStandaloneMmCorePeCoffSections (
-  IN        VOID                          *TeData,
-  IN  OUT   PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext,
-  OUT   EFI_PHYSICAL_ADDRESS              *ImageBase,
-  IN  OUT   UINT32                        *SectionHeaderOffset,
-  IN  OUT   UINT16                        *NumberOfSections
-  );
-
 /**
   Privileged firmware assigns RO & Executable attributes to all memory occupied
   by the Boot Firmware Volume. This function locates the Standalone MM Core
@@ -220,10 +170,10 @@ GetStandaloneMmCorePeCoffSections (
 **/
 EFI_STATUS
 EFIAPI
-LocateStandaloneMmCorePeCoffData (
+LocateStandaloneMmCoreUefiImage (
   IN        EFI_FIRMWARE_VOLUME_HEADER  *BfvAddress,
   IN  OUT   VOID                        **TeData,
-  IN  OUT   UINTN                       *TeDataSize
+  IN  OUT   UINT32                      *TeDataSize
   );
 
 /**

ArmPkg/Library/ArmStandaloneMmCoreEntryPoint/ArmStandaloneMmCoreEntryPoint.c

@@ -37,6 +37,7 @@
 #include <Library/SerialPortLib.h>
 #include <Library/StandaloneMmMmuLib.h>
 #include <Library/PcdLib.h>
+#include <Library/MemoryAllocationLib.h>
 
 #include <IndustryStandard/ArmStdSmc.h>
 #include <IndustryStandard/ArmMmSvc.h>
@@ -1010,13 +1011,12 @@ CEntryPoint (
   PE_COFF_LOADER_IMAGE_CONTEXT        ImageContext;
   ARM_SVC_ARGS                        EventCompleteSvcArgs;
   EFI_STATUS                          Status;
-  UINT32                              SectionHeaderOffset;
-  UINT16                              NumberOfSections;
   COMM_PROTOCOL                       CommProtocol;
   VOID                                *HobStart;
   VOID                                *TeData;
-  UINTN                               TeDataSize;
+  UINT32                              TeDataSize;
-  EFI_PHYSICAL_ADDRESS                ImageBase;
+  UINT32                              SectionIndex;
+  UEFI_IMAGE_RECORD                   *ImageRecord;
   EDKII_PI_MM_CPU_DRIVER_EP_PROTOCOL  *PiMmCpuDriverEpProtocol;
   EDKII_PI_MM_CPU_DRIVER_ENTRYPOINT   CpuDriverEntryPoint;
   EFI_HOB_FIRMWARE_VOLUME             *FvHob;
@@ -1051,7 +1051,7 @@ CEntryPoint (
   }
 
   // Locate PE/COFF File information for the Standalone MM core module
-  Status = LocateStandaloneMmCorePeCoffData (
+  Status = LocateStandaloneMmCoreUefiImage (
              (EFI_FIRMWARE_VOLUME_HEADER *)(UINTN)FvHob->BaseAddress,
              &TeData,
              &TeDataSize
@@ -1061,53 +1061,46 @@ CEntryPoint (
     goto finish;
   }
 
+  DEBUG ((DEBUG_INFO, "Found Standalone MM PE data - 0x%x\n", TeData));
+
   // Obtain the PE/COFF Section information for the Standalone MM core module
-  Status = GetStandaloneMmCorePeCoffSections (
+  Status = UefiImageInitializeContext (&ImageContext, TeData, TeDataSize);
-             TeData,
-             &ImageContext,
-             &ImageBase,
-             &SectionHeaderOffset,
-             &NumberOfSections
-             );
-
   if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "Unable to locate Standalone MM Core PE-COFF Section information - %r\n", Status));
     goto finish;
   }
 
-  //
+  ImageRecord = UefiImageLoaderGetImageRecord (&ImageContext);
-  // ImageBase may deviate from ImageContext.ImageAddress if we are dealing
-  // with a TE image, in which case the latter points to the actual offset
-  // of the image, whereas ImageBase refers to the address where the image
-  // would start if the stripped PE headers were still in place. In either
-  // case, we need to fix up ImageBase so it refers to the actual current
-  // load address.
-  //
-  ImageBase += (UINTN)TeData - ImageContext.ImageAddress;
 
-  // Update the memory access permissions of individual sections in the
+  if (ImageRecord == NULL) {
-  // Standalone MM core module
-  Status = UpdateMmFoundationPeCoffPermissions (
-             &ImageContext,
-             ImageBase,
-             SectionHeaderOffset,
-             NumberOfSections,
-             ArmSetMemoryRegionNoExec,
-             ArmSetMemoryRegionReadOnly,
-             ArmClearMemoryRegionReadOnly
-             );
-  if (EFI_ERROR (Status)) {
     goto finish;
   }
 
-  if (ImageContext.ImageAddress != (UINTN)TeData) {
+  UINT32 Address = 0;
-    ImageContext.ImageAddress = (UINTN)TeData;
+  for (SectionIndex = 0; SectionIndex < ImageRecord->NumSegments; ++ SectionIndex) {
-    ArmSetMemoryRegionNoExec (ImageBase, SIZE_4KB);
+    if ((ImageRecord->Segments[SectionIndex].Attributes & EFI_MEMORY_XP) != 0) {
-    ArmClearMemoryRegionReadOnly (ImageBase, SIZE_4KB);
+      ArmSetMemoryRegionNoExec (
+        Address,
+        ImageRecord->Segments[SectionIndex].Size
+        );
+    }
 
-    Status = PeCoffLoaderRelocateImage (&ImageContext);
+    if ((ImageRecord->Segments[SectionIndex].Attributes & EFI_MEMORY_RO) == 0) {
-    ASSERT_EFI_ERROR (Status);
+      ArmClearMemoryRegionReadOnly (
+        Address,
+        ImageRecord->Segments[SectionIndex].Size
+        );
+    }
+
+    Address += ImageRecord->Segments[SectionIndex].Size;
   }
 
+  FreePool (ImageRecord);
+
+  // FIXME: Should relocation not be performed with all of the Image writable?
+  Status = UefiImageRelocateImageInplaceForExecution (&ImageContext);
+  ASSERT_EFI_ERROR (Status);
+
   // Set the gHobList to point to the HOB list passed by TF-A.
   // This will be used by StandaloneMmCoreHobLib in early stage.
   gHobList = HobStart;

ArmPkg/Library/ArmStandaloneMmCoreEntryPoint/ArmStandaloneMmCoreEntryPoint.inf

@@ -63,11 +63,9 @@
   gArmTokenSpaceGuid.PcdStMmStackSize
 
 #
-# This configuration fails for CLANGPDB, which does not support PIE in the GCC
+# This configuration fails for CLANGDPB, which does not support PIE in the GCC
-# sense. Such however is required for ARM family StandaloneMmCore
+# sense. Such however is required for AArch64 StandaloneMmCore self-relocation,
-# self-relocation, and thus the CLANGPDB toolchain is unsupported for ARM and
+# and thus the CLANGPDB toolchain is unsupported for AArch64 for this module.
-# AARCH64 for this module.
 #
 [BuildOptions]
-  GCC:*_*_ARM_CC_FLAGS = -fpie
   GCC:*_*_AARCH64_CC_FLAGS = -fpie

ArmPkg/Library/ArmStandaloneMmCoreEntryPoint/SetPermissions.c

@@ -24,162 +24,23 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 
 #include <IndustryStandard/ArmStdSmc.h>
 
-/**
-  Privileged firmware assigns RO & Executable attributes to all memory occupied
-  by the Boot Firmware Volume. This function sets the correct permissions of
-  sections in the Standalone MM Core module to be able to access RO and RW data
-  and make further progress in the boot process.
-
-  @param  [in] ImageContext           Pointer to PE/COFF image context
-  @param  [in] ImageBase              Base of image in memory
-  @param  [in] SectionHeaderOffset    Offset of PE/COFF image section header
-  @param  [in] NumberOfSections       Number of Sections
-  @param  [in] TextUpdater            Function to change code permissions
-  @param  [in] ReadOnlyUpdater        Function to change RO permissions
-  @param  [in] ReadWriteUpdater       Function to change RW permissions
-
-**/
-EFI_STATUS
-EFIAPI
-UpdateMmFoundationPeCoffPermissions (
-  IN  CONST PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext,
-  IN  EFI_PHYSICAL_ADDRESS                ImageBase,
-  IN  UINT32                              SectionHeaderOffset,
-  IN  CONST  UINT16                       NumberOfSections,
-  IN  REGION_PERMISSION_UPDATE_FUNC       TextUpdater,
-  IN  REGION_PERMISSION_UPDATE_FUNC       ReadOnlyUpdater,
-  IN  REGION_PERMISSION_UPDATE_FUNC       ReadWriteUpdater
-  )
-{
-  EFI_IMAGE_SECTION_HEADER  SectionHeader;
-  RETURN_STATUS             Status;
-  EFI_PHYSICAL_ADDRESS      Base;
-  UINTN                     Size;
-  UINTN                     ReadSize;
-  UINTN                     Index;
-
-  ASSERT (ImageContext != NULL);
-
-  //
-  // Iterate over the sections
-  //
-  for (Index = 0; Index < NumberOfSections; Index++) {
-    //
-    // Read section header from file
-    //
-    Size     = sizeof (EFI_IMAGE_SECTION_HEADER);
-    ReadSize = Size;
-    Status   = ImageContext->ImageRead (
-                               ImageContext->Handle,
-                               SectionHeaderOffset,
-                               &Size,
-                               &SectionHeader
-                               );
-
-    if (RETURN_ERROR (Status) || (Size != ReadSize)) {
-      DEBUG ((
-        DEBUG_ERROR,
-        "%a: ImageContext->ImageRead () failed (Status = %r)\n",
-        __func__,
-        Status
-        ));
-      return Status;
-    }
-
-    DEBUG ((
-      DEBUG_INFO,
-      "%a: Section %d of image at 0x%lx has 0x%x permissions\n",
-      __func__,
-      Index,
-      ImageContext->ImageAddress,
-      SectionHeader.Characteristics
-      ));
-    DEBUG ((
-      DEBUG_INFO,
-      "%a: Section %d of image at 0x%lx has %a name\n",
-      __func__,
-      Index,
-      ImageContext->ImageAddress,
-      SectionHeader.Name
-      ));
-    DEBUG ((
-      DEBUG_INFO,
-      "%a: Section %d of image at 0x%lx has 0x%x address\n",
-      __func__,
-      Index,
-      ImageContext->ImageAddress,
-      ImageContext->ImageAddress + SectionHeader.VirtualAddress
-      ));
-    DEBUG ((
-      DEBUG_INFO,
-      "%a: Section %d of image at 0x%lx has 0x%x data\n",
-      __func__,
-      Index,
-      ImageContext->ImageAddress,
-      SectionHeader.PointerToRawData
-      ));
-
-    //
-    // If the section is marked as XN then remove the X attribute. Furthermore,
-    // if it is a writeable section then mark it appropriately as well.
-    //
-    if ((SectionHeader.Characteristics & EFI_IMAGE_SCN_MEM_EXECUTE) == 0) {
-      Base = ImageBase + SectionHeader.VirtualAddress;
-
-      TextUpdater (Base, SectionHeader.Misc.VirtualSize);
-
-      if ((SectionHeader.Characteristics & EFI_IMAGE_SCN_MEM_WRITE) != 0) {
-        ReadWriteUpdater (Base, SectionHeader.Misc.VirtualSize);
-        DEBUG ((
-          DEBUG_INFO,
-          "%a: Mapping section %d of image at 0x%lx with RW-XN permissions\n",
-          __func__,
-          Index,
-          ImageContext->ImageAddress
-          ));
-      } else {
-        DEBUG ((
-          DEBUG_INFO,
-          "%a: Mapping section %d of image at 0x%lx with RO-XN permissions\n",
-          __func__,
-          Index,
-          ImageContext->ImageAddress
-          ));
-      }
-    } else {
-      DEBUG ((
-        DEBUG_INFO,
-        "%a: Ignoring section %d of image at 0x%lx with 0x%x permissions\n",
-        __func__,
-        Index,
-        ImageContext->ImageAddress,
-        SectionHeader.Characteristics
-        ));
-    }
-
-    SectionHeaderOffset += sizeof (EFI_IMAGE_SECTION_HEADER);
-  }
-
-  return RETURN_SUCCESS;
-}
-
 /**
   Privileged firmware assigns RO & Executable attributes to all memory occupied
   by the Boot Firmware Volume. This function locates the Standalone MM Core
   module PE/COFF image in the BFV and returns this information.
 
-  @param  [in]      BfvAddress         Base Address of Boot Firmware Volume
+  @param  [in]      BfvAddress        Base Address of Boot Firmware Volume
-  @param  [in, out] TeData             Pointer to address for allocating memory
+  @param  [in, out] UefiImage         Pointer to address for allocating memory
-                                       for PE/COFF image data
+                                      for PE/COFF image data
-  @param  [in, out] TeDataSize         Pointer to size of PE/COFF image data
+  @param  [in, out] UefiImageSize     Pointer to size of PE/COFF image data
 
 **/
 EFI_STATUS
 EFIAPI
-LocateStandaloneMmCorePeCoffData (
+LocateStandaloneMmCoreUefiImage (
   IN        EFI_FIRMWARE_VOLUME_HEADER  *BfvAddress,
-  IN  OUT   VOID                        **TeData,
+  IN  OUT   VOID                            **UefiImage,
-  IN  OUT   UINTN                       *TeDataSize
+  IN  OUT   UINT32                          *UefiImageSize
   )
 {
   EFI_FFS_FILE_HEADER  *FileHeader;
@@ -201,9 +62,9 @@ LocateStandaloneMmCorePeCoffData (
     return Status;
   }
 
-  Status = FfsFindSectionData (EFI_SECTION_PE32, FileHeader, TeData, TeDataSize);
+  Status = FfsFindSectionData (EFI_SECTION_PE32, FileHeader, UefiImage, UefiImageSize);
   if (EFI_ERROR (Status)) {
-    Status = FfsFindSectionData (EFI_SECTION_TE, FileHeader, TeData, TeDataSize);
+    Status = FfsFindSectionData (EFI_SECTION_TE, FileHeader, UefiImage, UefiImageSize);
     if (EFI_ERROR (Status)) {
       DEBUG ((
         DEBUG_ERROR,
@@ -214,170 +75,6 @@ LocateStandaloneMmCorePeCoffData (
     }
   }
 
-  DEBUG ((DEBUG_INFO, "Found Standalone MM PE data - 0x%x\n", *TeData));
+  DEBUG ((DEBUG_INFO, "Found Standalone MM PE data - 0x%x\n", *UefiImage));
-  return Status;
-}
-
-/**
-  Returns the PC COFF section information.
-
-  @param  [in, out] ImageContext         Pointer to PE/COFF image context
-  @param  [out]     ImageBase            Base of image in memory
-  @param  [out]     SectionHeaderOffset  Offset of PE/COFF image section header
-  @param  [out]     NumberOfSections     Number of Sections
-
-**/
-STATIC
-EFI_STATUS
-GetPeCoffSectionInformation (
-  IN  OUT   PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext,
-  OUT   EFI_PHYSICAL_ADDRESS              *ImageBase,
-  OUT   UINT32                            *SectionHeaderOffset,
-  OUT   UINT16                            *NumberOfSections
-  )
-{
-  RETURN_STATUS                        Status;
-  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
-  EFI_IMAGE_OPTIONAL_HEADER_UNION      HdrData;
-  UINTN                                Size;
-  UINTN                                ReadSize;
-
-  ASSERT (ImageContext != NULL);
-  ASSERT (SectionHeaderOffset != NULL);
-  ASSERT (NumberOfSections != NULL);
-
-  Status = PeCoffLoaderGetImageInfo (ImageContext);
-  if (RETURN_ERROR (Status)) {
-    DEBUG ((
-      DEBUG_ERROR,
-      "%a: PeCoffLoaderGetImageInfo () failed (Status == %r)\n",
-      __func__,
-      Status
-      ));
-    return Status;
-  }
-
-  if (ImageContext->SectionAlignment < EFI_PAGE_SIZE) {
-    //
-    // The sections need to be at least 4 KB aligned, since that is the
-    // granularity at which we can tighten permissions.
-    //
-    if (!ImageContext->IsTeImage) {
-      DEBUG ((
-        DEBUG_WARN,
-        "%a: non-TE Image at 0x%lx has SectionAlignment < 4 KB (%lu)\n",
-        __func__,
-        ImageContext->ImageAddress,
-        ImageContext->SectionAlignment
-        ));
-      return RETURN_UNSUPPORTED;
-    }
-
-    ImageContext->SectionAlignment = EFI_PAGE_SIZE;
-  }
-
-  //
-  // Read the PE/COFF Header. For PE32 (32-bit) this will read in too much
-  // data, but that should not hurt anything. Hdr.Pe32->OptionalHeader.Magic
-  // determines if this is a PE32 or PE32+ image. The magic is in the same
-  // location in both images.
-  //
-  Hdr.Union = &HdrData;
-  Size      = sizeof (EFI_IMAGE_OPTIONAL_HEADER_UNION);
-  ReadSize  = Size;
-  Status    = ImageContext->ImageRead (
-                              ImageContext->Handle,
-                              ImageContext->PeCoffHeaderOffset,
-                              &Size,
-                              Hdr.Pe32
-                              );
-
-  if (RETURN_ERROR (Status) || (Size != ReadSize)) {
-    DEBUG ((
-      DEBUG_ERROR,
-      "%a: TmpContext->ImageRead () failed (Status = %r)\n",
-      __func__,
-      Status
-      ));
-    return Status;
-  }
-
-  *ImageBase = ImageContext->ImageAddress;
-  if (!ImageContext->IsTeImage) {
-    ASSERT (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE);
-
-    *SectionHeaderOffset = ImageContext->PeCoffHeaderOffset + sizeof (UINT32) +
-                           sizeof (EFI_IMAGE_FILE_HEADER);
-    *NumberOfSections = Hdr.Pe32->FileHeader.NumberOfSections;
-
-    switch (Hdr.Pe32->OptionalHeader.Magic) {
-      case EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC:
-        *SectionHeaderOffset += Hdr.Pe32->FileHeader.SizeOfOptionalHeader;
-        break;
-      case EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC:
-        *SectionHeaderOffset += Hdr.Pe32Plus->FileHeader.SizeOfOptionalHeader;
-        break;
-      default:
-        ASSERT (FALSE);
-    }
-  } else {
-    *SectionHeaderOffset = (UINTN)(sizeof (EFI_TE_IMAGE_HEADER));
-    *NumberOfSections    = Hdr.Te->NumberOfSections;
-    *ImageBase          -= (UINT32)Hdr.Te->StrippedSize - sizeof (EFI_TE_IMAGE_HEADER);
-  }
-
-  return RETURN_SUCCESS;
-}
-
-/**
-  Privileged firmware assigns RO & Executable attributes to all memory occupied
-  by the Boot Firmware Volume. This function locates the section information of
-  the Standalone MM Core module to be able to change permissions of the
-  individual sections later in the boot process.
-
-  @param  [in]      TeData                Pointer to PE/COFF image data
-  @param  [in, out] ImageContext          Pointer to PE/COFF image context
-  @param  [out]     ImageBase             Pointer to ImageBase variable
-  @param  [in, out] SectionHeaderOffset   Offset of PE/COFF image section header
-  @param  [in, out] NumberOfSections      Number of Sections
-
-**/
-EFI_STATUS
-EFIAPI
-GetStandaloneMmCorePeCoffSections (
-  IN        VOID                          *TeData,
-  IN  OUT   PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext,
-  OUT   EFI_PHYSICAL_ADDRESS              *ImageBase,
-  IN  OUT   UINT32                        *SectionHeaderOffset,
-  IN  OUT   UINT16                        *NumberOfSections
-  )
-{
-  EFI_STATUS  Status;
-
-  // Initialize the Image Context
-  ZeroMem (ImageContext, sizeof (PE_COFF_LOADER_IMAGE_CONTEXT));
-  ImageContext->Handle    = TeData;
-  ImageContext->ImageRead = PeCoffLoaderImageReadFromMemory;
-
-  DEBUG ((DEBUG_INFO, "Found Standalone MM PE data - 0x%x\n", TeData));
-
-  Status = GetPeCoffSectionInformation (
-             ImageContext,
-             ImageBase,
-             SectionHeaderOffset,
-             NumberOfSections
-             );
-  if (EFI_ERROR (Status)) {
-    DEBUG ((DEBUG_ERROR, "Unable to locate Standalone MM Core PE-COFF Section information - %r\n", Status));
-    return Status;
-  }
-
-  DEBUG ((
-    DEBUG_INFO,
-    "Standalone MM Core PE-COFF SectionHeaderOffset - 0x%x, NumberOfSections - %d\n",
-    *SectionHeaderOffset,
-    *NumberOfSections
-    ));
-
   return Status;
 }

ArmPkg/Library/DebugAgentSymbolsBaseLib/DebugAgentSymbolsBaseLib.c

@@ -13,8 +13,8 @@
 #include <Library/DebugLib.h>
 #include <Library/DebugAgentLib.h>
 #include <Library/PcdLib.h>
-#include <Library/PeCoffExtraActionLib.h>
+#include <Library/UefiImageExtraActionLib.h>
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
 
 #include <Pi/PiFirmwareFile.h>
 #include <Pi/PiFirmwareVolume.h>
@@ -165,34 +165,39 @@ GetFfsFile (
 
 EFI_STATUS
 GetImageContext (
-  IN  EFI_FFS_FILE_HEADER           *FfsHeader,
+  IN  EFI_FFS_FILE_HEADER              *FfsHeader,
-  OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   )
 {
-  EFI_STATUS                       Status;
+  EFI_STATUS                              Status;
-  UINTN                            ParsedLength;
+  UINTN                                   ParsedLength;
-  UINTN                            SectionSize;
+  UINT32                                  SectionSize;
-  UINTN                            SectionLength;
+  UINT32                                  SectionLength;
-  EFI_COMMON_SECTION_HEADER        *Section;
+  EFI_COMMON_SECTION_HEADER               *Section;
-  VOID                             *EfiImage;
+  VOID                                    *EfiImage;
-  UINTN                            ImageAddress;
-  EFI_IMAGE_DEBUG_DIRECTORY_ENTRY  *DebugEntry;
-  VOID                             *CodeViewEntryPointer;
 
-  Section      = (EFI_COMMON_SECTION_HEADER *)(FfsHeader + 1);
+  Section       = (EFI_COMMON_SECTION_HEADER *)(FfsHeader + 1);
-  SectionSize  = *(UINT32 *)(FfsHeader->Size) & 0x00FFFFFF;
+  SectionLength = 0;
-  SectionSize -= sizeof (EFI_FFS_FILE_HEADER);
+  SectionSize   = *(UINT32 *)(FfsHeader->Size) & 0x00FFFFFF;
-  ParsedLength = 0;
+  SectionSize  -= sizeof (EFI_FFS_FILE_HEADER);
-  EfiImage     = NULL;
+  ParsedLength  = 0;
+  EfiImage      = NULL;
 
   while (ParsedLength < SectionSize) {
+    //
+    // Size is 24 bits wide so mask upper 8 bits.
+    //
+    SectionLength = *(UINT32 *)Section->Size & 0x00FFFFFF;
+    if (SectionLength < sizeof (*Section)) {
+      return EFI_VOLUME_CORRUPTED;
+    }
+
     if ((Section->Type == EFI_SECTION_PE32) || (Section->Type == EFI_SECTION_TE)) {
-      EfiImage = (EFI_IMAGE_OPTIONAL_HEADER_UNION *)(Section + 1);
+      EfiImage = (Section + 1);
       break;
     }
 
     //
-    // Size is 24 bits wide so mask upper 8 bits.
     // SectionLength is adjusted it is 4 byte aligned.
     // Go to the next section
     //
@@ -208,34 +213,10 @@ GetImageContext (
   }
 
   // Initialize the Image Context
-  ZeroMem (ImageContext, sizeof (PE_COFF_LOADER_IMAGE_CONTEXT));
+  // FIXME: Common FFS API with size checks
-  ImageContext->Handle    = EfiImage;
+  Status = UefiImageInitializeContext (ImageContext, EfiImage, SectionLength - sizeof (*Section));
-  ImageContext->ImageRead = PeCoffLoaderImageReadFromMemory;
+  if (!EFI_ERROR(Status)) {
-
+    Status = UefiImageLoadImageInplace( ImageContext);
-  Status =  PeCoffLoaderGetImageInfo (ImageContext);
-  if (!EFI_ERROR (Status) && ((VOID *)(UINTN)ImageContext->DebugDirectoryEntryRva != NULL)) {
-    ImageAddress = ImageContext->ImageAddress;
-    if (ImageContext->IsTeImage) {
-      ImageAddress += sizeof (EFI_TE_IMAGE_HEADER) - ((EFI_TE_IMAGE_HEADER *)EfiImage)->StrippedSize;
-    }
-
-    DebugEntry = (EFI_IMAGE_DEBUG_DIRECTORY_ENTRY *)(ImageAddress + ImageContext->DebugDirectoryEntryRva);
-    if (DebugEntry->Type == EFI_IMAGE_DEBUG_TYPE_CODEVIEW) {
-      CodeViewEntryPointer = (VOID *)(ImageAddress + (UINTN)DebugEntry->RVA);
-      switch (*(UINT32 *)CodeViewEntryPointer) {
-        case CODEVIEW_SIGNATURE_NB10:
-          ImageContext->PdbPointer = (CHAR8 *)CodeViewEntryPointer + sizeof (EFI_IMAGE_DEBUG_CODEVIEW_NB10_ENTRY);
-          break;
-        case CODEVIEW_SIGNATURE_RSDS:
-          ImageContext->PdbPointer = (CHAR8 *)CodeViewEntryPointer + sizeof (EFI_IMAGE_DEBUG_CODEVIEW_RSDS_ENTRY);
-          break;
-        case CODEVIEW_SIGNATURE_MTOC:
-          ImageContext->PdbPointer = (CHAR8 *)CodeViewEntryPointer + sizeof (EFI_IMAGE_DEBUG_CODEVIEW_MTOC_ENTRY);
-          break;
-        default:
-          break;
-      }
-    }
   }
 
   return Status;
@@ -271,9 +252,9 @@ InitializeDebugAgent (
   IN DEBUG_AGENT_CONTINUE  Function  OPTIONAL
   )
 {
-  EFI_STATUS                    Status;
+  EFI_STATUS                      Status;
-  EFI_FFS_FILE_HEADER           *FfsHeader;
+  EFI_FFS_FILE_HEADER             *FfsHeader;
-  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT ImageContext;
 
   // We use InitFlag to know if DebugAgent has been initialized from
   // Sec (DEBUG_AGENT_INIT_PREMEM_SEC) or PrePi (DEBUG_AGENT_INIT_POSTMEM_SEC)
@@ -286,7 +267,7 @@ InitializeDebugAgent (
     if (!EFI_ERROR (Status)) {
       Status = GetImageContext (FfsHeader, &ImageContext);
       if (!EFI_ERROR (Status)) {
-        PeCoffLoaderRelocateImageExtraAction (&ImageContext);
+        UefiImageLoaderRelocateImageExtraAction (&ImageContext);
       }
     }
   } else if (InitFlag == DEBUG_AGENT_INIT_POSTMEM_SEC) {
@@ -297,7 +278,7 @@ InitializeDebugAgent (
     if (!EFI_ERROR (Status)) {
       Status = GetImageContext (FfsHeader, &ImageContext);
       if (!EFI_ERROR (Status)) {
-        PeCoffLoaderRelocateImageExtraAction (&ImageContext);
+        UefiImageLoaderRelocateImageExtraAction (&ImageContext);
       }
     }
 
@@ -308,7 +289,7 @@ InitializeDebugAgent (
     if (!EFI_ERROR (Status)) {
       Status = GetImageContext (FfsHeader, &ImageContext);
       if (!EFI_ERROR (Status)) {
-        PeCoffLoaderRelocateImageExtraAction (&ImageContext);
+        UefiImageLoaderRelocateImageExtraAction (&ImageContext);
       }
     }
   }

ArmPkg/Library/DebugAgentSymbolsBaseLib/DebugAgentSymbolsBaseLib.inf

@@ -24,8 +24,8 @@
 [LibraryClasses]
   DebugLib
   PcdLib
-  PeCoffExtraActionLib
+  UefiImageExtraActionLib
-  PeCoffLib
+  UefiImageLib
 
 [Pcd]
   gArmTokenSpaceGuid.PcdSecureFvBaseAddress

ArmPkg/Library/DebugUefiImageExtraActionLib/DebugUefiImageExtraActionLib.c

@@ -0,0 +1,141 @@
+/**@file
+
+Copyright (c) 2006 - 2009, Intel Corporation. All rights reserved.<BR>
+Portions copyright (c) 2008 - 2010, Apple Inc. All rights reserved.<BR>
+Portions copyright (c) 2011 - 2012, ARM Ltd. All rights reserved.<BR>
+
+SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include <PiDxe.h>
+#include <Library/UefiImageLib.h>
+
+#include <Library/BaseLib.h>
+#include <Library/DebugLib.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/UefiImageExtraActionLib.h>
+#include <Library/PrintLib.h>
+
+/**
+  If the build is done on cygwin the paths are cygpaths.
+  /cygdrive/c/tmp.txt vs c:\tmp.txt so we need to convert
+  them to work with RVD commands
+
+  @param  Name  Path to convert if needed
+
+**/
+CONST CHAR8 *
+DeCygwinPathIfNeeded (
+  IN  CONST CHAR8   *Name,
+  IN  CHAR8         *Temp,
+  IN  UINTN         Size
+  )
+{
+  CHAR8  *Ptr;
+  UINTN  Index;
+  UINTN  Index2;
+
+  Ptr = AsciiStrStr (Name, "/cygdrive/");
+  if (Ptr == NULL) {
+    return Name;
+  }
+
+  for (Index = 9, Index2 = 0; (Index < (Size + 9)) && (Ptr[Index] != '\0'); Index++, Index2++) {
+    Temp[Index2] = Ptr[Index];
+    if (Temp[Index2] == '/') {
+      Temp[Index2] = '\\';
+    }
+
+    if (Index2 == 1) {
+      Temp[Index2 - 1] = Ptr[Index];
+      Temp[Index2]     = ':';
+    }
+  }
+
+  return Temp;
+}
+
+/**
+  Performs additional actions after a PE/COFF image has been loaded and relocated.
+
+  If ImageContext is NULL, then ASSERT().
+
+  @param  ImageContext  Pointer to the image context structure that describes the
+                        PE/COFF image that has already been loaded and relocated.
+
+**/
+VOID
+EFIAPI
+UefiImageLoaderRelocateImageExtraAction (
+  IN CONST UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
+  )
+{
+  RETURN_STATUS Status;
+  CONST CHAR8   *PdbPath;
+  UINT32        PdbPathSize;
+#if defined (__CC_ARM) || defined (__GNUC__)
+  CHAR8         Temp[512];
+#endif
+
+  Status = UefiImageGetSymbolsPath (ImageContext, &PdbPath, &PdbPathSize);
+
+  if (!RETURN_ERROR (Status)) {
+ #ifdef __CC_ARM
+ #if (__ARMCC_VERSION < 500000)
+    // Print out the command for the RVD debugger to load symbols for this image
+    DEBUG ((DEBUG_LOAD | DEBUG_INFO, "load /a /ni /np %a &0x%p\n", DeCygwinPathIfNeeded (PdbPath, Temp, sizeof (Temp)), UefiImageLoaderGetImageAddress (ImageContext)));
+ #else
+    // Print out the command for the DS-5 to load symbols for this image
+    DEBUG ((DEBUG_LOAD | DEBUG_INFO, "add-symbol-file %a -o 0x%p\n", DeCygwinPathIfNeeded (PdbPath, Temp, sizeof (Temp)), UefiImageLoaderGetImageAddress (ImageContext)));
+ #endif
+ #elif __GNUC__
+    // This may not work correctly if you generate PE/COFF directly as then the Offset would not be required
+    DEBUG ((DEBUG_LOAD | DEBUG_INFO, "add-symbol-file %a -o 0x%p\n", DeCygwinPathIfNeeded (PdbPath, Temp, sizeof (Temp)), UefiImageLoaderGetImageAddress (ImageContext)));
+ #else
+    DEBUG ((DEBUG_LOAD | DEBUG_INFO, "Loading driver at 0x%11p EntryPoint=0x%11p\n", (VOID *)(UINTN)UefiImageLoaderGetImageAddress (ImageContext), FUNCTION_ENTRY_POINT (UefiImageLoaderGetImageEntryPoint (ImageContext))));
+ #endif
+  } else {
+    DEBUG ((DEBUG_LOAD | DEBUG_INFO, "Loading driver at 0x%11p EntryPoint=0x%11p\n", (VOID *)(UINTN)UefiImageLoaderGetImageAddress (ImageContext), FUNCTION_ENTRY_POINT (UefiImageLoaderGetImageEntryPoint (ImageContext))));
+  }
+}
+
+/**
+  Performs additional actions just before a PE/COFF image is unloaded.  Any resources
+  that were allocated by UefiImageLoaderRelocateImageExtraAction() must be freed.
+
+  If ImageContext is NULL, then ASSERT().
+
+  @param  ImageContext  Pointer to the image context structure that describes the
+                        PE/COFF image that is being unloaded.
+
+**/
+VOID
+EFIAPI
+UefiImageLoaderUnloadImageExtraAction (
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
+  )
+{
+  RETURN_STATUS Status;
+  CONST CHAR8   *PdbPath;
+  UINT32        PdbPathSize;
+#if defined (__CC_ARM) || defined (__GNUC__)
+  CHAR8         Temp[512];
+#endif
+
+  Status = UefiImageGetSymbolsPath (ImageContext, &PdbPath, &PdbPathSize);
+
+  if (!RETURN_ERROR (Status)) {
+ #ifdef __CC_ARM
+    // Print out the command for the RVD debugger to load symbols for this image
+    DEBUG ((DEBUG_LOAD | DEBUG_INFO, "unload symbols_only %a\n", DeCygwinPathIfNeeded (PdbPath, Temp, sizeof (Temp))));
+ #elif __GNUC__
+    // This may not work correctly if you generate PE/COFF directly as then the Offset would not be required
+    DEBUG ((DEBUG_LOAD | DEBUG_INFO, "remove-symbol-file %a 0x%08x\n", DeCygwinPathIfNeeded (PdbPath, Temp, sizeof (Temp)), (UINTN)UefiImageLoaderGetImageAddress (ImageContext)));
+ #else
+    DEBUG ((DEBUG_LOAD | DEBUG_INFO, "Unloading %a\n", PdbPath));
+ #endif
+  } else {
+    DEBUG ((DEBUG_LOAD | DEBUG_INFO, "Unloading driver at 0x%11p\n", (VOID *)(UINTN)UefiImageLoaderGetImageAddress (ImageContext)));
+  }
+}

ArmPkg/Library/DebugUefiImageExtraActionLib/DebugUefiImageExtraActionLib.inf

@@ -0,0 +1,33 @@
+#/** @file
+# UEFI Image extra action library for DXE phase that run Unix emulator.
+#
+# Lib to provide memory journal status code reporting Routines
+# Copyright (c) 2007 - 2010, Intel Corporation. All rights reserved.<BR>
+# Portions copyright (c) 2010, Apple Inc. All rights reserved.<BR>
+# SPDX-License-Identifier: BSD-2-Clause-Patent
+#
+#
+#**/
+
+[Defines]
+  INF_VERSION                    = 0x00010005
+  BASE_NAME                      = DebugUnixUefiImageExtraActionLib
+  FILE_GUID                      = C3E9448E-1726-42fb-9368-41F75B038C0C
+  MODULE_TYPE                    = BASE
+  VERSION_STRING                 = 1.0
+  LIBRARY_CLASS                  = UefiImageExtraActionLib
+
+#
+# The following information is for reference only and not required by the build tools.
+#
+#  VALID_ARCHITECTURES           = ARM
+#
+
+[Sources.common]
+  DebugUefiImageExtraActionLib.c
+
+[Packages]
+  MdePkg/MdePkg.dec
+
+[LibraryClasses]
+  DebugLib

ArmPkg/Library/DefaultExceptionHandlerLib/AArch64/DefaultExceptionHandler.c

@@ -12,7 +12,6 @@
 #include <Library/UefiLib.h>
 #include <Library/BaseLib.h>
 #include <Library/DebugLib.h>
-#include <Library/PeCoffGetEntryPointLib.h>
 #include <Library/PrintLib.h>
 #include <Library/SerialPortLib.h>
 #include <Library/UefiBootServicesTableLib.h>
@@ -36,11 +35,10 @@ STATIC CHAR8  *gExceptionTypeString[] = {
 
 STATIC BOOLEAN  mRecursiveException;
 
-CHAR8 *
+CONST CHAR8 *
 GetImageName (
   IN  UINTN  FaultAddress,
-  OUT UINTN  *ImageBase,
+  OUT UINTN  *ImageBase
-  OUT UINTN  *PeCoffSizeOfHeaders
   );
 
 STATIC
@@ -211,14 +209,13 @@ DefaultExceptionHandler (
   UnicodeSPrintAsciiFormat (UnicodeBuffer, MAX_PRINT_CHARS, Buffer);
 
   DEBUG_CODE_BEGIN ();
-  CHAR8   *Pdb, *PrevPdb;
+  CONST CHAR8  *Pdb, *PrevPdb;
-  UINTN   ImageBase;
+  UINTN        ImageBase;
-  UINTN   PeCoffSizeOfHeader;
+  UINT64       *Fp;
-  UINT64  *Fp;
+  UINT64       RootFp[2];
-  UINT64  RootFp[2];
+  UINTN        Idx;
-  UINTN   Idx;
 
-  PrevPdb = Pdb = GetImageName (SystemContext.SystemContextAArch64->ELR, &ImageBase, &PeCoffSizeOfHeader);
+  PrevPdb = Pdb = GetImageName (SystemContext.SystemContextAArch64->ELR, &ImageBase);
   if (Pdb != NULL) {
     DEBUG ((
       DEBUG_ERROR,
@@ -243,7 +240,7 @@ DefaultExceptionHandler (
     }
 
     for (Fp = RootFp; Fp[0] != 0; Fp = (UINT64 *)Fp[0]) {
-      Pdb = GetImageName (Fp[1], &ImageBase, &PeCoffSizeOfHeader);
+      Pdb = GetImageName (Fp[1], &ImageBase);
       if (Pdb != NULL) {
         if (Pdb != PrevPdb) {
           Idx++;
@@ -264,14 +261,14 @@ DefaultExceptionHandler (
       }
     }
 
-    PrevPdb = Pdb = GetImageName (SystemContext.SystemContextAArch64->ELR, &ImageBase, &PeCoffSizeOfHeader);
+    PrevPdb = Pdb = GetImageName (SystemContext.SystemContextAArch64->ELR, &ImageBase);
     if (Pdb != NULL) {
       DEBUG ((DEBUG_ERROR, "\n[ 0] %a\n", Pdb));
     }
 
     Idx = 0;
     for (Fp = RootFp; Fp[0] != 0; Fp = (UINT64 *)Fp[0]) {
-      Pdb = GetImageName (Fp[1], &ImageBase, &PeCoffSizeOfHeader);
+      Pdb = GetImageName (Fp[1], &ImageBase);
       if ((Pdb != NULL) && (Pdb != PrevPdb)) {
         DEBUG ((DEBUG_ERROR, "[% 2d] %a\n", ++Idx, Pdb));
         PrevPdb = Pdb;

ArmPkg/Library/DefaultExceptionHandlerLib/Arm/DefaultExceptionHandler.c

@@ -11,7 +11,6 @@
 #include <Uefi.h>
 #include <Library/BaseLib.h>
 #include <Library/DebugLib.h>
-#include <Library/PeCoffGetEntryPointLib.h>
 #include <Library/PrintLib.h>
 #include <Library/SerialPortLib.h>
 #include <Library/UefiBootServicesTableLib.h>
@@ -53,11 +52,10 @@ STATIC CONST CPSR_CHAR  mCpsrChar[] = {
   { 0,  '?' }
 };
 
-CHAR8 *
+CONST CHAR8 *
 GetImageName (
   IN  UINTN  FaultAddress,
-  OUT UINTN  *ImageBase,
+  OUT UINTN  *ImageBase
-  OUT UINTN  *PeCoffSizeOfHeaders
   );
 
 /**
@@ -225,16 +223,16 @@ DefaultExceptionHandler (
   UnicodeSPrintAsciiFormat (UnicodeBuffer, MAX_PRINT_CHARS, Buffer);
 
   DEBUG_CODE_BEGIN ();
-  CHAR8   *Pdb;
+  CONST CHAR8   *Pdb;
-  UINT32  ImageBase;
+  UINT32        ImageBase;
-  UINT32  PeCoffSizeOfHeader;
+  UINT32        Offset;
-  UINT32  Offset;
+  CHAR8         CpsrStr[CPSR_STRING_SIZE];  // char per bit. Lower 5-bits are mode
-  CHAR8   CpsrStr[CPSR_STRING_SIZE];
+                                            // that is a 3 char string
 
   CpsrString (SystemContext.SystemContextArm->CPSR, CpsrStr);
   DEBUG ((DEBUG_ERROR, "%a\n", CpsrStr));
 
-  Pdb    = GetImageName (SystemContext.SystemContextArm->PC, &ImageBase, &PeCoffSizeOfHeader);
+  Pdb    = GetImageName (SystemContext.SystemContextArm->PC, &ImageBase);
   Offset = SystemContext.SystemContextArm->PC - ImageBase;
   if (Pdb != NULL) {
     DEBUG ((DEBUG_ERROR, "%a\n", Pdb));
@@ -247,7 +245,9 @@ DefaultExceptionHandler (
     // you need to subtract out the size of the PE/COFF header to get
     // get the offset that matches the link map.
     //
-    DEBUG ((DEBUG_ERROR, "loaded at 0x%08x (PE/COFF offset) 0x%x (ELF or Mach-O offset) 0x%x", ImageBase, Offset, Offset - PeCoffSizeOfHeader));
+    // FIXME: Used to have  (ELF or Mach-O offset) 0x%x
+    //        Substitute with .text address (better + may be needed for GDB symbols?)
+    DEBUG ((EFI_D_ERROR, "loaded at 0x%08x (PE/COFF offset) 0x%x", ImageBase, Offset));
   }
 
   DEBUG_CODE_END ();

ArmPkg/Library/DefaultExceptionHandlerLib/DefaultExceptionHandlerLib.inf

@@ -34,7 +34,6 @@
   BaseLib
   PrintLib
   DebugLib
-  PeCoffGetEntryPointLib
   SerialPortLib
   UefiBootServicesTableLib
 

ArmPkg/Library/DefaultExceptionHandlerLib/DefaultExceptionHandlerUefi.c

@@ -7,7 +7,6 @@
 **/
 
 #include <Uefi.h>
-#include <Library/PeCoffGetEntryPointLib.h>
 #include <Library/UefiLib.h>
 
 #include <Guid/DebugImageInfoTable.h>
@@ -20,17 +19,15 @@
 
   @param  FaultAddress         Address to find PE/COFF image for.
   @param  ImageBase            Return load address of found image
-  @param  PeCoffSizeOfHeaders  Return the size of the PE/COFF header for the image that was found
 
   @retval NULL                 FaultAddress not in a loaded PE/COFF image.
   @retval                      Path and file name of PE/COFF image.
 
 **/
-CHAR8 *
+CONST CHAR8 *
 GetImageName (
   IN  UINTN  FaultAddress,
-  OUT UINTN  *ImageBase,
+  OUT UINTN  *ImageBase
-  OUT UINTN  *PeCoffSizeOfHeaders
   )
 {
   EFI_STATUS                         Status;
@@ -59,8 +56,7 @@ GetImageName (
             (Address <= ((CHAR8 *)DebugTable->NormalImage->LoadedImageProtocolInstance->ImageBase + DebugTable->NormalImage->LoadedImageProtocolInstance->ImageSize)))
         {
           *ImageBase           = (UINTN)DebugTable->NormalImage->LoadedImageProtocolInstance->ImageBase;
-          *PeCoffSizeOfHeaders = PeCoffGetSizeOfHeaders ((VOID *)(UINTN)*ImageBase);
+          return DebugTable->NormalImage->PdbPath;
-          return PeCoffLoaderGetPdbPointer (DebugTable->NormalImage->LoadedImageProtocolInstance->ImageBase);
         }
       }
     }

ArmPlatformPkg/ArmPlatformPkg.dsc

@@ -52,8 +52,9 @@
   MemoryAllocationLib|MdePkg/Library/UefiMemoryAllocationLib/UefiMemoryAllocationLib.inf
   MemoryInitPeiLib|ArmPlatformPkg/MemoryInitPei/MemoryInitPeiLib.inf
   PcdLib|MdePkg/Library/BasePcdLibNull/BasePcdLibNull.inf
-  PeCoffExtraActionLib|MdePkg/Library/BasePeCoffExtraActionLibNull/BasePeCoffExtraActionLibNull.inf
+  UefiImageExtraActionLib|MdePkg/Library/BaseUefiImageExtraActionLibNull/BaseUefiImageExtraActionLibNull.inf
-  PeCoffLib|MdePkg/Library/BasePeCoffLib/BasePeCoffLib.inf
+  PeCoffLib2|MdePkg/Library/BasePeCoffLib2/BasePeCoffLib2.inf
+  UefiImageLib|MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibPeCoff.inf
   PerformanceLib|MdePkg/Library/BasePerformanceLibNull/BasePerformanceLibNull.inf
   PL011UartClockLib|ArmPlatformPkg/Library/PL011UartClockLib/PL011UartClockLib.inf
   PL011UartLib|ArmPlatformPkg/Library/PL011UartLib/PL011UartLib.inf

ArmVirtPkg/ArmVirt.dsc.inc

@@ -61,8 +61,8 @@
   SynchronizationLib|MdePkg/Library/BaseSynchronizationLib/BaseSynchronizationLib.inf
   PerformanceLib|MdePkg/Library/BasePerformanceLibNull/BasePerformanceLibNull.inf
   PrintLib|MdePkg/Library/BasePrintLib/BasePrintLib.inf
-  PeCoffGetEntryPointLib|MdePkg/Library/BasePeCoffGetEntryPointLib/BasePeCoffGetEntryPointLib.inf
+  PeCoffLib2|MdePkg/Library/BasePeCoffLib2/BasePeCoffLib2.inf
-  PeCoffLib|MdePkg/Library/BasePeCoffLib/BasePeCoffLib.inf
+  UefiImageLib|MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibPeCoff.inf
   IoLib|MdePkg/Library/BaseIoLibIntrinsic/BaseIoLibIntrinsicArmVirt.inf
   UefiDecompressLib|MdePkg/Library/BaseUefiDecompressLib/BaseUefiDecompressLib.inf
   CpuLib|MdePkg/Library/BaseCpuLib/BaseCpuLib.inf
@@ -125,8 +125,8 @@
   SerialPortLib|ArmVirtPkg/Library/FdtPL011SerialPortLib/FdtPL011SerialPortLib.inf
   FdtSerialPortAddressLib|OvmfPkg/Library/FdtSerialPortAddressLib/FdtSerialPortAddressLib.inf
 
-  PeCoffExtraActionLib|ArmPkg/Library/DebugPeCoffExtraActionLib/DebugPeCoffExtraActionLib.inf
+  UefiImageExtraActionLib|ArmPkg/Library/DebugUefiImageExtraActionLib/DebugUefiImageExtraActionLib.inf
-  #PeCoffExtraActionLib|MdePkg/Library/BasePeCoffExtraActionLibNull/BasePeCoffExtraActionLibNull.inf
+  #UefiImageExtraActionLib|MdePkg/Library/BaseUefiImageExtraActionLibNull/BaseUefiImageExtraActionLibNull.inf
 
   DebugAgentLib|MdeModulePkg/Library/DebugAgentLibNull/DebugAgentLibNull.inf
   DebugAgentTimerLib|EmbeddedPkg/Library/DebugAgentTimerLibNull/DebugAgentTimerLibNull.inf
@@ -205,7 +205,6 @@
   PeiCoreEntryPoint|MdePkg/Library/PeiCoreEntryPoint/PeiCoreEntryPoint.inf
   PerformanceLib|MdeModulePkg/Library/PeiPerformanceLib/PeiPerformanceLib.inf
   OemHookStatusCodeLib|MdeModulePkg/Library/OemHookStatusCodeLibNull/OemHookStatusCodeLibNull.inf
-  PeCoffGetEntryPointLib|MdePkg/Library/BasePeCoffGetEntryPointLib/BasePeCoffGetEntryPointLib.inf
   ExtractGuidedSectionLib|MdePkg/Library/PeiExtractGuidedSectionLib/PeiExtractGuidedSectionLib.inf
 
   PeiServicesTablePointerLib|ArmPkg/Library/PeiServicesTablePointerLib/PeiServicesTablePointerLib.inf
@@ -223,7 +222,6 @@
   PeimEntryPoint|MdePkg/Library/PeimEntryPoint/PeimEntryPoint.inf
   PerformanceLib|MdeModulePkg/Library/PeiPerformanceLib/PeiPerformanceLib.inf
   OemHookStatusCodeLib|MdeModulePkg/Library/OemHookStatusCodeLibNull/OemHookStatusCodeLibNull.inf
-  PeCoffGetEntryPointLib|MdePkg/Library/BasePeCoffGetEntryPointLib/BasePeCoffGetEntryPointLib.inf
   ExtractGuidedSectionLib|MdePkg/Library/PeiExtractGuidedSectionLib/PeiExtractGuidedSectionLib.inf
 
   PeiServicesTablePointerLib|ArmPkg/Library/PeiServicesTablePointerLib/PeiServicesTablePointerLib.inf

ArmVirtPkg/PrePi/AArch64/ModuleEntryPoint.S

@@ -149,9 +149,8 @@ ASM_PFX(DiscoverDramFromDt):
   // window at the beginning of the FD image as a temp stack.
   //
   mov   x0, x7
-  adr   x1, PeCoffLoaderImageReadFromMemory
   mov   sp, x7
-  bl    RelocatePeCoffImage
+  bl    RelocateUefiImage
 
   //
   // Discover the memory size and offset from the DTB, and record in the

ArmVirtPkg/PrePi/Arm/ModuleEntryPoint.S

@@ -146,9 +146,8 @@ ASM_PFX(ArmPlatformPeiBootAction):
   // window at the beginning of the FD image as a temp stack.
   //
   mov   r0, r5
-  ADRL  (r1, PeCoffLoaderImageReadFromMemory)
   mov   sp, r5
-  bl    RelocatePeCoffImage
+  bl    RelocateUefiImage
 
   //
   // Discover the memory size and offset from the DTB, and record in the

ArmVirtPkg/PrePi/ArmVirtPrePiUniCoreRelocatable.inf

@@ -46,7 +46,7 @@
   TimerLib
   SerialPortLib
   ExtractGuidedSectionLib
-  PeCoffLib
+  UefiImageLib
   PrePiLib
   MemoryAllocationLib
   HobLib

ArmVirtPkg/PrePi/PrePi.c

@@ -9,7 +9,7 @@
 #include <PiPei.h>
 #include <Pi/PiBootMode.h>
 
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
 #include <Library/PrePiLib.h>
 #include <Library/PrintLib.h>
 #include <Library/PrePiHobListPointerLib.h>
@@ -129,15 +129,15 @@ CEntryPoint (
 }
 
 VOID
-RelocatePeCoffImage (
+RelocateUefiImage (
-  IN  EFI_PEI_FV_HANDLE         FwVolHeader,
+  IN  EFI_PEI_FV_HANDLE             FwVolHeader
-  IN  PE_COFF_LOADER_READ_FILE  ImageRead
   )
 {
-  EFI_PEI_FILE_HANDLE           FileHandle;
+  EFI_PEI_FILE_HANDLE             FileHandle;
-  VOID                          *SectionData;
+  VOID                            *SectionData;
-  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;
+  UINT32                          SectionSize;
-  EFI_STATUS                    Status;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT ImageContext;
+  EFI_STATUS                      Status;
 
   FileHandle = NULL;
   Status     = FfsFindNextFile (
@@ -147,21 +147,16 @@ RelocatePeCoffImage (
                  );
   ASSERT_EFI_ERROR (Status);
 
-  Status = FfsFindSectionData (EFI_SECTION_PE32, FileHandle, &SectionData);
+  Status = FfsFindSectionData (EFI_SECTION_PE32, FileHandle, &SectionData, &SectionSize);
   if (EFI_ERROR (Status)) {
-    Status = FfsFindSectionData (EFI_SECTION_TE, FileHandle, &SectionData);
+    Status = FfsFindSectionData (EFI_SECTION_TE, FileHandle, &SectionData, &SectionSize);
   }
 
   ASSERT_EFI_ERROR (Status);
 
-  ZeroMem (&ImageContext, sizeof ImageContext);
+  Status = UefiImageInitializeContext (&ImageContext, SectionData, SectionSize);
+  ASSERT_RETURN_ERROR (Status);
 
-  ImageContext.Handle    = (EFI_HANDLE)SectionData;
+  Status = UefiImageRelocateImageInplaceForExecution (&ImageContext);
-  ImageContext.ImageRead = ImageRead;
+  ASSERT_RETURN_ERROR (Status);
-  PeCoffLoaderGetImageInfo (&ImageContext);
-
-  if (ImageContext.ImageAddress != (UINTN)SectionData) {
-    ImageContext.ImageAddress = (UINTN)SectionData;
-    PeCoffLoaderRelocateImage (&ImageContext);
-  }
 }

BaseTools/Source/C/GenFv/GenFvInternalLib.c

@@ -54,6 +54,10 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
                                     (((Imm32) >> 2) & 0x7fffff))
 #define ARM_JUMP_TO_THUMB(Offset)  _ARM_JUMP_TO_THUMB((Offset) - 8)
 
+#define ALIGN_VALUE_ADDEND(Value, Alignment)  (((Alignment) - (Value)) & ((Alignment) - 1U))
+#define ALIGN_VALUE(Value, Alignment) ((Value) + ALIGN_VALUE_ADDEND (Value, Alignment))
+#define IS_ALIGNED(Value, Alignment)  (((Value) & ((Alignment) - 1U)) == 0U)
+
 /*
  * Arm instruction to return from exception (MOVS PC, LR)
  */
@@ -1308,11 +1312,17 @@ Returns:
       // Rebase the PE or TE image in FileBuffer of FFS file for XIP
       // Rebase for the debug genfvmap tool
       //
-      Status = FfsRebase (FvInfo, FvInfo->FvFiles[Index], (EFI_FFS_FILE_HEADER *) FileBuffer, (UINTN) *VtfFileImage - (UINTN) FvImage->FileImage, FvMapFile);
+      Status = FfsRebase (FvInfo, FvInfo->FvFiles[Index], (EFI_FFS_FILE_HEADER **)&FileBuffer, &FileSize, (UINTN) *VtfFileImage - (UINTN) FvImage->FileImage, FvMapFile);
       if (EFI_ERROR (Status)) {
         Error (NULL, 0, 3000, "Invalid", "Could not rebase %s.", FvInfo->FvFiles[Index]);
         return Status;
       }
+
+      if (FileSize > (UINTN) ((UINTN) *VtfFileImage - (UINTN) FvImage->CurrentFilePointer)) {
+        free (FileBuffer);
+        Error (NULL, 0, 4002, "Resource", "FV space is full, not enough room to add file %s after ImageBase aligning.", FvInfo->FvFiles[Index]);
+        return EFI_OUT_OF_RESOURCES;
+      }
       //
       // copy VTF File
       //
@@ -1354,11 +1364,17 @@ Returns:
     // Rebase the PE or TE image in FileBuffer of FFS file for XIP.
     // Rebase Bs and Rt drivers for the debug genfvmap tool.
     //
-    Status = FfsRebase (FvInfo, FvInfo->FvFiles[Index], (EFI_FFS_FILE_HEADER *) FileBuffer, (UINTN) FvImage->CurrentFilePointer - (UINTN) FvImage->FileImage, FvMapFile);
+    Status = FfsRebase (FvInfo, FvInfo->FvFiles[Index], (EFI_FFS_FILE_HEADER **)&FileBuffer, &FileSize, (UINTN) FvImage->CurrentFilePointer - (UINTN) FvImage->FileImage, FvMapFile);
-  if (EFI_ERROR (Status)) {
+    if (EFI_ERROR (Status)) {
-    Error (NULL, 0, 3000, "Invalid", "Could not rebase %s.", FvInfo->FvFiles[Index]);
+      Error (NULL, 0, 3000, "Invalid", "Could not rebase %s.", FvInfo->FvFiles[Index]);
-    return Status;
+      return Status;
-  }
+    }
+
+    if (FileSize > (UINTN) ((UINTN) *VtfFileImage - (UINTN) FvImage->CurrentFilePointer)) {
+      free (FileBuffer);
+      Error (NULL, 0, 4002, "Resource", "FV space is full, not enough room to add file %s after ImageBase aligning.", FvInfo->FvFiles[Index]);
+      return EFI_OUT_OF_RESOURCES;
+    }
     //
     // Copy the file
     //
@@ -3456,11 +3472,109 @@ Returns:
   return EFI_SUCCESS;
 }
 
+EFI_PHYSICAL_ADDRESS
+AddPadSection (
+  IN OUT EFI_PHYSICAL_ADDRESS     *BaseAddress,
+  IN     UINT32                   Alignment,
+  IN OUT EFI_FFS_FILE_HEADER      **FfsFile,
+  IN OUT UINTN                    *FileSize,
+  IN OUT EFI_FILE_SECTION_POINTER *Section
+  )
+{
+  EFI_COMMON_SECTION_HEADER *NewSection;
+  UINT32                    FfsHeaderLength;
+  UINT32                    FfsFileLength;
+  UINT32                    PadSize;
+  UINTN                     PadAddress;
+  UINT8                     *FfsPart;
+  UINT32                    PartSize;
+  UINT32                    Offset;
+  EFI_FFS_INTEGRITY_CHECK   *IntegrityCheck;
+
+  PadAddress = ALIGN_VALUE (*BaseAddress + sizeof (EFI_COMMON_SECTION_HEADER), Alignment);
+  PadSize    = PadAddress - *BaseAddress;
+
+  Offset = (UINT32)((UINTN)((*Section).Pe32Section) - (UINTN)(*FfsFile));
+  PartSize = GetFfsFileLength (*FfsFile) - Offset;
+  FfsPart = calloc (1, PartSize);
+  if (FfsPart == NULL) {
+    fprintf (stderr, "GenFv: Could not allocate memory for FfsPart\n");
+    return EFI_OUT_OF_RESOURCES;
+  }
+
+  CopyMem (FfsPart, (UINT8 *)(UINTN)((*Section).Pe32Section), PartSize);
+
+  FfsFileLength = GetFfsFileLength (*FfsFile) + PadSize;
+  *FfsFile = realloc (*FfsFile, FfsFileLength);
+  if (*FfsFile == NULL) {
+    fprintf (stderr, "GenFv: Could not reallocate memory for FfsFile\n");
+    free (FfsPart);
+    return EFI_OUT_OF_RESOURCES;
+  }
+  *FileSize += PadSize;
+
+  NewSection = (EFI_COMMON_SECTION_HEADER *)((UINTN)(*FfsFile) + Offset);
+
+  NewSection->Size[0] = (UINT8)(PadSize & 0xff);
+  NewSection->Size[1] = (UINT8)((PadSize & 0xff00) >> 8);
+  NewSection->Size[2] = (UINT8)((PadSize & 0xff0000) >> 16);
+  NewSection->Type    = EFI_SECTION_RAW;
+  ++NewSection;
+  ZeroMem ((VOID *)NewSection, PadSize - sizeof (EFI_COMMON_SECTION_HEADER));
+
+  *Section = (EFI_FILE_SECTION_POINTER)(EFI_PE32_SECTION *)((UINT8 *)NewSection + PadSize - sizeof (EFI_COMMON_SECTION_HEADER));
+
+  CopyMem (
+    (UINT8 *)((*Section).Pe32Section),
+    FfsPart,
+    PartSize
+    );
+
+  FfsHeaderLength = GetFfsHeaderLength(*FfsFile);
+  if (FfsHeaderLength > sizeof(EFI_FFS_FILE_HEADER)) {
+    ((EFI_FFS_FILE_HEADER2 *)(*FfsFile))->ExtendedSize = FfsFileLength;
+  } else {
+    (*FfsFile)->Size[0] = (UINT8)(FfsFileLength & 0x000000FF);
+    (*FfsFile)->Size[1] = (UINT8)((FfsFileLength & 0x0000FF00) >> 8);
+    (*FfsFile)->Size[2] = (UINT8)((FfsFileLength & 0x00FF0000) >> 16);
+  }
+
+  //
+  // Recalculate the FFS header checksum. Instead of setting Header and State
+  // both to zero, set Header to (UINT8)(-State) so State preserves its original
+  // value
+  //
+  IntegrityCheck = &(*FfsFile)->IntegrityCheck;
+  IntegrityCheck->Checksum.Header = (UINT8) (0x100 - (*FfsFile)->State);
+  IntegrityCheck->Checksum.File = 0;
+
+  IntegrityCheck->Checksum.Header = CalculateChecksum8 (
+                                      (UINT8 *)(*FfsFile), FfsHeaderLength);
+
+  if ((*FfsFile)->Attributes & FFS_ATTRIB_CHECKSUM) {
+    //
+    // Ffs header checksum = zero, so only need to calculate ffs body.
+    //
+    IntegrityCheck->Checksum.File = CalculateChecksum8 (
+                                      (UINT8 *)(*FfsFile) + FfsHeaderLength,
+                                      FfsFileLength - FfsHeaderLength);
+  } else {
+    IntegrityCheck->Checksum.File = FFS_FIXED_CHECKSUM;
+  }
+
+  *BaseAddress = PadAddress;
+
+  free (FfsPart);
+
+  return EFI_SUCCESS;
+}
+
 EFI_STATUS
 FfsRebase (
   IN OUT  FV_INFO               *FvInfo,
   IN      CHAR8                 *FileName,
-  IN OUT  EFI_FFS_FILE_HEADER   *FfsFile,
+  IN OUT  EFI_FFS_FILE_HEADER   **FfsFile,
+  IN OUT  UINTN                 *FileSize,
   IN      UINTN                 XipOffset,
   IN      FILE                  *FvMapFile
   )
@@ -3533,13 +3647,12 @@ Returns:
     return EFI_SUCCESS;
   }
 
-
   XipBase = FvInfo->BaseAddress + XipOffset;
 
   //
   // We only process files potentially containing PE32 sections.
   //
-  switch (FfsFile->Type) {
+  switch ((*FfsFile)->Type) {
     case EFI_FV_FILETYPE_SECURITY_CORE:
     case EFI_FV_FILETYPE_PEI_CORE:
     case EFI_FV_FILETYPE_PEIM:
@@ -3551,7 +3664,7 @@ Returns:
       //
       // Rebase the inside FvImage.
       //
-      GetChildFvFromFfs (FvInfo, FfsFile, XipOffset);
+      GetChildFvFromFfs (FvInfo, *FfsFile, XipOffset);
 
       //
       // Search PE/TE section in FV sectin.
@@ -3561,7 +3674,7 @@ Returns:
       return EFI_SUCCESS;
   }
 
-  FfsHeaderSize = GetFfsHeaderLength(FfsFile);
+  FfsHeaderSize = GetFfsHeaderLength(*FfsFile);
   //
   // Rebase each PE32 section
   //
@@ -3575,7 +3688,7 @@ Returns:
     //
     // Find Pe Image
     //
-    Status = GetSectionByType (FfsFile, EFI_SECTION_PE32, Index, &CurrentPe32Section);
+    Status = GetSectionByType (*FfsFile, EFI_SECTION_PE32, Index, &CurrentPe32Section);
     if (EFI_ERROR (Status)) {
       break;
     }
@@ -3624,7 +3737,7 @@ Returns:
     //
     // Calculate the PE32 base address, based on file type
     //
-    switch (FfsFile->Type) {
+    switch ((*FfsFile)->Type) {
       case EFI_FV_FILETYPE_SECURITY_CORE:
       case EFI_FV_FILETYPE_PEI_CORE:
       case EFI_FV_FILETYPE_PEIM:
@@ -3702,7 +3815,7 @@ Returns:
           ImageContext.RelocationsStripped = FALSE;
         }
 
-        NewPe32BaseAddress = XipBase + (UINTN) CurrentPe32Section.Pe32Section + CurSecHdrSize - (UINTN)FfsFile;
+        NewPe32BaseAddress = XipBase + (UINTN) CurrentPe32Section.Pe32Section + CurSecHdrSize - (UINTN)(*FfsFile);
         break;
 
       case EFI_FV_FILETYPE_DRIVER:
@@ -3717,7 +3830,7 @@ Returns:
           Error (NULL, 0, 3000, "Invalid", "PE image Section-Alignment and File-Alignment do not match : %s.", FileName);
           return EFI_ABORTED;
         }
-        NewPe32BaseAddress = XipBase + (UINTN) CurrentPe32Section.Pe32Section + CurSecHdrSize - (UINTN)FfsFile;
+        NewPe32BaseAddress = XipBase + (UINTN) CurrentPe32Section.Pe32Section + CurSecHdrSize - (UINTN)(*FfsFile);
         break;
 
       default:
@@ -3741,15 +3854,28 @@ Returns:
     //
     // Load and Relocate Image Data
     //
-    MemoryImagePointer = (UINT8 *) malloc ((UINTN) ImageContext.ImageSize + ImageContext.SectionAlignment);
+    MemoryImagePointer = (UINT8 *) calloc (1, (UINTN)ImageContext.ImageSize + ImageContext.SectionAlignment);
     if (MemoryImagePointer == NULL) {
       Error (NULL, 0, 4001, "Resource", "memory cannot be allocated on rebase of %s", FileName);
       return EFI_OUT_OF_RESOURCES;
     }
-    memset ((VOID *) MemoryImagePointer, 0, (UINTN) ImageContext.ImageSize + ImageContext.SectionAlignment);
+    ImageContext.ImageAddress = ALIGN_VALUE ((UINTN)MemoryImagePointer, ImageContext.SectionAlignment);
-    ImageContext.ImageAddress = ((UINTN) MemoryImagePointer + ImageContext.SectionAlignment - 1) & (~((UINTN) ImageContext.SectionAlignment - 1));
 
-    Status =  PeCoffLoaderLoadImage (&ImageContext);
+    if (!(IS_ALIGNED (NewPe32BaseAddress, ImageContext.SectionAlignment))) {
+      Status = AddPadSection (&NewPe32BaseAddress, ImageContext.SectionAlignment, FfsFile, FileSize, &CurrentPe32Section);
+      if (EFI_ERROR (Status)) {
+        free ((VOID *) MemoryImagePointer);
+        return Status;
+      }
+
+      CurSecHdrSize       = GetSectionHeaderLength (CurrentPe32Section.CommonHeader);
+      ImageContext.Handle = (VOID *) ((UINTN) CurrentPe32Section.Pe32Section + CurSecHdrSize);
+      PdbPointer          = PeCoffLoaderGetPdbPointer (ImageContext.Handle);
+
+      ImgHdr = (EFI_IMAGE_OPTIONAL_HEADER_UNION *)((UINTN) CurrentPe32Section.Pe32Section + CurSecHdrSize + ImageContext.PeCoffHeaderOffset);
+    }
+
+    Status = PeCoffLoaderLoadImage (&ImageContext);
     if (EFI_ERROR (Status)) {
       Error (NULL, 0, 3000, "Invalid", "LocateImage() call failed on rebase of %s", FileName);
       free ((VOID *) MemoryImagePointer);
@@ -3757,7 +3883,8 @@ Returns:
     }
 
     ImageContext.DestinationAddress = NewPe32BaseAddress;
-    Status                          = PeCoffLoaderRelocateImage (&ImageContext);
+
+    Status = PeCoffLoaderRelocateImage (&ImageContext);
     if (EFI_ERROR (Status)) {
       Error (NULL, 0, 3000, "Invalid", "RelocateImage() call failed on rebase of %s Status=%d", FileName, Status);
       free ((VOID *) MemoryImagePointer);
@@ -3807,15 +3934,15 @@ Returns:
     //
     // Now update file checksum
     //
-    if (FfsFile->Attributes & FFS_ATTRIB_CHECKSUM) {
+    if ((*FfsFile)->Attributes & FFS_ATTRIB_CHECKSUM) {
-      SavedState  = FfsFile->State;
+      SavedState  = (*FfsFile)->State;
-      FfsFile->IntegrityCheck.Checksum.File = 0;
+      (*FfsFile)->IntegrityCheck.Checksum.File = 0;
-      FfsFile->State                        = 0;
+      (*FfsFile)->State                        = 0;
-      FfsFile->IntegrityCheck.Checksum.File = CalculateChecksum8 (
+      (*FfsFile)->IntegrityCheck.Checksum.File = CalculateChecksum8 (
-                                                (UINT8 *) ((UINT8 *)FfsFile + FfsHeaderSize),
+                                                (UINT8 *) ((UINT8 *)(*FfsFile) + FfsHeaderSize),
-                                                GetFfsFileLength (FfsFile) - FfsHeaderSize
+                                                GetFfsFileLength (*FfsFile) - FfsHeaderSize
                                                 );
-      FfsFile->State = SavedState;
+      (*FfsFile)->State = SavedState;
     }
 
     //
@@ -3829,14 +3956,14 @@ Returns:
       PdbPointer = FileName;
     }
 
-    WriteMapFile (FvMapFile, PdbPointer, FfsFile, NewPe32BaseAddress, &OrigImageContext);
+    WriteMapFile (FvMapFile, PdbPointer, *FfsFile, NewPe32BaseAddress, &OrigImageContext);
   }
 
-  if (FfsFile->Type != EFI_FV_FILETYPE_SECURITY_CORE &&
+  if ((*FfsFile)->Type != EFI_FV_FILETYPE_SECURITY_CORE &&
-      FfsFile->Type != EFI_FV_FILETYPE_PEI_CORE &&
+      (*FfsFile)->Type != EFI_FV_FILETYPE_PEI_CORE &&
-      FfsFile->Type != EFI_FV_FILETYPE_PEIM &&
+      (*FfsFile)->Type != EFI_FV_FILETYPE_PEIM &&
-      FfsFile->Type != EFI_FV_FILETYPE_COMBINED_PEIM_DRIVER &&
+      (*FfsFile)->Type != EFI_FV_FILETYPE_COMBINED_PEIM_DRIVER &&
-      FfsFile->Type != EFI_FV_FILETYPE_FIRMWARE_VOLUME_IMAGE
+      (*FfsFile)->Type != EFI_FV_FILETYPE_FIRMWARE_VOLUME_IMAGE
       ) {
     //
     // Only Peim code may have a TE section
@@ -3853,7 +3980,7 @@ Returns:
     //
     // Find Te Image
     //
-    Status = GetSectionByType (FfsFile, EFI_SECTION_TE, Index, &CurrentPe32Section);
+    Status = GetSectionByType (*FfsFile, EFI_SECTION_TE, Index, &CurrentPe32Section);
     if (EFI_ERROR (Status)) {
       break;
     }
@@ -3901,7 +4028,7 @@ Returns:
     // Set new rebased address.
     //
     NewPe32BaseAddress = XipBase + (UINTN) TEImageHeader + sizeof (EFI_TE_IMAGE_HEADER) \
-                         - TEImageHeader->StrippedSize - (UINTN) FfsFile;
+                         - TEImageHeader->StrippedSize - (UINTN) (*FfsFile);
 
     //
     // if reloc is stripped, try to get the original efi image to get reloc info.
@@ -4044,15 +4171,15 @@ Returns:
     //
     // Now update file checksum
     //
-    if (FfsFile->Attributes & FFS_ATTRIB_CHECKSUM) {
+    if ((*FfsFile)->Attributes & FFS_ATTRIB_CHECKSUM) {
-      SavedState  = FfsFile->State;
+      SavedState  = (*FfsFile)->State;
-      FfsFile->IntegrityCheck.Checksum.File = 0;
+      (*FfsFile)->IntegrityCheck.Checksum.File = 0;
-      FfsFile->State                        = 0;
+      (*FfsFile)->State                        = 0;
-      FfsFile->IntegrityCheck.Checksum.File = CalculateChecksum8 (
+      (*FfsFile)->IntegrityCheck.Checksum.File = CalculateChecksum8 (
-                                                (UINT8 *)((UINT8 *)FfsFile + FfsHeaderSize),
+                                                (UINT8 *)((UINT8 *)(*FfsFile) + FfsHeaderSize),
-                                                GetFfsFileLength (FfsFile) - FfsHeaderSize
+                                                GetFfsFileLength (*FfsFile) - FfsHeaderSize
                                                 );
-      FfsFile->State = SavedState;
+      (*FfsFile)->State = SavedState;
     }
     //
     // Get this module function address from ModulePeMapFile and add them into FvMap file
@@ -4068,7 +4195,7 @@ Returns:
     WriteMapFile (
       FvMapFile,
       PdbPointer,
-      FfsFile,
+      *FfsFile,
       NewPe32BaseAddress,
       &OrigImageContext
       );

BaseTools/Source/C/GenFv/GenFvInternalLib.h

@@ -331,7 +331,8 @@ EFI_STATUS
 FfsRebase (
   IN OUT  FV_INFO               *FvInfo,
   IN      CHAR8                 *FileName,
-  IN OUT  EFI_FFS_FILE_HEADER   *FfsFile,
+  IN OUT  EFI_FFS_FILE_HEADER   **FfsFile,
+  IN OUT  UINTN                 *FileSize,
   IN      UINTN                 XipOffset,
   IN      FILE                  *FvMapFile
   );

BaseTools/Source/Python/build/build.py

@@ -72,6 +72,24 @@ gSupportedTarget = ['all', 'genc', 'genmake', 'modules', 'libraries', 'fds', 'cl
 TemporaryTablePattern = re.compile(r'^_\d+_\d+_[a-fA-F0-9]+$')
 TmpTableDict = {}
 
+## Return the biggest multiple of alignment that is smaller than or equal to
+#  value.
+#
+#   @param value      The value to align down.
+#   @param alignment  The boundary to align down to.
+#
+def AlignDown(value, alignment):
+    return value - (value % alignment)
+
+## Return the smallest multiple of alignment that is bigger than or equal to
+#  value.
+#
+#   @param value      The value to align up.
+#   @param alignment  The boundary to align up to.
+#
+def AlignUp(value, alignment):
+    return AlignDown(value + alignment - 1, alignment)
+
 ## Check environment PATH variable to make sure the specified tool is found
 #
 #   If the tool is found in the PATH, then True is returned
@@ -706,7 +724,7 @@ class PeImageInfo():
         self.OutputDir        = OutputDir
         self.DebugDir         = DebugDir
         self.Image            = ImageClass
-        self.Image.Size       = (self.Image.Size // 0x1000 + 1) * 0x1000
+        self.Image.Size       = AlignUp(self.Image.Size, 0x1000)
 
 ## The class implementing the EDK2 build process
 #
@@ -1520,12 +1538,15 @@ class Build():
             ## for SMM module in SMRAM, the SMRAM will be allocated from base to top.
             if not ModeIsSmm:
                 BaseAddress = BaseAddress - ModuleInfo.Image.Size
+                BaseAddress = AlignDown(BaseAddress, ModuleInfo.Image.SectionAlignment)
                 #
                 # Update Image to new BaseAddress by GenFw tool
                 #
                 LaunchCommand(["GenFw", "--rebase", str(BaseAddress), "-r", ModuleOutputImage], ModuleInfo.OutputDir)
                 LaunchCommand(["GenFw", "--rebase", str(BaseAddress), "-r", ModuleDebugImage], ModuleInfo.DebugDir)
+            ## for SMM module in SMRAM, the SMRAM will be allocated from base to top.
             else:
+                BaseAddress = AlignUp(BaseAddress, ModuleInfo.Image.SectionAlignment)
                 #
                 # Set new address to the section header only for SMM driver.
                 #

EmbeddedPkg/EmbeddedPkg.dec

@@ -81,7 +81,6 @@
   gEmbeddedDeviceGuid =   { 0xbf4b9d10, 0x13ec, 0x43dd, { 0x88, 0x80, 0xe9, 0xb, 0x71, 0x8f, 0x27, 0xde } }
   gEmbeddedExternalDeviceProtocolGuid = { 0x735F8C64, 0xD696, 0x44D0, { 0xBD, 0xF2, 0x44, 0x7F, 0xD0, 0x5A, 0x54, 0x06 }}
   gEmbeddedGpioProtocolGuid           = { 0x17a0a3d7, 0xc0a5, 0x4635, { 0xbb, 0xd5, 0x07, 0x21, 0x87, 0xdf, 0xe2, 0xee }}
-  gPeCoffLoaderProtocolGuid =  { 0xB323179B, 0x97FB, 0x477E, { 0xB0, 0xFE, 0xD8, 0x85, 0x91, 0xFA, 0x11, 0xAB } }
   gEmbeddedMmcHostProtocolGuid         = { 0x3e591c00, 0x9e4a, 0x11df, {0x92, 0x44, 0x00, 0x02, 0xA5, 0xD5, 0xC5, 0x1B }}
   gAndroidFastbootTransportProtocolGuid = { 0x74bd9fe0, 0x8902, 0x11e3, {0xb9, 0xd3, 0xf7, 0x22, 0x38, 0xfc, 0x9a, 0x31}}
   gAndroidFastbootPlatformProtocolGuid =  { 0x524685a0, 0x89a0, 0x11e3, {0x9d, 0x4d, 0xbf, 0xa9, 0xf6, 0xa4, 0x03, 0x08}}

EmbeddedPkg/EmbeddedPkg.dsc

@@ -58,9 +58,9 @@
 
   ReportStatusCodeLib|MdeModulePkg/Library/PeiDxeDebugLibReportStatusCode/PeiDxeDebugLibReportStatusCode.inf
 
-  PeCoffGetEntryPointLib|MdePkg/Library/BasePeCoffGetEntryPointLib/BasePeCoffGetEntryPointLib.inf
+  PeCoffLib2|MdePkg/Library/BasePeCoffLib2/BasePeCoffLib2.inf
-  PeCoffLib|MdePkg/Library/BasePeCoffLib/BasePeCoffLib.inf
+  UefiImageLib|MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibPeCoff.inf
-  PeCoffExtraActionLib|MdePkg/Library/BasePeCoffExtraActionLibNull/BasePeCoffExtraActionLibNull.inf
+  UefiImageExtraActionLib|MdePkg/Library/BaseUefiImageExtraActionLibNull/BaseUefiImageExtraActionLibNull.inf
   CacheMaintenanceLib|MdePkg/Library/BaseCacheMaintenanceLib/BaseCacheMaintenanceLib.inf
   PrePiLib|EmbeddedPkg/Library/PrePiLib/PrePiLib.inf
 

EmbeddedPkg/GdbStub/GdbStub.c

@@ -807,166 +807,6 @@ gXferObjectReadResponse (
   return Count;
 }
 
-/**
-  Note: This should be a library function.  In the Apple case you have to add
-  the size of the PE/COFF header into the starting address to make things work
-  right as there is no way to pad the Mach-O for the size of the PE/COFF header.
-
-
-  Returns a pointer to the PDB file name for a PE/COFF image that has been
-  loaded into system memory with the PE/COFF Loader Library functions.
-
-  Returns the PDB file name for the PE/COFF image specified by Pe32Data.  If
-  the PE/COFF image specified by Pe32Data is not a valid, then NULL is
-  returned.  If the PE/COFF image specified by Pe32Data does not contain a
-  debug directory entry, then NULL is returned.  If the debug directory entry
-  in the PE/COFF image specified by Pe32Data does not contain a PDB file name,
-  then NULL is returned.
-  If Pe32Data is NULL, then ASSERT().
-
-  @param  Pe32Data   Pointer to the PE/COFF image that is loaded in system
-                     memory.
-  @param  DebugBase  Address that the debugger would use as the base of the image
-
-  @return The PDB file name for the PE/COFF image specified by Pe32Data or NULL
-          if it cannot be retrieved. DebugBase is only valid if PDB file name is
-          valid.
-
-**/
-VOID *
-EFIAPI
-PeCoffLoaderGetDebuggerInfo (
-  IN VOID   *Pe32Data,
-  OUT VOID  **DebugBase
-  )
-{
-  EFI_IMAGE_DOS_HEADER                 *DosHdr;
-  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
-  EFI_IMAGE_DATA_DIRECTORY             *DirectoryEntry;
-  EFI_IMAGE_DEBUG_DIRECTORY_ENTRY      *DebugEntry;
-  UINTN                                DirCount;
-  VOID                                 *CodeViewEntryPointer;
-  INTN                                 TEImageAdjust;
-  UINT32                               NumberOfRvaAndSizes;
-  UINT16                               Magic;
-  UINTN                                SizeOfHeaders;
-
-  ASSERT (Pe32Data   != NULL);
-
-  TEImageAdjust       = 0;
-  DirectoryEntry      = NULL;
-  DebugEntry          = NULL;
-  NumberOfRvaAndSizes = 0;
-  SizeOfHeaders       = 0;
-
-  DosHdr = (EFI_IMAGE_DOS_HEADER *)Pe32Data;
-  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
-    //
-    // DOS image header is present, so read the PE header after the DOS image header.
-    //
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINTN)Pe32Data + (UINTN)((DosHdr->e_lfanew) & 0x0ffff));
-  } else {
-    //
-    // DOS image header is not present, so PE header is at the image base.
-    //
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)Pe32Data;
-  }
-
-  if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {
-    if (Hdr.Te->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress != 0) {
-      DirectoryEntry = &Hdr.Te->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG];
-      TEImageAdjust  = sizeof (EFI_TE_IMAGE_HEADER) - Hdr.Te->StrippedSize;
-      DebugEntry     = (EFI_IMAGE_DEBUG_DIRECTORY_ENTRY *)((UINTN)Hdr.Te +
-                                                           Hdr.Te->DataDirectory[EFI_TE_IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress +
-                                                           TEImageAdjust);
-    }
-
-    SizeOfHeaders = sizeof (EFI_TE_IMAGE_HEADER) + (UINTN)Hdr.Te->BaseOfCode - (UINTN)Hdr.Te->StrippedSize;
-
-    // __APPLE__ check this math...
-    *DebugBase = ((CHAR8 *)Pe32Data) -  TEImageAdjust;
-  } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {
-    *DebugBase = Pe32Data;
-
-    //
-    // NOTE: We use Machine field to identify PE32/PE32+, instead of Magic.
-    //       It is due to backward-compatibility, for some system might
-    //       generate PE32+ image with PE32 Magic.
-    //
-    switch (Hdr.Pe32->FileHeader.Machine) {
-      case EFI_IMAGE_MACHINE_IA32:
-        //
-        // Assume PE32 image with IA32 Machine field.
-        //
-        Magic = EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC;
-        break;
-      case EFI_IMAGE_MACHINE_X64:
-      case EFI_IMAGE_MACHINE_IA64:
-        //
-        // Assume PE32+ image with X64 or IPF Machine field
-        //
-        Magic = EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;
-        break;
-      default:
-        //
-        // For unknown Machine field, use Magic in optional Header
-        //
-        Magic = Hdr.Pe32->OptionalHeader.Magic;
-    }
-
-    if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
-      //
-      // Use PE32 offset get Debug Directory Entry
-      //
-      SizeOfHeaders       = Hdr.Pe32->OptionalHeader.SizeOfHeaders;
-      NumberOfRvaAndSizes = Hdr.Pe32->OptionalHeader.NumberOfRvaAndSizes;
-      DirectoryEntry      = (EFI_IMAGE_DATA_DIRECTORY *)&(Hdr.Pe32->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_DEBUG]);
-      DebugEntry          = (EFI_IMAGE_DEBUG_DIRECTORY_ENTRY *)((UINTN)Pe32Data + DirectoryEntry->VirtualAddress);
-    } else if (Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
-      //
-      // Use PE32+ offset get Debug Directory Entry
-      //
-      SizeOfHeaders       = Hdr.Pe32Plus->OptionalHeader.SizeOfHeaders;
-      NumberOfRvaAndSizes = Hdr.Pe32Plus->OptionalHeader.NumberOfRvaAndSizes;
-      DirectoryEntry      = (EFI_IMAGE_DATA_DIRECTORY *)&(Hdr.Pe32Plus->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_DEBUG]);
-      DebugEntry          = (EFI_IMAGE_DEBUG_DIRECTORY_ENTRY *)((UINTN)Pe32Data + DirectoryEntry->VirtualAddress);
-    }
-
-    if (NumberOfRvaAndSizes <= EFI_IMAGE_DIRECTORY_ENTRY_DEBUG) {
-      DirectoryEntry = NULL;
-      DebugEntry     = NULL;
-    }
-  } else {
-    return NULL;
-  }
-
-  if ((DebugEntry == NULL) || (DirectoryEntry == NULL)) {
-    return NULL;
-  }
-
-  for (DirCount = 0; DirCount < DirectoryEntry->Size; DirCount += sizeof (EFI_IMAGE_DEBUG_DIRECTORY_ENTRY), DebugEntry++) {
-    if (DebugEntry->Type == EFI_IMAGE_DEBUG_TYPE_CODEVIEW) {
-      if (DebugEntry->SizeOfData > 0) {
-        CodeViewEntryPointer = (VOID *)((UINTN)DebugEntry->RVA + ((UINTN)Pe32Data) + (UINTN)TEImageAdjust);
-        switch (*(UINT32 *)CodeViewEntryPointer) {
-          case CODEVIEW_SIGNATURE_NB10:
-            return (VOID *)((CHAR8 *)CodeViewEntryPointer + sizeof (EFI_IMAGE_DEBUG_CODEVIEW_NB10_ENTRY));
-          case CODEVIEW_SIGNATURE_RSDS:
-            return (VOID *)((CHAR8 *)CodeViewEntryPointer + sizeof (EFI_IMAGE_DEBUG_CODEVIEW_RSDS_ENTRY));
-          case CODEVIEW_SIGNATURE_MTOC:
-            *DebugBase = (VOID *)(UINTN)((UINTN)DebugBase - SizeOfHeaders);
-            return (VOID *)((CHAR8 *)CodeViewEntryPointer + sizeof (EFI_IMAGE_DEBUG_CODEVIEW_MTOC_ENTRY));
-          default:
-            break;
-        }
-      }
-    }
-  }
-
-  (void)SizeOfHeaders;
-  return NULL;
-}
-
 /**
   Process "qXfer:object:read:annex:offset,length" request.
 
@@ -1007,7 +847,7 @@ QxferLibrary (
   )
 {
   VOID   *LoadAddress;
-  CHAR8  *Pdb;
+  CONST CHAR8                       *Pdb;
   UINTN  Size;
 
   if (Offset != gPacketqXferLibraryOffset) {
@@ -1035,12 +875,8 @@ QxferLibrary (
     for ( ; gEfiDebugImageTableEntry < gDebugImageTableHeader->TableSize; gEfiDebugImageTableEntry++, gDebugTable++) {
       if (gDebugTable->NormalImage != NULL) {
         if ((gDebugTable->NormalImage->ImageInfoType == EFI_DEBUG_IMAGE_INFO_TYPE_NORMAL) &&
-            (gDebugTable->NormalImage->LoadedImageProtocolInstance != NULL))
+            (gDebugTable->NormalImage->LoadedImageProtocolInstance != NULL)) {
-        {
+          Pdb = gDebugTable->NormalImage->PdbPath;
-          Pdb = PeCoffLoaderGetDebuggerInfo (
-                  gDebugTable->NormalImage->LoadedImageProtocolInstance->ImageBase,
-                  &LoadAddress
-                  );
           if (Pdb != NULL) {
             Size = AsciiSPrint (
                      gXferLibraryBuffer,

EmbeddedPkg/GdbStub/GdbStubInternal.h

@@ -26,7 +26,7 @@
 #include <Protocol/LoadedImage.h>
 #include <Protocol/LoadedImage.h>
 #include <Guid/DebugImageInfoTable.h>
-#include <IndustryStandard/PeImage.h>
+#include <IndustryStandard/PeImage2.h>
 
 extern CONST CHAR8  mHexToStr[];
 

EmbeddedPkg/Include/Library/PrePiLib.h

@@ -71,6 +71,7 @@ EFI_STATUS
   @param  FileHeader            A pointer to the file header that contains the set of sections to
                                 be searched.
   @param  SectionData           A pointer to the discovered section, if successful.
+  @param  SectionSize           A pointer to the size of the discovered section, if successful.
 
   @retval EFI_SUCCESS           The section was found.
   @retval EFI_NOT_FOUND         The section was not found.
@@ -82,7 +83,8 @@ FfsFindSectionDataWithHook (
   IN EFI_SECTION_TYPE        SectionType,
   IN FFS_CHECK_SECTION_HOOK  SectionCheckHook,
   IN EFI_PEI_FILE_HANDLE     FileHandle,
-  OUT VOID                   **SectionData
+  OUT VOID                   **SectionData,
+  OUT UINT32                 *SectionSize
   );
 
 /**
@@ -92,6 +94,7 @@ FfsFindSectionDataWithHook (
   @param  FileHandle            A pointer to the file header that contains the set of sections to
                                 be searched.
   @param  SectionData           A pointer to the discovered section, if successful.
+  @param  SectionSize           A pointer to the size of the discovered section, if successful.
 
   @retval EFI_SUCCESS           The section was found.
   @retval EFI_NOT_FOUND         The section was not found.
@@ -102,7 +105,8 @@ EFIAPI
 FfsFindSectionData (
   IN EFI_SECTION_TYPE     SectionType,
   IN EFI_PEI_FILE_HANDLE  FileHandle,
-  OUT VOID                **SectionData
+  OUT VOID                **SectionData,
+  OUT UINT32              *SectionSize
   );
 
 /**
@@ -675,7 +679,7 @@ BuildExtractSectionHob (
 
 VOID
 EFIAPI
-BuildPeCoffLoaderHob (
+BuildUefiLoaderHob (
   VOID
   );
 
@@ -760,10 +764,11 @@ AllocateAlignedPages (
 
 EFI_STATUS
 EFIAPI
-LoadPeCoffImage (
+LoadUefiImage (
-  IN  VOID                  *PeCoffImage,
+  IN  VOID                                      *UefiImage,
+  IN  UINT32                                    UefiImageSize,
   OUT EFI_PHYSICAL_ADDRESS  *ImageAddress,
-  OUT UINT64                *ImageSize,
+  OUT UINT32                *ImageSize,
   OUT EFI_PHYSICAL_ADDRESS  *EntryPoint
   );
 

EmbeddedPkg/Library/PrePiHobLib/Hob.c

@@ -12,12 +12,11 @@
 #include <Library/BaseLib.h>
 #include <Library/BaseMemoryLib.h>
 #include <Library/DebugLib.h>
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
 #include <Library/HobLib.h>
 #include <Library/PcdLib.h>
 #include <Library/PrePiHobListPointerLib.h>
 
-#include <Protocol/PeCoffLoader.h>
 #include <Guid/ExtractSection.h>
 #include <Guid/MemoryTypeInformation.h>
 #include <Guid/MemoryAllocationHob.h>
@@ -821,27 +820,6 @@ BuildExtractSectionHob (
   BuildGuidDataHob (Guid, &Data, sizeof (Data));
 }
 
-PE_COFF_LOADER_PROTOCOL  gPeCoffProtocol = {
-  PeCoffLoaderGetImageInfo,
-  PeCoffLoaderLoadImage,
-  PeCoffLoaderRelocateImage,
-  PeCoffLoaderImageReadFromMemory,
-  PeCoffLoaderRelocateImageForRuntime,
-  PeCoffLoaderUnloadImage
-};
-
-VOID
-EFIAPI
-BuildPeCoffLoaderHob (
-  VOID
-  )
-{
-  VOID  *Ptr;
-
-  Ptr = &gPeCoffProtocol;
-  BuildGuidDataHob (&gPeCoffLoaderProtocolGuid, &Ptr, sizeof (VOID *));
-}
-
 // May want to put this into a library so you only need the PCD settings if you are using the feature?
 VOID
 BuildMemoryTypeInformationHob (

EmbeddedPkg/Library/PrePiLib/FwVol.c

@@ -278,7 +278,8 @@ FfsProcessSection (
   IN FFS_CHECK_SECTION_HOOK     SectionCheckHook,
   IN EFI_COMMON_SECTION_HEADER  *Section,
   IN UINTN                      SectionSize,
-  OUT VOID                      **OutputBuffer
+  OUT VOID                      **OutputBuffer,
+  OUT UINT32                    *OutputSize
   )
 {
   EFI_STATUS                Status;
@@ -298,6 +299,7 @@ FfsProcessSection (
 
   Found         = FALSE;
   *OutputBuffer = NULL;
+  *OutputSize   = 0;
   ParsedLength  = 0;
   Status        = EFI_NOT_FOUND;
   while (ParsedLength < SectionSize) {
@@ -313,10 +315,13 @@ FfsProcessSection (
       }
 
       if (Found) {
+        // FIXME: Use common API with size checks
         if (IS_SECTION2 (Section)) {
           *OutputBuffer = (VOID *)((UINT8 *)Section + sizeof (EFI_COMMON_SECTION_HEADER2));
+          *OutputSize   = SECTION2_SIZE (Section) - sizeof (EFI_COMMON_SECTION_HEADER2);
         } else {
           *OutputBuffer = (VOID *)((UINT8 *)Section + sizeof (EFI_COMMON_SECTION_HEADER));
+          *OutputSize   = SECTION_SIZE (Section) - sizeof (EFI_COMMON_SECTION_HEADER);
         }
 
         return EFI_SUCCESS;
@@ -432,7 +437,8 @@ FfsProcessSection (
                  SectionCheckHook,
                  DstBuffer,
                  DstBufferSize,
-                 OutputBuffer
+                 OutputBuffer,
+                 OutputSize
                  );
       }
     }
@@ -466,6 +472,7 @@ CheckNextSection:
   @param  FileHandle            A pointer to the file header that contains the set of sections to
                                 be searched.
   @param  SectionData           A pointer to the discovered section, if successful.
+  @param  SectionSize           A pointer to the size of the discovered section, if successful.
 
   @retval EFI_SUCCESS           The section was found.
   @retval EFI_NOT_FOUND         The section was not found.
@@ -477,7 +484,8 @@ FfsFindSectionDataWithHook (
   IN EFI_SECTION_TYPE        SectionType,
   IN FFS_CHECK_SECTION_HOOK  SectionCheckHook,
   IN EFI_PEI_FILE_HANDLE     FileHandle,
-  OUT VOID                   **SectionData
+  OUT VOID                   **SectionData,
+  OUT UINT32                 *SectionSize
   )
 {
   EFI_FFS_FILE_HEADER        *FfsFileHeader;
@@ -500,7 +508,8 @@ FfsFindSectionDataWithHook (
            SectionCheckHook,
            Section,
            FileSize,
-           SectionData
+           SectionData,
+           SectionSize
            );
 }
 
@@ -511,6 +520,7 @@ FfsFindSectionDataWithHook (
   @param  FileHandle            A pointer to the file header that contains the set of sections to
                                 be searched.
   @param  SectionData           A pointer to the discovered section, if successful.
+  @param  SectionSize           A pointer to the size of the discovered section, if successful.
 
   @retval EFI_SUCCESS           The section was found.
   @retval EFI_NOT_FOUND         The section was not found.
@@ -521,10 +531,11 @@ EFIAPI
 FfsFindSectionData (
   IN EFI_SECTION_TYPE     SectionType,
   IN EFI_PEI_FILE_HANDLE  FileHandle,
-  OUT VOID                **SectionData
+  OUT VOID                **SectionData,
+  OUT UINT32              *SectionSize
   )
 {
-  return FfsFindSectionDataWithHook (SectionType, NULL, FileHandle, SectionData);
+  return FfsFindSectionDataWithHook (SectionType, NULL, FileHandle, SectionData, SectionSize);
 }
 
 /**
@@ -816,6 +827,7 @@ FfsProcessFvFile (
 {
   EFI_STATUS            Status;
   EFI_PEI_FV_HANDLE     FvImageHandle;
+  UINT32                FvImageHandleSize;
   EFI_FV_INFO           FvImageInfo;
   UINT32                FvAlignment;
   VOID                  *FvBuffer;
@@ -842,7 +854,13 @@ FfsProcessFvFile (
   //
   // Find FvImage in FvFile
   //
-  Status = FfsFindSectionDataWithHook (EFI_SECTION_FIRMWARE_VOLUME_IMAGE, NULL, FvFileHandle, (VOID **)&FvImageHandle);
+  Status = FfsFindSectionDataWithHook (
+    EFI_SECTION_FIRMWARE_VOLUME_IMAGE,
+    NULL,
+    FvFileHandle,
+    (VOID **)&FvImageHandle,
+    &FvImageHandleSize
+    );
   if (EFI_ERROR (Status)) {
     return Status;
   }

EmbeddedPkg/Library/PrePiLib/PrePi.h

@@ -18,7 +18,7 @@
 #include <Library/DebugLib.h>
 #include <Library/BaseMemoryLib.h>
 #include <Library/UefiDecompressLib.h>
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
 #include <Library/CacheMaintenanceLib.h>
 #include <Library/MemoryAllocationLib.h>
 #include <Library/TimerLib.h>

EmbeddedPkg/Library/PrePiLib/PrePiLib.c

@@ -6,6 +6,7 @@
 
 **/
 
+#include "ProcessorBind.h"
 #include <PrePi.h>
 
 //
@@ -56,54 +57,45 @@ AllocateCodePages (
 
 EFI_STATUS
 EFIAPI
-LoadPeCoffImage (
+LoadUefiImage (
-  IN  VOID                  *PeCoffImage,
+  IN  VOID                  *UefiImage,
+  IN  UINT32                UefiImageSize,
   OUT EFI_PHYSICAL_ADDRESS  *ImageAddress,
-  OUT UINT64                *ImageSize,
+  OUT UINT32                *DestinationSize,
   OUT EFI_PHYSICAL_ADDRESS  *EntryPoint
   )
 {
-  RETURN_STATUS                 Status;
+  RETURN_STATUS                   Status;
-  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT ImageContext;
-  VOID                          *Buffer;
+  UINT32                          ImageSize;
+  VOID                            *Buffer;
+  UINT32                          BufferSize;
+  UINT32                          BufferPages;
+  UINT32                          BufferAlignment;
 
-  ZeroMem (&ImageContext, sizeof (ImageContext));
+  Status = UefiImageInitializeContext (&ImageContext, UefiImage, UefiImageSize);
-
-  ImageContext.Handle    = PeCoffImage;
-  ImageContext.ImageRead = PeCoffLoaderImageReadFromMemory;
-
-  Status = PeCoffLoaderGetImageInfo (&ImageContext);
   ASSERT_EFI_ERROR (Status);
 
+  ImageSize       = UefiImageGetImageSize (&ImageContext);
+  BufferPages     = EFI_SIZE_TO_PAGES (ImageSize);
+  BufferSize      = EFI_PAGES_TO_SIZE (BufferPages);
+  BufferAlignment = UefiImageGetSegmentAlignment (&ImageContext);
+
   //
   // Allocate Memory for the image
   //
-  Buffer = AllocateCodePages (EFI_SIZE_TO_PAGES ((UINT32)ImageContext.ImageSize));
+  Buffer = AllocateAlignedCodePages (BufferPages, BufferAlignment);
   ASSERT (Buffer != 0);
 
-  ImageContext.ImageAddress = (EFI_PHYSICAL_ADDRESS)(UINTN)Buffer;
-
   //
-  // Load the image to our new buffer
+  // Load and relocate the image to our new buffer
   //
-  Status = PeCoffLoaderLoadImage (&ImageContext);
+  Status = UefiImageLoadImageForExecution (&ImageContext, Buffer, BufferSize, NULL, 0);
   ASSERT_EFI_ERROR (Status);
 
-  //
+  *ImageAddress    = (UINTN) Buffer;
-  // Relocate the image in our new buffer
+  *DestinationSize = BufferSize;
-  //
+  *EntryPoint      = (UINTN) UefiImageLoaderGetImageEntryPoint (&ImageContext);
-  Status = PeCoffLoaderRelocateImage (&ImageContext);
-  ASSERT_EFI_ERROR (Status);
-
-  *ImageAddress = ImageContext.ImageAddress;
-  *ImageSize    = ImageContext.ImageSize;
-  *EntryPoint   = ImageContext.EntryPoint;
-
-  //
-  // Flush not needed for all architectures. We could have a processor specific
-  // function in this library that does the no-op if needed.
-  //
-  InvalidateInstructionCacheRange ((VOID *)(UINTN)*ImageAddress, (UINTN)*ImageSize);
 
   return Status;
 }
@@ -122,22 +114,23 @@ LoadDxeCoreFromFfsFile (
   )
 {
   EFI_STATUS            Status;
-  VOID                  *PeCoffImage;
+  VOID                  *UefiImage;
+  UINT32                UefiImageSize;
   EFI_PHYSICAL_ADDRESS  ImageAddress;
-  UINT64                ImageSize;
+  UINT32                DestinationSize;
   EFI_PHYSICAL_ADDRESS  EntryPoint;
   VOID                  *BaseOfStack;
   VOID                  *TopOfStack;
   VOID                  *Hob;
   EFI_FV_FILE_INFO      FvFileInfo;
 
-  Status = FfsFindSectionDataWithHook (EFI_SECTION_PE32, NULL, FileHandle, &PeCoffImage);
+  Status = FfsFindSectionDataWithHook (EFI_SECTION_PE32, NULL, FileHandle, &UefiImage, &UefiImageSize);
   if (EFI_ERROR (Status)) {
     return Status;
   }
 
-  Status = LoadPeCoffImage (PeCoffImage, &ImageAddress, &ImageSize, &EntryPoint);
+  Status = LoadUefiImage (UefiImage, UefiImageSize, &ImageAddress, &DestinationSize, &EntryPoint);
-  // For NT32 Debug  Status = SecWinNtPeiLoadFile (PeCoffImage, &ImageAddress, &ImageSize, &EntryPoint);
+  // For NT32 Debug  Status = SecWinNtPeiLoadFile (UefiImage, &ImageAddress, &ImageSize, &EntryPoint);
   ASSERT_EFI_ERROR (Status);
 
   //
@@ -146,7 +139,7 @@ LoadDxeCoreFromFfsFile (
   Status = FfsGetFileInfo (FileHandle, &FvFileInfo);
   ASSERT_EFI_ERROR (Status);
 
-  BuildModuleHob (&FvFileInfo.FileName, (EFI_PHYSICAL_ADDRESS)(UINTN)ImageAddress, EFI_SIZE_TO_PAGES ((UINT32)ImageSize) * EFI_PAGE_SIZE, EntryPoint);
+  BuildModuleHob (&FvFileInfo.FileName, (EFI_PHYSICAL_ADDRESS)(UINTN)ImageAddress, DestinationSize, EntryPoint);
 
   DEBUG ((DEBUG_INFO | DEBUG_LOAD, "Loading DxeCore at 0x%10p EntryPoint=0x%10p\n", (VOID *)(UINTN)ImageAddress, (VOID *)(UINTN)EntryPoint));
 

EmbeddedPkg/Library/PrePiLib/PrePiLib.inf

@@ -41,7 +41,7 @@
   DebugLib
   BaseMemoryLib
   UefiDecompressLib
-  PeCoffLib
+  UefiImageLib
   CacheMaintenanceLib
   PrintLib
   SerialPortLib
@@ -53,9 +53,6 @@
 [Guids]
   gEfiMemoryTypeInformationGuid
 
-[Protocols]
-  gPeCoffLoaderProtocolGuid
-
 
 [FixedPcd.common]
   gEmbeddedTokenSpaceGuid.PcdMemoryTypeEfiACPIReclaimMemory

EmulatorPkg/EmulatorPkg.dsc

@@ -62,8 +62,8 @@
   SynchronizationLib|MdePkg/Library/BaseSynchronizationLib/BaseSynchronizationLib.inf
   PrintLib|MdePkg/Library/BasePrintLib/BasePrintLib.inf
   CacheMaintenanceLib|MdePkg/Library/BaseCacheMaintenanceLib/BaseCacheMaintenanceLib.inf
-  PeCoffLib|MdePkg/Library/BasePeCoffLib/BasePeCoffLib.inf
+  PeCoffLib2|MdePkg/Library/BasePeCoffLib2/BasePeCoffLib2.inf
-  PeCoffGetEntryPointLib|MdePkg/Library/BasePeCoffGetEntryPointLib/BasePeCoffGetEntryPointLib.inf
+  UefiImageLib|MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibPeCoff.inf
   BaseMemoryLib|MdePkg/Library/BaseMemoryLib/BaseMemoryLib.inf
   FrameBufferBltLib|MdeModulePkg/Library/FrameBufferBltLib/FrameBufferBltLib.inf
 
@@ -145,8 +145,7 @@
 [LibraryClasses.common.SEC]
   PeiServicesLib|EmulatorPkg/Library/SecPeiServicesLib/SecPeiServicesLib.inf
   PcdLib|MdePkg/Library/BasePcdLibNull/BasePcdLibNull.inf
-  PeCoffGetEntryPointLib|EmulatorPkg/Library/PeiEmuPeCoffGetEntryPointLib/PeiEmuPeCoffGetEntryPointLib.inf
+  UefiImageExtraActionLib|EmulatorPkg/Library/PeiEmuUefiImageExtraActionLib/PeiEmuUefiImageExtraActionLib.inf
-  PeCoffExtraActionLib|EmulatorPkg/Library/PeiEmuPeCoffExtraActionLib/PeiEmuPeCoffExtraActionLib.inf
   SerialPortLib|EmulatorPkg/Library/PeiEmuSerialPortLib/PeiEmuSerialPortLib.inf
   PpiListLib|EmulatorPkg/Library/SecPpiListLib/SecPpiListLib.inf
   DebugLib|MdePkg/Library/BaseDebugLibSerialPort/BaseDebugLibSerialPort.inf
@@ -154,13 +153,13 @@
 
 [LibraryClasses.common.USER_DEFINED, LibraryClasses.common.BASE]
   DebugLib|MdePkg/Library/BaseDebugLibNull/BaseDebugLibNull.inf
-  PeCoffExtraActionLib|MdePkg/Library/BasePeCoffExtraActionLibNull/BasePeCoffExtraActionLibNull.inf
+  UefiImageExtraActionLib|MdePkg/Library/BaseUefiImageExtraActionLibNull/BaseUefiImageExtraActionLibNull.inf
+  HobLib|MdePkg/Library/PeiHobLib/PeiHobLib.inf
   MemoryAllocationLib|MdePkg/Library/PeiMemoryAllocationLib/PeiMemoryAllocationLib.inf
   PcdLib|MdePkg/Library/BasePcdLibNull/BasePcdLibNull.inf
   PpiListLib|EmulatorPkg/Library/SecPpiListLib/SecPpiListLib.inf
   ThunkPpiList|EmulatorPkg/Library/ThunkPpiList/ThunkPpiList.inf
   ThunkProtocolList|EmulatorPkg/Library/ThunkProtocolList/ThunkProtocolList.inf
-  PeCoffGetEntryPointLib|MdePkg/Library/BasePeCoffGetEntryPointLib/BasePeCoffGetEntryPointLib.inf
   PpiListLib|EmulatorPkg/Library/SecPpiListLib/SecPpiListLib.inf
   PeiServicesLib|EmulatorPkg/Library/SecPeiServicesLib/SecPeiServicesLib.inf
 
@@ -169,8 +168,7 @@
   HobLib|MdePkg/Library/PeiHobLib/PeiHobLib.inf
   MemoryAllocationLib|MdePkg/Library/PeiMemoryAllocationLib/PeiMemoryAllocationLib.inf
   ReportStatusCodeLib|MdeModulePkg/Library/PeiReportStatusCodeLib/PeiReportStatusCodeLib.inf
-  PeCoffGetEntryPointLib|EmulatorPkg/Library/PeiEmuPeCoffGetEntryPointLib/PeiEmuPeCoffGetEntryPointLib.inf
+  UefiImageExtraActionLib|EmulatorPkg/Library/PeiEmuUefiImageExtraActionLib/PeiEmuUefiImageExtraActionLib.inf
-  PeCoffExtraActionLib|EmulatorPkg/Library/PeiEmuPeCoffExtraActionLib/PeiEmuPeCoffExtraActionLib.inf
   ExtractGuidedSectionLib|MdePkg/Library/PeiExtractGuidedSectionLib/PeiExtractGuidedSectionLib.inf
   SerialPortLib|EmulatorPkg/Library/PeiEmuSerialPortLib/PeiEmuSerialPortLib.inf
   ReportStatusCodeLib|MdeModulePkg/Library/PeiReportStatusCodeLib/PeiReportStatusCodeLib.inf
@@ -186,7 +184,7 @@
   HobLib|MdePkg/Library/DxeCoreHobLib/DxeCoreHobLib.inf
   MemoryAllocationLib|MdeModulePkg/Library/DxeCoreMemoryAllocationLib/DxeCoreMemoryAllocationLib.inf
   ReportStatusCodeLib|MdeModulePkg/Library/DxeReportStatusCodeLib/DxeReportStatusCodeLib.inf
-  PeCoffExtraActionLib|EmulatorPkg/Library/DxeEmuPeCoffExtraActionLib/DxeEmuPeCoffExtraActionLib.inf
+  UefiImageExtraActionLib|EmulatorPkg/Library/DxeEmuUefiImageExtraActionLib/DxeEmuUefiImageExtraActionLib.inf
   ExtractGuidedSectionLib|MdePkg/Library/DxeExtractGuidedSectionLib/DxeExtractGuidedSectionLib.inf
   PcdLib|MdePkg/Library/BasePcdLibNull/BasePcdLibNull.inf
   TimerLib|EmulatorPkg/Library/DxeCoreTimerLib/DxeCoreTimerLib.inf
@@ -208,7 +206,7 @@
   MemoryAllocationLib|MdePkg/Library/UefiMemoryAllocationLib/UefiMemoryAllocationLib.inf
   ReportStatusCodeLib|MdeModulePkg/Library/DxeReportStatusCodeLib/DxeReportStatusCodeLib.inf
   EmuThunkLib|EmulatorPkg/Library/DxeEmuLib/DxeEmuLib.inf
-  PeCoffExtraActionLib|EmulatorPkg/Library/DxeEmuPeCoffExtraActionLib/DxeEmuPeCoffExtraActionLib.inf
+  UefiImageExtraActionLib|EmulatorPkg/Library/DxeEmuUefiImageExtraActionLib/DxeEmuUefiImageExtraActionLib.inf
   ReportStatusCodeLib|MdeModulePkg/Library/DxeReportStatusCodeLib/DxeReportStatusCodeLib.inf
   TimerLib|EmulatorPkg/Library/DxeTimerLib/DxeTimerLib.inf
 

EmulatorPkg/Include/Library/EmuMagicPageLib.h

@@ -23,7 +23,7 @@ typedef struct {
   // Used by SecPeiServicesLib
   EFI_PEI_PPI_DESCRIPTOR    *PpiList;
 
-  // Needed by PEI PEI PeCoffLoaderExtraActionLib
+  // Needed by PEI PEI UefiImageLoaderExtraActionLib
   EMU_THUNK_PROTOCOL        *Thunk;
 } EMU_MAGIC_PAGE_LAYOUT;
 

EmulatorPkg/Include/Protocol/EmuThunk.h

@@ -17,7 +17,7 @@
 // neded for things like EFI_TIME_CAPABILITIES
 #include <Uefi.h>
 
-#include <Library/PeCoffExtraActionLib.h>
+#include <Library/UefiImageExtraActionLib.h>
 
 #include <Protocol/EmuIoThunk.h>
 #include <Protocol/DevicePath.h>
@@ -88,20 +88,21 @@ BOOLEAN
 typedef
 EFI_STATUS
 (EFIAPI *EMU_PE_COFF_GET_ENTRY_POINT)(
-  IN     VOID  *Pe32Data,
+  IN     VOID   *Pe32Data,
-  IN OUT VOID  **EntryPoint
+  IN     UINT32 Pe32Size,
+  IN OUT VOID   **EntryPoint
   );
 
 typedef
 VOID
 (EFIAPI *EMU_PE_COFF_RELOCATE_EXTRA_ACTION)(
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   );
 
 typedef
 VOID
 (EFIAPI *EMU_PE_COFF_UNLOAD_EXTRA_ACTION)(
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   );
 
 typedef
@@ -217,9 +218,9 @@ struct _EMU_THUNK_PROTOCOL {
   ///
   /// PE/COFF loader hooks to get symbols loaded
   ///
-  EMU_PE_COFF_GET_ENTRY_POINT          PeCoffGetEntryPoint;
+  EMU_PE_COFF_GET_ENTRY_POINT       UefiImageGetEntryPoint;
-  EMU_PE_COFF_RELOCATE_EXTRA_ACTION    PeCoffRelocateImageExtraAction;
+  EMU_PE_COFF_RELOCATE_EXTRA_ACTION UefiImageRelocateImageExtraAction;
-  EMU_PE_COFF_UNLOAD_EXTRA_ACTION      PeCoffUnloadImageExtraAction;
+  EMU_PE_COFF_UNLOAD_EXTRA_ACTION   UefiImageUnloadImageExtraAction;
 
   ///
   /// DXE Architecture Protocol Services

EmulatorPkg/Library/DxeEmuPeCoffExtraActionLib/DxeEmuPeCoffExtraActionLib.c

@@ -1,3 +1,5 @@
+#ifndef DISABLE_NEW_DEPRECATED_INTERFACES
+
 /** @file
   Provides services to perform additional actions to relocate and unload
   PE/Coff image for Emu environment specific purpose such as souce level debug.
@@ -23,7 +25,8 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 //
 // Cache of UnixThunk protocol
 //
-EMU_THUNK_PROTOCOL  *mThunk = NULL;
+EMU_THUNK_PROTOCOL   *mThunk = NULL;
+
 
 /**
   The constructor function gets  the pointer of the WinNT thunk functions
@@ -39,7 +42,7 @@ DxeEmuPeCoffLibExtraActionConstructor (
   IN EFI_SYSTEM_TABLE  *SystemTable
   )
 {
-  EFI_HOB_GUID_TYPE  *GuidHob;
+   EFI_HOB_GUID_TYPE        *GuidHob;
 
   //
   // Retrieve EmuThunkProtocol from GUID'ed HOB
@@ -72,6 +75,8 @@ PeCoffLoaderRelocateImageExtraAction (
   }
 }
 
+
+
 /**
   Performs additional actions just before a PE/COFF image is unloaded.  Any resources
   that were allocated by PeCoffLoaderRelocateImageExtraAction() must be freed.
@@ -92,3 +97,5 @@ PeCoffLoaderUnloadImageExtraAction (
     mThunk->PeCoffUnloadImageExtraAction (ImageContext);
   }
 }
+
+#endif // DISABLE_NEW_DEPRECATED_INTERFACES

EmulatorPkg/Library/DxeEmuUefiImageExtraActionLib/DxeEmuUefiImageExtraActionLib.c

@@ -0,0 +1,94 @@
+/** @file
+  Provides services to perform additional actions to relocate and unload
+  PE/Coff image for Emu environment specific purpose such as souce level debug.
+  This version only works for DXE phase
+
+Copyright (c) 2006 - 2009, Intel Corporation. All rights reserved.<BR>
+Portions copyright (c) 2008 - 2011, Apple Inc. All rights reserved.<BR>
+SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include <PiDxe.h>
+
+#include <Protocol/EmuThunk.h>
+
+#include <Library/UefiImageLib.h>
+#include <Library/BaseLib.h>
+#include <Library/DebugLib.h>
+#include <Library/HobLib.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/UefiImageExtraActionLib.h>
+
+//
+// Cache of UnixThunk protocol
+//
+EMU_THUNK_PROTOCOL  *mThunk = NULL;
+
+/**
+  The constructor function gets  the pointer of the WinNT thunk functions
+  It will ASSERT() if Unix thunk protocol is not installed.
+
+  @retval EFI_SUCCESS   Unix thunk protocol is found and cached.
+
+**/
+EFI_STATUS
+EFIAPI
+DxeEmuUefiImageLibExtraActionConstructor (
+  IN EFI_HANDLE        ImageHandle,
+  IN EFI_SYSTEM_TABLE  *SystemTable
+  )
+{
+  EFI_HOB_GUID_TYPE  *GuidHob;
+
+  //
+  // Retrieve EmuThunkProtocol from GUID'ed HOB
+  //
+  GuidHob = GetFirstGuidHob (&gEmuThunkProtocolGuid);
+  ASSERT (GuidHob != NULL);
+  mThunk = (EMU_THUNK_PROTOCOL *)(*(UINTN *)(GET_GUID_HOB_DATA (GuidHob)));
+  ASSERT (mThunk != NULL);
+
+  return EFI_SUCCESS;
+}
+
+/**
+  Performs additional actions after a PE/COFF image has been loaded and relocated.
+
+  If ImageContext is NULL, then ASSERT().
+
+  @param  ImageContext  Pointer to the image context structure that describes the
+                        PE/COFF image that has already been loaded and relocated.
+
+**/
+VOID
+EFIAPI
+UefiImageLoaderRelocateImageExtraAction (
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
+  )
+{
+  if (mThunk != NULL) {
+    mThunk->UefiImageRelocateImageExtraAction (ImageContext);
+  }
+}
+
+/**
+  Performs additional actions just before a PE/COFF image is unloaded.  Any resources
+  that were allocated by UefiImageLoaderRelocateImageExtraAction() must be freed.
+
+  If ImageContext is NULL, then ASSERT().
+
+  @param  ImageContext  Pointer to the image context structure that describes the
+                        PE/COFF image that is being unloaded.
+
+**/
+VOID
+EFIAPI
+UefiImageLoaderUnloadImageExtraAction (
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
+  )
+{
+  if (mThunk != NULL) {
+    mThunk->UefiImageUnloadImageExtraAction (ImageContext);
+  }
+}

EmulatorPkg/Library/DxeEmuUefiImageExtraActionLib/DxeEmuUefiImageExtraActionLib.inf

@@ -0,0 +1,42 @@
+## @file
+# PeCoff extra action libary for DXE phase that run Emu emulator.
+#
+# Copyright (c) 2007 - 2010, Intel Corporation. All rights reserved.<BR>
+# Portions copyright (c) 2011, Apple Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause-Patent
+
+#
+#
+##
+
+[Defines]
+  INF_VERSION                    = 0x00010005
+  BASE_NAME                      = DxeEmuUefiImageExtraActionLib
+  FILE_GUID                      = 68FCD487-D230-6846-95B1-5E1F2EF942C4
+  MODULE_TYPE                    = DXE_DRIVER
+  VERSION_STRING                 = 1.0
+  LIBRARY_CLASS                  = UefiImageExtraActionLib|DXE_CORE DXE_DRIVER DXE_RUNTIME_DRIVER UEFI_DRIVER UEFI_APPLICATION
+
+  CONSTRUCTOR                    = DxeEmuUefiImageLibExtraActionConstructor
+
+#
+# The following information is for reference only and not required by the build tools.
+#
+#  VALID_ARCHITECTURES           = IA32
+#
+
+[Sources]
+  DxeEmuUefiImageExtraActionLib.c
+
+[Packages]
+  MdePkg/MdePkg.dec
+  EmulatorPkg/EmulatorPkg.dec
+
+[LibraryClasses]
+  DebugLib
+  HobLib
+  BaseMemoryLib
+
+[Protocols]
+  gEmuThunkProtocolGuid                     # PROTOCOL ALWAYS_CONSUMED
+

EmulatorPkg/Library/PeiEmuPeCoffExtraActionLib/PeiEmuPeCoffExtraActionLib.c

@@ -1,3 +1,5 @@
+#ifndef DISABLE_NEW_DEPRECATED_INTERFACES
+
 /** @file
   Provides services to perform additional actions to relocate and unload
   PE/Coff image for Emu environment specific purpose such as souce level debug.
@@ -22,7 +24,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 //
 // Cache of UnixThunk protocol
 //
-EMU_THUNK_PROTOCOL  *mThunk = NULL;
+EMU_THUNK_PROTOCOL   *mThunk = NULL;
 
 /**
   The function caches the pointer of the Unix thunk functions
@@ -36,21 +38,22 @@ EFIAPI
 EmuPeCoffGetThunkStucture (
   )
 {
-  EMU_THUNK_PPI  *ThunkPpi;
+  EMU_THUNK_PPI     *ThunkPpi;
-  EFI_STATUS     Status;
+  EFI_STATUS        Status;
+
 
   //
   // Locate Unix ThunkPpi for retrieving standard output handle
   //
   Status = PeiServicesLocatePpi (
-             &gEmuThunkPpiGuid,
+              &gEmuThunkPpiGuid,
-             0,
+              0,
-             NULL,
+              NULL,
-             (VOID **)&ThunkPpi
+              (VOID **) &ThunkPpi
-             );
+              );
   ASSERT_EFI_ERROR (Status);
 
-  EMU_MAGIC_PAGE ()->Thunk = (EMU_THUNK_PROTOCOL *)ThunkPpi->Thunk ();
+  EMU_MAGIC_PAGE()->Thunk = (EMU_THUNK_PROTOCOL *) ThunkPpi->Thunk ();
 
   return EFI_SUCCESS;
 }
@@ -70,12 +73,12 @@ PeCoffLoaderRelocateImageExtraAction (
   IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
   )
 {
-  if (EMU_MAGIC_PAGE ()->Thunk == NULL) {
+  if (EMU_MAGIC_PAGE()->Thunk == NULL) {
     EmuPeCoffGetThunkStucture ();
   }
+    EMU_MAGIC_PAGE()->Thunk->PeCoffRelocateImageExtraAction (ImageContext);
+  }
 
-  EMU_MAGIC_PAGE ()->Thunk->PeCoffRelocateImageExtraAction (ImageContext);
-}
 
 /**
   Performs additional actions just before a PE/COFF image is unloaded.  Any resources
@@ -93,9 +96,10 @@ PeCoffLoaderUnloadImageExtraAction (
   IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
   )
 {
-  if (EMU_MAGIC_PAGE ()->Thunk == NULL) {
+  if (EMU_MAGIC_PAGE()->Thunk == NULL) {
     EmuPeCoffGetThunkStucture ();
   }
-
+  EMU_MAGIC_PAGE()->Thunk->PeCoffUnloadImageExtraAction (ImageContext);
-  EMU_MAGIC_PAGE ()->Thunk->PeCoffUnloadImageExtraAction (ImageContext);
 }
+
+#endif // DISABLE_NEW_DEPRECATED_INTERFACES

EmulatorPkg/Library/PeiEmuPeCoffGetEntryPointLib/PeiEmuPeCoffGetEntryPointLib.c

@@ -7,7 +7,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 **/
 
 #include "PiPei.h"
-#include <Library/PeCoffGetEntryPointLib.h>
+//#include <Library/PeCoffGetEntryPointLib.h>
 #include <Library/PeiServicesLib.h>
 #include <IndustryStandard/PeImage.h>
 #include <Library/DebugLib.h>
@@ -35,8 +35,9 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 RETURN_STATUS
 EFIAPI
 PeCoffLoaderGetEntryPoint (
-  IN     VOID  *Pe32Data,
+  IN     VOID   *Pe32Data,
-  IN OUT VOID  **EntryPoint
+  IN     UINT32 Pe32Size,
+  IN OUT VOID   **EntryPoint
   )
 {
   EMU_THUNK_PPI       *ThunkPpi;
@@ -56,7 +57,7 @@ PeCoffLoaderGetEntryPoint (
 
   Thunk = (EMU_THUNK_PROTOCOL *)ThunkPpi->Thunk ();
 
-  return Thunk->PeCoffGetEntryPoint (Pe32Data, EntryPoint);
+  return Thunk->PeCoffGetEntryPoint (Pe32Data, Pe32Size, EntryPoint);
 }
 
 /**

EmulatorPkg/Library/PeiEmuUefiImageExtraActionLib/PeiEmuUefiImageExtraActionLib.c

@@ -0,0 +1,101 @@
+/** @file
+  Provides services to perform additional actions to relocate and unload
+  PE/Coff image for Emu environment specific purpose such as souce level debug.
+  This version only works for PEI phase
+
+Copyright (c) 2006 - 2009, Intel Corporation. All rights reserved.<BR>
+Portions copyright (c) 2008 - 2011, Apple Inc. All rights reserved.<BR>
+SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+#include <PiPei.h>
+#include <Ppi/EmuThunk.h>
+#include <Protocol/EmuThunk.h>
+
+#include <Library/UefiImageLib.h>
+#include <Library/PeiServicesLib.h>
+#include <Library/DebugLib.h>
+#include <Library/BaseLib.h>
+#include <Library/UefiImageExtraActionLib.h>
+#include <Library/EmuMagicPageLib.h>
+
+//
+// Cache of UnixThunk protocol
+//
+EMU_THUNK_PROTOCOL  *mThunk = NULL;
+
+/**
+  The function caches the pointer of the Unix thunk functions
+  It will ASSERT() if Unix thunk ppi is not installed.
+
+  @retval EFI_SUCCESS   WinNT thunk protocol is found and cached.
+
+**/
+EFI_STATUS
+EFIAPI
+EmuUefiImageGetThunkStucture (
+  )
+{
+  EMU_THUNK_PPI  *ThunkPpi;
+  EFI_STATUS     Status;
+
+  //
+  // Locate Unix ThunkPpi for retrieving standard output handle
+  //
+  Status = PeiServicesLocatePpi (
+             &gEmuThunkPpiGuid,
+             0,
+             NULL,
+             (VOID **)&ThunkPpi
+             );
+  ASSERT_EFI_ERROR (Status);
+
+  EMU_MAGIC_PAGE ()->Thunk = (EMU_THUNK_PROTOCOL *)ThunkPpi->Thunk ();
+
+  return EFI_SUCCESS;
+}
+
+/**
+  Performs additional actions after a PE/COFF image has been loaded and relocated.
+
+  If ImageContext is NULL, then ASSERT().
+
+  @param  ImageContext  Pointer to the image context structure that describes the
+                        PE/COFF image that has already been loaded and relocated.
+
+**/
+VOID
+EFIAPI
+UefiImageLoaderRelocateImageExtraAction (
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
+  )
+{
+  if (EMU_MAGIC_PAGE ()->Thunk == NULL) {
+    EmuUefiImageGetThunkStucture ();
+  }
+
+  EMU_MAGIC_PAGE ()->Thunk->UefiImageRelocateImageExtraAction (ImageContext);
+}
+
+/**
+  Performs additional actions just before a PE/COFF image is unloaded.  Any resources
+  that were allocated by UefiImageLoaderRelocateImageExtraAction() must be freed.
+
+  If ImageContext is NULL, then ASSERT().
+
+  @param  ImageContext  Pointer to the image context structure that describes the
+                        PE/COFF image that is being unloaded.
+
+**/
+VOID
+EFIAPI
+UefiImageLoaderUnloadImageExtraAction (
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
+  )
+{
+  if (EMU_MAGIC_PAGE ()->Thunk == NULL) {
+    EmuUefiImageGetThunkStucture ();
+  }
+
+  EMU_MAGIC_PAGE ()->Thunk->UefiImageUnloadImageExtraAction (ImageContext);
+}

EmulatorPkg/Library/PeiEmuUefiImageExtraActionLib/PeiEmuUefiImageExtraActionLib.inf

@@ -0,0 +1,43 @@
+## @file
+# PeCoff extra action libary for Pei phase that run Emu emulator.
+#
+# Lib to provide memory journal status code reporting Routines
+# Copyright (c) 2007 - 2010, Intel Corporation. All rights reserved.<BR>
+# Portions copyright (c) 2011, Apple Inc. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause-Patent
+
+#
+#
+##
+
+[Defines]
+  INF_VERSION                    = 0x00010005
+  BASE_NAME                      = PeiEmuUefiImageExtraActionLib
+  FILE_GUID                      = 79C4E72A-730B-F040-8129-95877B3A97A8
+  MODULE_TYPE                    = PEIM
+  VERSION_STRING                 = 1.0
+  LIBRARY_CLASS                  = UefiImageExtraActionLib|PEI_CORE PEIM
+
+#
+# The following information is for reference only and not required by the build tools.
+#
+#  VALID_ARCHITECTURES           = IA32
+#
+
+[Sources]
+  PeiEmuUefiImageExtraActionLib.c
+
+[Packages]
+  MdePkg/MdePkg.dec
+  EmulatorPkg/EmulatorPkg.dec
+
+[LibraryClasses]
+  BaseLib
+  PeiServicesLib
+  DebugLib
+
+[Ppis]
+  gEmuThunkPpiGuid                          # PPI ALWAYS_CONSUMED
+
+[Pcd]
+  gEmulatorPkgTokenSpaceGuid.PcdPeiServicesTablePage

EmulatorPkg/Library/SecPeiServicesLib/FwVol.c

@@ -276,3 +276,52 @@ Returns:
 
   return EFI_NOT_FOUND;
 }
+
+// FIXME: Docs
+EFI_STATUS
+SecFfsFindSectionData2 (
+  IN EFI_SECTION_TYPE      SectionType,
+  IN EFI_FFS_FILE_HEADER   *FfsFileHeader,
+  IN OUT VOID              **SectionData,
+  OUT UINT32               *SectionDataSize
+  )
+{
+  UINT32                    FileSize;
+  EFI_COMMON_SECTION_HEADER *Section;
+  UINT32                    SectionLength;
+  UINT32                    ParsedLength;
+
+  //
+  // Size is 24 bits wide so mask upper 8 bits.
+  //    Does not include FfsFileHeader header size
+  // FileSize is adjusted to FileOccupiedSize as it is 8 byte aligned.
+  //
+  Section   = (EFI_COMMON_SECTION_HEADER *) (FfsFileHeader + 1);
+  FileSize  = *(UINT32 *) (FfsFileHeader->Size) & 0x00FFFFFF;
+  FileSize -= sizeof (EFI_FFS_FILE_HEADER);
+
+  *SectionData  = NULL;
+  ParsedLength  = 0;
+  while (ParsedLength < FileSize) {
+    // FIXME: Common API with size checks
+    //
+    // Size is 24 bits wide so mask upper 8 bits.
+    //
+    SectionLength = *(UINT32 *) Section->Size & 0x00FFFFFF;
+    if (Section->Type == SectionType) {
+      *SectionData = (VOID *) (Section + 1);
+      *SectionDataSize = SectionLength - sizeof (*Section);
+      return EFI_SUCCESS;
+    }
+    //
+    // SectionLength is adjusted it is 4 byte aligned.
+    // Go to the next section
+    //
+    SectionLength = GET_OCCUPIED_SIZE (SectionLength, 4);
+
+    ParsedLength += SectionLength;
+    Section = (EFI_COMMON_SECTION_HEADER *) ((UINT8 *) Section + SectionLength);
+  }
+
+  return EFI_NOT_FOUND;
+}

EmulatorPkg/Library/SecPeiServicesLib/PeiServicesLib.c

@@ -26,6 +26,14 @@ SecFfsFindSectionData (
   OUT VOID                **SectionData
   );
 
+EFI_STATUS
+SecFfsFindSectionData2 (
+  IN EFI_SECTION_TYPE      SectionType,
+  IN EFI_FFS_FILE_HEADER   *FfsFileHeader,
+  IN OUT VOID              **SectionData,
+  OUT UINT32               *SectionDataSize
+  );
+
 /**
   This service enables a given PEIM to register an interface into the PEI Foundation.
 
@@ -304,6 +312,38 @@ PeiServicesFfsFindSectionData (
   return SecFfsFindSectionData (SectionType, FileHandle, SectionData);
 }
 
+/**
+  This service enables PEIMs to discover sections of a given instance and type within a valid FFS file.
+
+  @param  SectionType           The value of the section type to find.
+  @param  SectionInstance       Section instance to find.
+  @param  FileHandle            A pointer to the file header that contains the set
+                                of sections to be searched.
+  @param  SectionData           A pointer to the discovered section, if successful.
+  @param  SectionDataSize       The size of the discovered section, if successful.
+  @param  AuthenticationStatus  A pointer to the authentication status for this section.
+
+  @retval EFI_SUCCESS           The section was found.
+  @retval EFI_NOT_FOUND         The section was not found.
+
+**/
+EFI_STATUS
+EFIAPI
+PeiServicesFfsFindSectionData4 (
+  IN EFI_SECTION_TYPE           SectionType,
+  IN UINTN                      SectionInstance,
+  IN EFI_PEI_FILE_HANDLE        FileHandle,
+  OUT VOID                      **SectionData,
+  OUT UINT32                    *SectionDataSize,
+  OUT UINT32                    *AuthenticationStatus
+  )
+{
+  ASSERT (SectionInstance == 0);
+
+  *AuthenticationStatus = 0;
+  return SecFfsFindSectionData2 (SectionType, FileHandle, SectionData, SectionDataSize);
+}
+
 /**
   This service enables PEIMs to register the permanent memory configuration
   that has been initialized with the PEI Foundation.

EmulatorPkg/Sec/Sec.c

@@ -9,6 +9,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 **/
 
 #include "Sec.h"
+#include <Ppi/EmuThunk.h>
 
 EFI_PEI_TEMPORARY_RAM_SUPPORT_PPI  mSecTemporaryRamSupportPpi = {
   SecTemporaryRamSupport
@@ -22,6 +23,53 @@ EFI_PEI_PPI_DESCRIPTOR  gPrivateDispatchTable[] = {
   }
 };
 
+//FIXME:
+/**
+  Retrieves and returns a pointer to the entry point to a PE/COFF image that has been loaded
+  into system memory with the PE/COFF Loader Library functions.
+
+  Retrieves the entry point to the PE/COFF image specified by Pe32Data and returns this entry
+  point in EntryPoint.  If the entry point could not be retrieved from the PE/COFF image, then
+  return RETURN_INVALID_PARAMETER.  Otherwise return RETURN_SUCCESS.
+  If Pe32Data is NULL, then ASSERT().
+  If EntryPoint is NULL, then ASSERT().
+
+  @param  Pe32Data                  The pointer to the PE/COFF image that is loaded in system memory.
+  @param  EntryPoint                The pointer to entry point to the PE/COFF image to return.
+
+  @retval RETURN_SUCCESS            EntryPoint was returned.
+  @retval RETURN_INVALID_PARAMETER  The entry point could not be found in the PE/COFF image.
+
+**/
+RETURN_STATUS
+EFIAPI
+UefiImageLoaderGetEntryPoint (
+  IN     VOID   *Pe32Data,
+  IN     UINT32 Pe32Size,
+  IN OUT VOID   **EntryPoint
+  )
+{
+  EMU_THUNK_PPI           *ThunkPpi;
+  EFI_STATUS              Status;
+  EMU_THUNK_PROTOCOL      *Thunk;
+
+  //
+  // Locate EmuThunkPpi for retrieving standard output handle
+  //
+  Status = PeiServicesLocatePpi (
+              &gEmuThunkPpiGuid,
+              0,
+              NULL,
+              (VOID **) &ThunkPpi
+             );
+  ASSERT_EFI_ERROR (Status);
+
+  Thunk  = (EMU_THUNK_PROTOCOL *)ThunkPpi->Thunk ();
+
+  return Thunk->UefiImageGetEntryPoint (Pe32Data, Pe32Size, EntryPoint);
+}
+
+
 /**
   The entry point of PE/COFF Image for the PEI Core, that has been hijacked by this
   SEC that sits on top of an OS application. So the entry and exit of this module
@@ -64,13 +112,15 @@ _ModuleEntryPoint (
   EFI_STATUS                Status;
   EFI_PEI_FV_HANDLE         VolumeHandle;
   EFI_PEI_FILE_HANDLE       FileHandle;
-  VOID                      *PeCoffImage;
+  VOID                      *UefiImage;
+  UINT32                    UefiImageSize;
   EFI_PEI_CORE_ENTRY_POINT  EntryPoint;
   EFI_PEI_PPI_DESCRIPTOR    *Ppi;
   EFI_PEI_PPI_DESCRIPTOR    *SecPpiList;
   UINTN                     SecReseveredMemorySize;
   UINTN                     Index;
   EFI_PEI_PPI_DESCRIPTOR    PpiArray[10];
+  UINT32                    AuthenticationStatus;
 
   EMU_MAGIC_PAGE ()->PpiList = PpiList;
   ProcessLibraryConstructorList ();
@@ -118,10 +168,10 @@ _ModuleEntryPoint (
   Status       = PeiServicesFfsFindNextFile (EFI_FV_FILETYPE_PEI_CORE, VolumeHandle, &FileHandle);
   ASSERT_EFI_ERROR (Status);
 
-  Status = PeiServicesFfsFindSectionData (EFI_SECTION_PE32, FileHandle, &PeCoffImage);
+  Status = PeiServicesFfsFindSectionData4 (EFI_SECTION_PE32, 0, FileHandle, &UefiImage, &UefiImageSize, &AuthenticationStatus);
   ASSERT_EFI_ERROR (Status);
 
-  Status = PeCoffLoaderGetEntryPoint (PeCoffImage, (VOID **)&EntryPoint);
+  Status = UefiImageLoaderGetEntryPoint (UefiImage, UefiImageSize, (VOID **)&EntryPoint);
   ASSERT_EFI_ERROR (Status);
 
   // Transfer control to PEI Core

EmulatorPkg/Sec/Sec.h

@@ -15,7 +15,6 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/EmuMagicPageLib.h>
 #include <Library/DebugLib.h>
 #include <Library/PeiServicesLib.h>
-#include <Library/PeCoffGetEntryPointLib.h>
 #include <Library/BaseMemoryLib.h>
 
 #include <Ppi/TemporaryRamSupport.h>

EmulatorPkg/Sec/Sec.inf

@@ -36,7 +36,6 @@
 
 [LibraryClasses]
   DebugLib
-  PeCoffGetEntryPointLib
   PeiServicesLib
   PpiListLib
   BaseMemoryLib

EmulatorPkg/Unix/GdbRun.sh

@@ -41,7 +41,7 @@ set $SymbolFileChangesCount = 0
 #
 define AddFirmwareSymbolFile
   if $SymbolFileChangesCount < $arg0
-    add-symbol-file $arg1 $arg2
+    add-symbol-file $arg1 -o $arg2
     set $SymbolFileChangesCount = $arg0
   end
 end

EmulatorPkg/Unix/Host/EmuThunk.c

@@ -415,9 +415,9 @@ EMU_THUNK_PROTOCOL  gEmuThunkProtocol = {
   GasketSecMalloc,
   GasketSecValloc,
   GasketSecFree,
-  GasketSecPeCoffGetEntryPoint,
+  GasketSecUefiImageGetEntryPoint,
-  GasketSecPeCoffRelocateImageExtraAction,
+  GasketSecUefiImageRelocateImageExtraAction,
-  GasketSecPeCoffUnloadImageExtraAction,
+  GasketSecUefiImageUnloadImageExtraAction,
   GasketSecEnableInterrupt,
   GasketSecDisableInterrupt,
   GasketQueryPerformanceFrequency,

EmulatorPkg/Unix/Host/Gasket.h

@@ -67,21 +67,22 @@ GasketSecFree (
 
 RETURN_STATUS
 EFIAPI
-GasketSecPeCoffGetEntryPoint (
+GasketSecUefiImageGetEntryPoint (
   IN     VOID  *Pe32Data,
+  IN     UINT32 Pe32Size,
   IN OUT VOID  **EntryPoint
   );
 
 VOID
 EFIAPI
-GasketSecPeCoffRelocateImageExtraAction (
+GasketSecUefiImageRelocateImageExtraAction (
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   );
 
 VOID
 EFIAPI
-GasketSecPeCoffUnloadImageExtraAction (
+GasketSecUefiImageUnloadImageExtraAction (
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   );
 
 VOID

EmulatorPkg/Unix/Host/Host.c

@@ -44,9 +44,6 @@ EMU_FD_INFO  *gFdInfo;
 UINTN              gSystemMemoryCount = 0;
 EMU_SYSTEM_MEMORY  *gSystemMemory;
 
-UINTN                        mImageContextModHandleArraySize = 0;
-IMAGE_CONTEXT_TO_MOD_HANDLE  *mImageContextModHandleArray    = NULL;
-
 EFI_PEI_PPI_DESCRIPTOR  *gPpiList;
 
 int  gInXcode = 0;
@@ -98,11 +95,13 @@ main (
   BOOLEAN               Done;
   EFI_PEI_FILE_HANDLE   FileHandle;
   VOID                  *SecFile;
+  UINT32                SecFileSize;
   CHAR16                *MemorySizeStr;
   CHAR16                *FirmwareVolumesStr;
   UINTN                 *StackPointer;
   FILE                  *GdbTempFile;
   EMU_THUNK_PPI         *SecEmuThunkPpi;
+  UINT32                AuthenticationStatus;
 
   //
   // Xcode does not support sourcing gdb scripts directly, so the Xcode XML
@@ -111,8 +110,7 @@ main (
   SecGdbConfigBreak ();
 
   //
-  // If dlopen doesn't work, then we build a gdb script to allow the
+  // We build a gdb script to allow the symbols to be loaded.
-  // symbols to be loaded.
   //
   Index               = strlen (*Argv);
   gGdbWorkingFileName = AllocatePool (Index + strlen (".gdb") + 1);
@@ -284,7 +282,14 @@ main (
                      &FileHandle
                      );
       if (!EFI_ERROR (Status)) {
-        Status = PeiServicesFfsFindSectionData (EFI_SECTION_PE32, FileHandle, &SecFile);
+        Status = PeiServicesFfsFindSectionData4 (
+                   EFI_SECTION_PE32,
+                   0,
+                   FileHandle,
+                   &SecFile,
+                   &SecFileSize,
+                   &AuthenticationStatus
+                   );
         if (!EFI_ERROR (Status)) {
           PeiIndex = Index;
           printf (" contains SEC Core");
@@ -330,7 +335,7 @@ main (
   //
   // Hand off to SEC
   //
-  SecLoadFromCore ((UINTN)InitialStackMemory, (UINTN)InitialStackMemorySize, (UINTN)gFdInfo[0].Address, SecFile);
+  SecLoadFromCore ((UINTN)InitialStackMemory, (UINTN)InitialStackMemorySize, (UINTN)gFdInfo[0].Address, SecFile, SecFileSize);
 
   //
   // If we get here, then the SEC Core returned. This is an error as SEC should
@@ -546,7 +551,8 @@ SecLoadFromCore (
   IN  UINTN  LargestRegion,
   IN  UINTN  LargestRegionSize,
   IN  UINTN  BootFirmwareVolumeBase,
-  IN  VOID   *PeiCorePe32File
+  IN  VOID    *PeiCorePe32File,
+  IN  UINT32  PeiCorePe32Size
   )
 {
   EFI_STATUS            Status;
@@ -597,7 +603,7 @@ SecLoadFromCore (
   //
   // Find the SEC Core Entry Point
   //
-  Status = SecPeCoffGetEntryPoint (PeiCorePe32File, (VOID **)&PeiCoreEntryPoint);
+  Status = SecUefiImageGetEntryPoint (PeiCorePe32File, PeiCorePe32Size, (VOID **)&PeiCoreEntryPoint);
   if (EFI_ERROR (Status)) {
     return;
   }
@@ -731,53 +737,28 @@ SecEmuThunkAddress (
 
 RETURN_STATUS
 EFIAPI
-SecPeCoffGetEntryPoint (
+SecUefiImageGetEntryPoint (
-  IN     VOID  *Pe32Data,
+  IN     VOID   *Pe32Data,
-  IN OUT VOID  **EntryPoint
+  IN     UINT32 Pe32Size,
+  IN OUT VOID   **EntryPoint
   )
 {
-  EFI_STATUS                    Status;
+  EFI_STATUS                      Status;
-  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT ImageContext;
 
-  ZeroMem (&ImageContext, sizeof (ImageContext));
+  Status                  = UefiImageInitializeContext (&ImageContext, Pe32Data, Pe32Size);
-  ImageContext.Handle    = Pe32Data;
-  ImageContext.ImageRead = (PE_COFF_LOADER_READ_FILE)SecImageRead;
-
-  Status = PeCoffLoaderGetImageInfo (&ImageContext);
   if (EFI_ERROR (Status)) {
     return Status;
   }
 
-  if (ImageContext.ImageAddress != (UINTN)Pe32Data) {
+  // FIXME: Why cannot the Image be in-place already?
-    //
+  Status = UefiImageRelocateImageInplaceForExecution (&ImageContext);
-    // Relocate image to match the address where it resides
+  if (EFI_ERROR (Status)) {
-    //
+    return Status;
-    ImageContext.ImageAddress = (UINTN)Pe32Data;
-    Status                    = PeCoffLoaderLoadImage (&ImageContext);
-    if (EFI_ERROR (Status)) {
-      return Status;
-    }
-
-    Status = PeCoffLoaderRelocateImage (&ImageContext);
-    if (EFI_ERROR (Status)) {
-      return Status;
-    }
-  } else {
-    //
-    // Or just return image entry point
-    //
-    ImageContext.PdbPointer = PeCoffLoaderGetPdbPointer (Pe32Data);
-    Status                  = PeCoffLoaderGetEntryPoint (Pe32Data, EntryPoint);
-    if (EFI_ERROR (Status)) {
-      return Status;
-    }
-
-    ImageContext.EntryPoint = (UINTN)*EntryPoint;
   }
 
-  // On Unix a dlopen is done that will change the entry point
+  SecUefiImageRelocateImageExtraAction (&ImageContext);
-  SecPeCoffRelocateImageExtraAction (&ImageContext);
+  *EntryPoint = (VOID *) (UefiImageLoaderGetImageEntryPoint (&ImageContext));
-  *EntryPoint = (VOID *)(UINTN)ImageContext.EntryPoint;
 
   return Status;
 }
@@ -900,114 +881,9 @@ Returns:
   return EFI_SUCCESS;
 }
 
-/*++
-
-Routine Description:
-  Store the ModHandle in an array indexed by the Pdb File name.
-  The ModHandle is needed to unload the image.
-
-Arguments:
-  ImageContext - Input data returned from PE Loader Library. Used to find the
-                 .PDB file name of the PE Image.
-  ModHandle    - Returned from LoadLibraryEx() and stored for call to
-                 FreeLibrary().
-
-Returns:
-  EFI_SUCCESS - ModHandle was stored.
-
-**/
-EFI_STATUS
-AddHandle (
-  IN  PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext,
-  IN  VOID                          *ModHandle
-  )
-{
-  UINTN                        Index;
-  IMAGE_CONTEXT_TO_MOD_HANDLE  *Array;
-  UINTN                        PreviousSize;
-
-  Array = mImageContextModHandleArray;
-  for (Index = 0; Index < mImageContextModHandleArraySize; Index++, Array++) {
-    if (Array->ImageContext == NULL) {
-      //
-      // Make a copy of the string and store the ModHandle
-      //
-      Array->ImageContext = ImageContext;
-      Array->ModHandle    = ModHandle;
-      return EFI_SUCCESS;
-    }
-  }
-
-  //
-  // No free space in mImageContextModHandleArray so grow it by
-  // IMAGE_CONTEXT_TO_MOD_HANDLE entires. realloc will
-  // copy the old values to the new location. But it does
-  // not zero the new memory area.
-  //
-  PreviousSize                     = mImageContextModHandleArraySize * sizeof (IMAGE_CONTEXT_TO_MOD_HANDLE);
-  mImageContextModHandleArraySize += MAX_IMAGE_CONTEXT_TO_MOD_HANDLE_ARRAY_SIZE;
-
-  mImageContextModHandleArray = ReallocatePool (
-                                  (mImageContextModHandleArraySize - 1) * sizeof (IMAGE_CONTEXT_TO_MOD_HANDLE),
-                                  mImageContextModHandleArraySize * sizeof (IMAGE_CONTEXT_TO_MOD_HANDLE),
-                                  mImageContextModHandleArray
-                                  );
-  if (mImageContextModHandleArray == NULL) {
-    ASSERT (FALSE);
-    return EFI_OUT_OF_RESOURCES;
-  }
-
-  memset (mImageContextModHandleArray + PreviousSize, 0, MAX_IMAGE_CONTEXT_TO_MOD_HANDLE_ARRAY_SIZE * sizeof (IMAGE_CONTEXT_TO_MOD_HANDLE));
-
-  return AddHandle (ImageContext, ModHandle);
-}
-
-/*++
-
-Routine Description:
-  Return the ModHandle and delete the entry in the array.
-
-Arguments:
-  ImageContext - Input data returned from PE Loader Library. Used to find the
-                 .PDB file name of the PE Image.
-
-Returns:
-  ModHandle - ModHandle associated with ImageContext is returned
-  NULL      - No ModHandle associated with ImageContext
-
-**/
-VOID *
-RemoveHandle (
-  IN  PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
-  )
-{
-  UINTN                        Index;
-  IMAGE_CONTEXT_TO_MOD_HANDLE  *Array;
-
-  if (ImageContext->PdbPointer == NULL) {
-    //
-    // If no PDB pointer there is no ModHandle so return NULL
-    //
-    return NULL;
-  }
-
-  Array = mImageContextModHandleArray;
-  for (Index = 0; Index < mImageContextModHandleArraySize; Index++, Array++) {
-    if (Array->ImageContext == ImageContext) {
-      //
-      // If you find a match return it and delete the entry
-      //
-      Array->ImageContext = NULL;
-      return Array->ModHandle;
-    }
-  }
-
-  return NULL;
-}
-
 BOOLEAN
 IsPdbFile (
-  IN  CHAR8  *PdbFileName
+  IN  CONST CHAR8   *PdbFileName
   )
 {
   UINTN  Len;
@@ -1035,23 +911,29 @@ IsPdbFile (
 
 void
 PrintLoadAddress (
-  IN PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT       *ImageContext
   )
 {
-  if (ImageContext->PdbPointer == NULL) {
+  EFI_STATUS                   Status;
+  CONST CHAR8                  *PdbPath;
+  UINT32                       PdbPathSize;
+
+  Status = UefiImageGetSymbolsPath (ImageContext, &PdbPath, &PdbPathSize);
+
+  if (EFI_ERROR (Status)) {
     fprintf (
       stderr,
       "0x%08lx Loading NO DEBUG with entry point 0x%08lx\n",
-      (unsigned long)(ImageContext->ImageAddress),
+      (unsigned long) UefiImageLoaderGetImageAddress (ImageContext),
-      (unsigned long)ImageContext->EntryPoint
+      (unsigned long) UefiImageLoaderGetImageEntryPoint (ImageContext)
       );
   } else {
     fprintf (
       stderr,
       "0x%08lx Loading %s with entry point 0x%08lx\n",
-      (unsigned long)(ImageContext->ImageAddress + ImageContext->SizeOfHeaders),
+      (unsigned long) UefiImageLoaderGetImageAddress (ImageContext),
-      ImageContext->PdbPointer,
+      PdbPath,
-      (unsigned long)ImageContext->EntryPoint
+      (unsigned long) UefiImageLoaderGetImageEntryPoint (ImageContext)
       );
   }
 
@@ -1059,71 +941,12 @@ PrintLoadAddress (
   fflush (stderr);
 }
 
-/**
-  Loads the image using dlopen so symbols will be automatically
-  loaded by gdb.
-
-  @param  ImageContext  The PE/COFF image context
-
-  @retval TRUE - The image was successfully loaded
-  @retval FALSE - The image was successfully loaded
-
-**/
-BOOLEAN
-DlLoadImage (
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
-  )
-{
- #ifdef __APPLE__
-
-  return FALSE;
-
- #else
-
-  void  *Handle = NULL;
-  void  *Entry  = NULL;
-
-  if (ImageContext->PdbPointer == NULL) {
-    return FALSE;
-  }
-
-  if (!IsPdbFile (ImageContext->PdbPointer)) {
-    return FALSE;
-  }
-
-  fprintf (
-    stderr,
-    "Loading %s 0x%08lx - entry point 0x%08lx\n",
-    ImageContext->PdbPointer,
-    (unsigned long)ImageContext->ImageAddress,
-    (unsigned long)ImageContext->EntryPoint
-    );
-
-  Handle = dlopen (ImageContext->PdbPointer, RTLD_NOW);
-  if (Handle != NULL) {
-    Entry = dlsym (Handle, "_ModuleEntryPoint");
-    AddHandle (ImageContext, Handle);
-  } else {
-    printf ("%s\n", dlerror ());
-  }
-
-  if (Entry != NULL) {
-    ImageContext->EntryPoint = (UINTN)Entry;
-    printf ("Change %s Entrypoint to :0x%08lx\n", ImageContext->PdbPointer, (unsigned long)Entry);
-    return TRUE;
-  } else {
-    return FALSE;
-  }
-
- #endif
-}
-
 #ifdef __APPLE__
 __attribute__ ((noinline))
 #endif
 VOID
 SecGdbScriptBreak (
-  char               *FileName,
+  const char          *FileName,
   int                FileNameLength,
   long unsigned int  LoadAddress,
   int                AddSymbolFlag
@@ -1141,28 +964,37 @@ SecGdbScriptBreak (
 **/
 VOID
 GdbScriptAddImage (
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT      *ImageContext
   )
 {
+  EFI_STATUS  Status;
+  CONST CHAR8 *PdbPath;
+  UINT32      PdbPathSize;
+
   PrintLoadAddress (ImageContext);
 
-  if ((ImageContext->PdbPointer != NULL) && !IsPdbFile (ImageContext->PdbPointer)) {
+  Status = UefiImageGetSymbolsPath ((ImageContext, &PdbPath,) &PdbPathSize);
+  if (EFI_ERROR (Status)) {
+    return;
+  }
+
+  if (!IsPdbFile (PdbPath)) {
     FILE  *GdbTempFile;
     if (FeaturePcdGet (PcdEmulatorLazyLoadSymbols)) {
       GdbTempFile = fopen (gGdbWorkingFileName, "a");
       if (GdbTempFile != NULL) {
-        long unsigned int  SymbolsAddr = (long unsigned int)(ImageContext->ImageAddress + ImageContext->SizeOfHeaders);
+        long unsigned int SymbolsAddr = (long unsigned int)UefiImageLoaderGetImageAddress (ImageContext);
         mScriptSymbolChangesCount++;
         fprintf (
           GdbTempFile,
           "AddFirmwareSymbolFile 0x%x %s 0x%08lx\n",
           mScriptSymbolChangesCount,
-          ImageContext->PdbPointer,
+          PdbPath,
           SymbolsAddr
           );
         fclose (GdbTempFile);
         // This is for the lldb breakpoint only
-        SecGdbScriptBreak (ImageContext->PdbPointer, strlen (ImageContext->PdbPointer) + 1, (long unsigned int)(ImageContext->ImageAddress + ImageContext->SizeOfHeaders), 1);
+        SecGdbScriptBreak (PdbPath, PdbPathSize, (long unsigned int)UefiImageLoaderGetImageAddress (ImageContext), 1);
       } else {
         ASSERT (FALSE);
       }
@@ -1171,9 +1003,9 @@ GdbScriptAddImage (
       if (GdbTempFile != NULL) {
         fprintf (
           GdbTempFile,
-          "add-symbol-file %s 0x%08lx\n",
+          "add-symbol-file %s -o 0x%08lx\n",
-          ImageContext->PdbPointer,
+          PdbPath,
-          (long unsigned int)(ImageContext->ImageAddress + ImageContext->SizeOfHeaders)
+          (long unsigned int)UefiImageLoaderGetImageAddress (ImageContext)
           );
         fclose (GdbTempFile);
 
@@ -1183,7 +1015,7 @@ GdbScriptAddImage (
         // Also used for the lldb breakpoint script. The lldb breakpoint script does
         // not use the file, it uses the arguments.
         //
-        SecGdbScriptBreak (ImageContext->PdbPointer, strlen (ImageContext->PdbPointer) + 1, (long unsigned int)(ImageContext->ImageAddress + ImageContext->SizeOfHeaders), 1);
+        SecGdbScriptBreak (PdbPath, PdbPathSize, (long unsigned int)UefiImageLoaderGetImageAddress (ImageContext), 1);
       } else {
         ASSERT (FALSE);
       }
@@ -1193,13 +1025,11 @@ GdbScriptAddImage (
 
 VOID
 EFIAPI
-SecPeCoffRelocateImageExtraAction (
+SecUefiImageRelocateImageExtraAction (
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT      *ImageContext
   )
 {
-  if (!DlLoadImage (ImageContext)) {
+  GdbScriptAddImage (ImageContext);
-    GdbScriptAddImage (ImageContext);
-  }
 }
 
 /**
@@ -1211,15 +1041,23 @@ SecPeCoffRelocateImageExtraAction (
 **/
 VOID
 GdbScriptRemoveImage (
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT      *ImageContext
   )
 {
-  FILE  *GdbTempFile;
+  FILE        *GdbTempFile;
+  EFI_STATUS  Status;
+  CONST CHAR8 *PdbPath;
+  UINT32      PdbPathSize;
+
+  Status = UefiImageGetSymbolsPath (ImageContext, &PdbPath, &PdbPathSize);
+  if (EFI_ERROR (Status)) {
+    return;
+  }
 
   //
   // Need to skip .PDB files created from VC++
   //
-  if (IsPdbFile (ImageContext->PdbPointer)) {
+  if (IsPdbFile (PdbPath)) {
     return;
   }
 
@@ -1234,24 +1072,24 @@ GdbScriptRemoveImage (
         GdbTempFile,
         "RemoveFirmwareSymbolFile 0x%x %s\n",
         mScriptSymbolChangesCount,
-        ImageContext->PdbPointer
+        PdbPath
         );
       fclose (GdbTempFile);
-      SecGdbScriptBreak (ImageContext->PdbPointer, strlen (ImageContext->PdbPointer) + 1, 0, 0);
+      SecGdbScriptBreak (PdbPath, PdbPathSize, 0, 0);
     } else {
       ASSERT (FALSE);
     }
   } else {
     GdbTempFile = fopen (gGdbWorkingFileName, "w");
     if (GdbTempFile != NULL) {
-      fprintf (GdbTempFile, "remove-symbol-file %s\n", ImageContext->PdbPointer);
+      fprintf (GdbTempFile, "remove-symbol-file %s\n", PdbPath);
       fclose (GdbTempFile);
 
       //
       // Target for gdb breakpoint in a script that uses gGdbWorkingFileName to set a breakpoint.
       // Hey what can you say scripting in gdb is not that great....
       //
-      SecGdbScriptBreak (ImageContext->PdbPointer, strlen (ImageContext->PdbPointer) + 1, 0, 0);
+      SecGdbScriptBreak (PdbPath, PdbPathSize, 0, 0);
     } else {
       ASSERT (FALSE);
     }
@@ -1260,22 +1098,9 @@ GdbScriptRemoveImage (
 
 VOID
 EFIAPI
-SecPeCoffUnloadImageExtraAction (
+SecUefiImageUnloadImageExtraAction (
-  IN PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT      *ImageContext
   )
 {
-  VOID  *Handle;
-
-  //
-  // Check to see if the image symbols were loaded with gdb script, or dlopen
-  //
-  Handle = RemoveHandle (ImageContext);
-  if (Handle != NULL) {
- #ifndef __APPLE__
-    dlclose (Handle);
- #endif
-    return;
-  }
-
   GdbScriptRemoveImage (ImageContext);
 }

EmulatorPkg/Unix/Host/Host.h

@@ -85,7 +85,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <PiPei.h>
 #include <Uefi.h>
 
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
 #include <Library/BaseLib.h>
 #include <Library/BaseMemoryLib.h>
 #include <Library/MemoryAllocationLib.h>
@@ -97,7 +97,6 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/ThunkPpiList.h>
 #include <Library/ThunkProtocolList.h>
 #include <Library/PeiServicesLib.h>
-#include <Library/PeCoffGetEntryPointLib.h>
 #include <Library/EmuMagicPageLib.h>
 
 #include <Ppi/EmuThunk.h>
@@ -137,7 +136,7 @@ typedef struct {
 #define MAX_IMAGE_CONTEXT_TO_MOD_HANDLE_ARRAY_SIZE  0x100
 
 typedef struct {
-  PE_COFF_LOADER_IMAGE_CONTEXT    *ImageContext;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT   *ImageContext;
   VOID                            *ModHandle;
 } IMAGE_CONTEXT_TO_MOD_HANDLE;
 
@@ -162,7 +161,8 @@ SecLoadFromCore (
   IN  UINTN  LargestRegion,
   IN  UINTN  LargestRegionSize,
   IN  UINTN  BootFirmwareVolumeBase,
-  IN  VOID   *PeiCoreFile
+  IN  VOID    *PeiCoreFile,
+  IN  UINT32  PeiCorePe32Size
   );
 
 EFI_STATUS
@@ -194,17 +194,11 @@ SecFfsFindSectionData (
   );
 
 EFI_STATUS
-EFIAPI
+SecFfsFindSectionData2 (
-SecUnixPeCoffLoaderLoadAsDll (
+  IN EFI_SECTION_TYPE      SectionType,
-  IN CHAR8  *PdbFileName,
+  IN EFI_FFS_FILE_HEADER   *FfsFileHeader,
-  IN VOID   **ImageEntryPoint,
+  IN OUT VOID              **SectionData,
-  OUT VOID  **ModHandle
+  OUT UINT32               *SectionDataSize
-  );
-
-EFI_STATUS
-EFIAPI
-SecUnixPeCoffLoaderFreeLibrary (
-  OUT VOID  *ModHandle
   );
 
 EFI_STATUS
@@ -228,7 +222,7 @@ GasketSecUnixFdAddress (
 
 EFI_STATUS
 GetImageReadFunction (
-  IN PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT          *ImageContext,
   IN EFI_PHYSICAL_ADDRESS          *TopOfMemory
   );
 
@@ -273,21 +267,22 @@ GasketSecTemporaryRamSupport (
 
 RETURN_STATUS
 EFIAPI
-SecPeCoffGetEntryPoint (
+SecUefiImageGetEntryPoint (
   IN     VOID  *Pe32Data,
+  IN     UINT32 Pe32Size,
   IN OUT VOID  **EntryPoint
   );
 
 VOID
 EFIAPI
-SecPeCoffRelocateImageExtraAction (
+SecUefiImageRelocateImageExtraAction (
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   );
 
 VOID
 EFIAPI
-SecPeCoffLoaderUnloadImageExtraAction (
+SecUefiImageLoaderUnloadImageExtraAction (
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   );
 
 VOID

EmulatorPkg/Unix/Host/Host.inf

@@ -55,12 +55,11 @@
   PrintLib
   BaseMemoryLib
   BaseLib
-  PeCoffLib
+  UefiImageLib
   ThunkPpiList
   ThunkProtocolList
   PpiListLib
   PeiServicesLib
-  PeCoffGetEntryPointLib
   StackCheckLib
 
 [Ppis]

EmulatorPkg/Unix/Host/Ia32/Gasket.S

@@ -313,24 +313,26 @@ ASM_PFX(GasketSecGetNextProtocol):
 
 // PPIs produced by SEC
 
-ASM_GLOBAL ASM_PFX(GasketSecPeCoffGetEntryPoint)
+ASM_GLOBAL ASM_PFX(GasketSecUefiImageGetEntryPoint)
-ASM_PFX(GasketSecPeCoffGetEntryPoint):
+ASM_PFX(GasketSecUefiImageGetEntryPoint):
   pushl %ebp
   movl  %esp, %ebp
-  subl  $24, %esp      // sub extra 16 from the stack for alignment
+  subl  $40, %esp      // sub extra 16 from the stack for alignment
   and   $-16, %esp    // stack needs to end in 0xFFFFFFF0 before call
+  movl  16(%ebp), %eax
+  movl  %eax, 8(%esp)
   movl  12(%ebp), %eax
   movl  %eax, 4(%esp)
   movl  8(%ebp), %eax
   movl  %eax, (%esp)
 
-  call    ASM_PFX(SecPeCoffGetEntryPoint)
+  call    ASM_PFX(SecUefiImageGetEntryPoint)
 
   leave
   ret
 
-ASM_GLOBAL ASM_PFX(GasketSecPeCoffRelocateImageExtraAction)
+ASM_GLOBAL ASM_PFX(GasketSecUefiImageRelocateImageExtraAction)
-ASM_PFX(GasketSecPeCoffRelocateImageExtraAction):
+ASM_PFX(GasketSecUefiImageRelocateImageExtraAction):
   pushl %ebp
   movl  %esp, %ebp
   subl  $24, %esp      // sub extra 16 from the stack for alignment
@@ -338,13 +340,13 @@ ASM_PFX(GasketSecPeCoffRelocateImageExtraAction):
   movl  8(%ebp), %eax
   movl  %eax, (%esp)
 
-  call    ASM_PFX(SecPeCoffRelocateImageExtraAction)
+  call    ASM_PFX(SecUefiImageRelocateImageExtraAction)
 
   leave
   ret
 
-ASM_GLOBAL ASM_PFX(GasketSecPeCoffUnloadImageExtraAction)
+ASM_GLOBAL ASM_PFX(GasketSecUefiImageUnloadImageExtraAction)
-ASM_PFX(GasketSecPeCoffUnloadImageExtraAction):
+ASM_PFX(GasketSecUefiImageUnloadImageExtraAction):
   pushl %ebp
   movl  %esp, %ebp
   subl  $24, %esp      // sub extra 16 from the stack for alignment
@@ -352,7 +354,7 @@ ASM_PFX(GasketSecPeCoffUnloadImageExtraAction):
   movl  8(%ebp), %eax
   movl  %eax, (%esp)
 
-  call    ASM_PFX(SecPeCoffUnloadImageExtraAction)
+  call    ASM_PFX(SecUefiImageUnloadImageExtraAction)
 
   leave
   ret

EmulatorPkg/Unix/Host/X64/Gasket.S

@@ -358,8 +358,8 @@ ASM_PFX(GasketSecGetNextProtocol):
 
 // PPIs produced by SEC
 
-ASM_GLOBAL ASM_PFX(GasketSecPeCoffGetEntryPoint)
+ASM_GLOBAL ASM_PFX(GasketSecUefiImageGetEntryPoint)
-ASM_PFX(GasketSecPeCoffGetEntryPoint):
+ASM_PFX(GasketSecUefiImageGetEntryPoint):
   pushq   %rbp            // stack frame is for the debugger
   movq    %rsp, %rbp
 
@@ -368,16 +368,17 @@ ASM_PFX(GasketSecPeCoffGetEntryPoint):
 
   movq    %rcx, %rdi    // Swizzle args
   movq    %rdx, %rsi
+  movq    %r8,  %rdx
 
-  call    ASM_PFX(SecPeCoffGetEntryPoint)
+  call    ASM_PFX(SecUefiImageGetEntryPoint)
 
   popq    %rdi          // restore state
   popq    %rsi
   popq    %rbp
   ret
 
-ASM_GLOBAL ASM_PFX(GasketSecPeCoffRelocateImageExtraAction)
+ASM_GLOBAL ASM_PFX(GasketSecUefiImageRelocateImageExtraAction)
-ASM_PFX(GasketSecPeCoffRelocateImageExtraAction):
+ASM_PFX(GasketSecUefiImageRelocateImageExtraAction):
   pushq   %rbp            // stack frame is for the debugger
   movq    %rsp, %rbp
 
@@ -386,15 +387,15 @@ ASM_PFX(GasketSecPeCoffRelocateImageExtraAction):
 
   movq    %rcx, %rdi    // Swizzle args
 
-  call    ASM_PFX(SecPeCoffRelocateImageExtraAction)
+  call    ASM_PFX(SecUefiImageRelocateImageExtraAction)
 
   popq    %rdi          // restore state
   popq    %rsi
   popq    %rbp
   ret
 
-ASM_GLOBAL ASM_PFX(GasketSecPeCoffUnloadImageExtraAction)
+ASM_GLOBAL ASM_PFX(GasketSecUefiImageUnloadImageExtraAction)
-ASM_PFX(GasketSecPeCoffUnloadImageExtraAction):
+ASM_PFX(GasketSecUefiImageUnloadImageExtraAction):
   pushq   %rbp            // stack frame is for the debugger
   movq    %rsp, %rbp
 
@@ -403,7 +404,7 @@ ASM_PFX(GasketSecPeCoffUnloadImageExtraAction):
 
   movq    %rcx, %rdi    // Swizzle args
 
-  call    ASM_PFX(SecPeCoffUnloadImageExtraAction)
+  call    ASM_PFX(SecUefiImageUnloadImageExtraAction)
 
   popq    %rdi          // restore state
   popq    %rsi

EmulatorPkg/Unix/lldbefi.py

@@ -361,7 +361,7 @@ def LoadEmulatorEfiSymbols(frame, bp_loc , internal_dict):
     #
     # VOID
     # SecGdbScriptBreak (
-    #   char                *FileName,
+    #   const char          *FileName,
     #   int                 FileNameLength,
     #   long unsigned int   LoadAddress,
     #   int                 AddSymbolFlag
@@ -394,7 +394,7 @@ def LoadEmulatorEfiSymbols(frame, bp_loc , internal_dict):
 
     debugger = frame.thread.process.target.debugger
     if frame.FindVariable ("AddSymbolFlag").GetValueAsUnsigned() == 1:
-        LoadAddress = frame.FindVariable ("LoadAddress").GetValueAsUnsigned() - 0x240
+        LoadAddress = frame.FindVariable ("LoadAddress").GetValueAsUnsigned()
 
         debugger.HandleCommand ("target modules add  %s" % FileName)
         print("target modules load --slid 0x%x %s" % (LoadAddress, FileName))

EmulatorPkg/Win/Host/WinHost.c

@@ -444,6 +444,7 @@ Returns:
   BOOLEAN              Done;
   EFI_PEI_FILE_HANDLE  FileHandle;
   VOID                 *SecFile;
+  UINT32                SecFileSize;
   CHAR16               *MemorySizeStr;
   CHAR16               *FirmwareVolumesStr;
   UINTN                ProcessAffinityMask;
@@ -451,6 +452,7 @@ Returns:
   INT32                LowBit;
   UINTN                ResetJumpCode;
   EMU_THUNK_PPI        *SecEmuThunkPpi;
+  UINT32               AuthenticationStatus;
 
   //
   // If enabled use the magic page to communicate between modules
@@ -648,7 +650,14 @@ Returns:
                      &FileHandle
                      );
       if (!EFI_ERROR (Status)) {
-        Status = PeiServicesFfsFindSectionData (EFI_SECTION_PE32, FileHandle, &SecFile);
+        Status = PeiServicesFfsFindSectionData4 (
+                   EFI_SECTION_PE32,
+                   0,
+                   FileHandle,
+                   &SecFile,
+                   &SecFileSize,
+                   &AuthenticationStatus
+                   );
         if (!EFI_ERROR (Status)) {
           SecPrint (" contains SEC Core");
         }
@@ -679,7 +688,7 @@ Returns:
   //
   // Hand off to SEC Core
   //
-  SecLoadSecCore ((UINTN)TemporaryRam, TemporaryRamSize, gFdInfo[0].Address, gFdInfo[0].Size, SecFile);
+  SecLoadSecCore ((UINTN)TemporaryRam, TemporaryRamSize, gFdInfo[0].Address, gFdInfo[0].Size, SecFile, SecFileSize);
 
   //
   // If we get here, then the SEC Core returned. This is an error as SEC should
@@ -695,7 +704,8 @@ SecLoadSecCore (
   IN  UINTN  TemporaryRamSize,
   IN  VOID   *BootFirmwareVolumeBase,
   IN  UINTN  BootFirmwareVolumeSize,
-  IN  VOID   *SecCorePe32File
+  IN  VOID   *SecCorePe32File,
+  IN  UINT32  SecCorePe32Size
   )
 
 /*++
@@ -759,8 +769,9 @@ Returns:
   //
   // Load the PEI Core from a Firmware Volume
   //
-  Status = SecPeCoffGetEntryPoint (
+  Status = SecUefiImageGetEntryPoint (
              SecCorePe32File,
+            SecCorePe32Size,
              &SecCoreEntryPoint
              );
   if (EFI_ERROR (Status)) {
@@ -784,40 +795,47 @@ Returns:
 
 RETURN_STATUS
 EFIAPI
-SecPeCoffGetEntryPoint (
+SecUefiImageGetEntryPoint (
-  IN     VOID  *Pe32Data,
+  IN     VOID   *Pe32Data,
-  IN OUT VOID  **EntryPoint
+  IN     UINT32 Pe32Size,
+  IN OUT VOID   **EntryPoint
   )
 {
-  EFI_STATUS                    Status;
+  EFI_STATUS                       Status;
-  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT  ImageContext;
+  VOID                             *Dest;
+  UINT32                           DestSize;
 
-  ZeroMem (&ImageContext, sizeof (ImageContext));
+  Status                  = UefiImageInitializeContext (&ImageContext, Pe32Data, Pe32Size);
-  ImageContext.Handle = Pe32Data;
-
-  ImageContext.ImageRead = (PE_COFF_LOADER_READ_FILE)SecImageRead;
-
-  Status = PeCoffLoaderGetImageInfo (&ImageContext);
   if (EFI_ERROR (Status)) {
     return Status;
   }
 
   //
-  // XIP for SEC and PEI_CORE
+  // Allocate space in NT (not emulator) memory with ReadWrite and Execute attribute.
+  // Extra space is for alignment
   //
-  ImageContext.ImageAddress = (EFI_PHYSICAL_ADDRESS)(UINTN)Pe32Data;
+  Status = UefiImageLoaderGetDestinationSize(&ImageContext, &DestSize);
-
-  Status = PeCoffLoaderLoadImage (&ImageContext);
   if (EFI_ERROR (Status)) {
     return Status;
   }
 
-  Status = PeCoffLoaderRelocateImage (&ImageContext);
+  Dest = VirtualAlloc (NULL, (SIZE_T) DestSize, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
+  if (Dest == NULL) {
+    return EFI_OUT_OF_RESOURCES;
+  }
+
+  Status = UefiImageLoadImage (&ImageContext, Dest, DestSize);
   if (EFI_ERROR (Status)) {
     return Status;
   }
 
-  *EntryPoint = (VOID *)(UINTN)ImageContext.EntryPoint;
+  Status = UefiImageRelocateImage (&ImageContext, (UINTN) Dest, NULL, 0);
+  if (EFI_ERROR (Status)) {
+    return Status;
+  }
+
+  *EntryPoint   = (VOID *) (UefiImageLoaderGetImageEntryPoint (&ImageContext));
 
   return EFI_SUCCESS;
 }
@@ -951,7 +969,7 @@ Returns:
 --*/
 EFI_STATUS
 AddModHandle (
-  IN  PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext,
+  IN  UEFI_IMAGE_LOADER_IMAGE_CONTEXT      *ImageContext,
   IN  VOID                          *ModHandle
   )
 
@@ -1028,7 +1046,7 @@ AddModHandle (
 **/
 VOID *
 RemoveModHandle (
-  IN  PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN  UEFI_IMAGE_LOADER_IMAGE_CONTEXT      *ImageContext
   )
 {
   UINTN                   Index;
@@ -1064,8 +1082,8 @@ typedef struct {
 
 VOID
 EFIAPI
-PeCoffLoaderRelocateImageExtraAction (
+UefiImageLoaderRelocateImageExtraAction (
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT      *ImageContext
   )
 {
   EFI_STATUS                           Status;
@@ -1103,7 +1121,7 @@ PeCoffLoaderRelocateImageExtraAction (
   // Load the DLL if it's not an EBC image.
   //
   if ((ImageContext->PdbPointer != NULL) &&
-      (ImageContext->Machine != EFI_IMAGE_MACHINE_EBC))
+      (UefiImageGetMachine (ImageContext) != EFI_IMAGE_MACHINE_EBC))
   {
     //
     // Convert filename from ASCII to Unicode
@@ -1117,7 +1135,7 @@ PeCoffLoaderRelocateImageExtraAction (
       free (DllFileName);
 
       //
-      // Never return an error if PeCoffLoaderRelocateImage() succeeded.
+      // Never return an error if UefiImageRelocateImage() succeeded.
       // The image will run, but we just can't source level debug. If we
       // return an error the image will not run.
       //
@@ -1316,8 +1334,8 @@ PeCoffLoaderRelocateImageExtraAction (
 
 VOID
 EFIAPI
-PeCoffLoaderUnloadImageExtraAction (
+UefiImageLoaderUnloadImageExtraAction (
-  IN PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   )
 {
   VOID  *ModHandle;

EmulatorPkg/Win/Host/WinHost.h

@@ -21,7 +21,7 @@ Abstract:
 #include "WinInclude.h"
 
 #include <PiPei.h>
-#include <IndustryStandard/PeImage.h>
+#include <IndustryStandard/PeImage2.h>
 #include <Guid/FileInfo.h>
 #include <Guid/FileSystemInfo.h>
 #include <Guid/FileSystemVolumeLabelInfo.h>
@@ -35,7 +35,7 @@ Abstract:
 #include <Protocol/EmuSnp.h>
 
 #include <Library/BaseLib.h>
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
 #include <Library/DebugLib.h>
 #include <Library/MemoryAllocationLib.h>
 #include <Library/ThunkPpiList.h>
@@ -44,7 +44,7 @@ Abstract:
 #include <Library/PrintLib.h>
 #include <Library/BaseMemoryLib.h>
 #include <Library/PeiServicesLib.h>
-#include <Library/PeCoffExtraActionLib.h>
+#include <Library/UefiImageExtraActionLib.h>
 #include <Library/NetLib.h>
 
 #define TEMPORARY_RAM_SIZE  0x20000
@@ -61,8 +61,9 @@ typedef struct {
 
 RETURN_STATUS
 EFIAPI
-SecPeCoffGetEntryPoint (
+SecUefiImageGetEntryPoint (
   IN     VOID  *Pe32Data,
+  IN     UINT32 Pe32Size,
   IN OUT VOID  **EntryPoint
   );
 
@@ -72,7 +73,8 @@ SecLoadSecCore (
   IN  UINTN  TemporaryRamSize,
   IN  VOID   *BootFirmwareVolumeBase,
   IN  UINTN  BootFirmwareVolumeSize,
-  IN  VOID   *SecCorePe32File
+  IN  VOID   *SecCorePe32File,
+  IN  UINT32  SecCorePe32Size
   )
 
 /*++

EmulatorPkg/Win/Host/WinHost.inf

@@ -49,7 +49,7 @@
   PrintLib
   BaseMemoryLib
   BaseLib
-  PeCoffLib
+  UefiImageLib
   ThunkPpiList
   ThunkProtocolList
   PpiListLib

EmulatorPkg/Win/Host/WinThunk.c

@@ -592,9 +592,9 @@ EMU_THUNK_PROTOCOL  gEmuThunkProtocol = {
   SecAlloc,
   NULL,
   SecFree,
-  SecPeCoffGetEntryPoint,
+  SecUefiImageGetEntryPoint,
-  PeCoffLoaderRelocateImageExtraAction,
+  UefiImageLoaderRelocateImageExtraAction,
-  PeCoffLoaderUnloadImageExtraAction,
+  UefiImageLoaderUnloadImageExtraAction,
   SecEnableInterrupt,
   SecDisableInterrupt,
   SecQueryPerformanceFrequency,

IntelFsp2Pkg/Tools/Tests/QemuFspPkg.dsc

@@ -90,11 +90,11 @@
   PeiServicesTablePointerLib|MdePkg/Library/PeiServicesTablePointerLibIdt/PeiServicesTablePointerLibIdt.inf
   PeiServicesLib|MdePkg/Library/PeiServicesLib/PeiServicesLib.inf
   MemoryAllocationLib|MdePkg/Library/PeiMemoryAllocationLib/PeiMemoryAllocationLib.inf
-  PeCoffGetEntryPointLib|MdePkg/Library/BasePeCoffGetEntryPointLib/BasePeCoffGetEntryPointLib.inf
   ReportStatusCodeLib|MdeModulePkg/Library/PeiReportStatusCodeLib/PeiReportStatusCodeLib.inf
   CacheMaintenanceLib|MdePkg/Library/BaseCacheMaintenanceLib/BaseCacheMaintenanceLib.inf
-  PeCoffLib|MdePkg/Library/BasePeCoffLib/BasePeCoffLib.inf
+  PeCoffLib2|MdePkg/Library/BasePeCoffLib2/BasePeCoffLib2.inf
-  PeCoffExtraActionLib|MdePkg/Library/BasePeCoffExtraActionLibNull/BasePeCoffExtraActionLibNull.inf
+  UefiImageLib|MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibPeCoff.inf
+  UefiImageExtraActionLib|MdePkg/Library/BaseUefiImageExtraActionLibNull/BaseUefiImageExtraActionLibNull.inf
   UefiDecompressLib|MdePkg/Library/BaseUefiDecompressLib/BaseUefiDecompressLib.inf
   SynchronizationLib|MdePkg/Library/BaseSynchronizationLib/BaseSynchronizationLib.inf
   CpuLib|MdePkg/Library/BaseCpuLib/BaseCpuLib.inf

IntelFsp2WrapperPkg/FspWrapperNotifyDxe/FspWrapperNotifyDxe.inf

@@ -39,12 +39,13 @@
   BaseMemoryLib
   UefiLib
   FspWrapperApiLib
-  PeCoffLib
+  UefiImageLib
   CacheMaintenanceLib
   DxeServicesLib
   PerformanceLib
   HobLib
   FspWrapperPlatformLib
+  MemoryAllocationLib
 
 [Protocols]
   gEfiPciEnumerationCompleteProtocolGuid            ## CONSUMES

IntelFsp2WrapperPkg/FspWrapperNotifyDxe/LoadBelow4G.c

@@ -11,11 +11,13 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/UefiDriverEntryPoint.h>
 #include <Library/BaseMemoryLib.h>
 #include <Library/DebugLib.h>
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
 #include <Library/UefiBootServicesTableLib.h>
 #include <Library/DxeServicesLib.h>
 #include <Library/CacheMaintenanceLib.h>
 #include <Library/UefiLib.h>
+#include <Library/MemoryAllocationLib.h>
+#include <Library/MemoryAllocationLibEx.h>
 
 /**
   Relocate this image under 4G memory.
@@ -37,9 +39,11 @@ RelocateImageUnder4GIfNeeded (
   UINT8                         *Buffer;
   UINTN                         BufferSize;
   EFI_HANDLE                    NewImageHandle;
+  UINT32                        ImageSize;
+  UINT32                        ImageAlignment;
   UINTN                         Pages;
   EFI_PHYSICAL_ADDRESS          FfsBuffer;
-  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT ImageContext;
   VOID                          *Interface;
 
   //
@@ -83,43 +87,34 @@ RelocateImageUnder4GIfNeeded (
              &BufferSize
              );
   ASSERT_EFI_ERROR (Status);
-  ImageContext.Handle    = Buffer;
-  ImageContext.ImageRead = PeCoffLoaderImageReadFromMemory;
   //
   // Get information about the image being loaded
   //
-  Status = PeCoffLoaderGetImageInfo (&ImageContext);
+  Status = UefiImageInitializeContext (&ImageContext, Buffer, (UINT32) BufferSize);
   ASSERT_EFI_ERROR (Status);
-  if (ImageContext.SectionAlignment > EFI_PAGE_SIZE) {
+  ImageSize      = UefiImageGetImageSize (&ImageContext);
-    Pages = EFI_SIZE_TO_PAGES ((UINTN)(ImageContext.ImageSize + ImageContext.SectionAlignment));
+  ImageAlignment = UefiImageGetSegmentAlignment (&ImageContext);
-  } else {
+  Pages = EFI_SIZE_TO_PAGES (ImageSize);
-    Pages = EFI_SIZE_TO_PAGES ((UINTN)ImageContext.ImageSize);
-  }
 
   FfsBuffer = 0xFFFFFFFF;
-  Status    = gBS->AllocatePages (
+  Status    = AllocateAlignedPagesEx (
-                     AllocateMaxAddress,
+                AllocateMaxAddress,
-                     EfiBootServicesCode,
+                EfiBootServicesCode,
-                     Pages,
+                Pages,
-                     &FfsBuffer
+                ImageAlignment,
-                     );
+                &FfsBuffer
+                );
   ASSERT_EFI_ERROR (Status);
-  ImageContext.ImageAddress = (PHYSICAL_ADDRESS)(UINTN)FfsBuffer;
   //
-  // Align buffer on section boundary
+  // Load and relocate the image to our new buffer
   //
-  ImageContext.ImageAddress += ImageContext.SectionAlignment - 1;
+  Status = UefiImageLoadImageForExecution (
-  ImageContext.ImageAddress &= ~((EFI_PHYSICAL_ADDRESS)ImageContext.SectionAlignment - 1);
+             &ImageContext,
-  //
+             (VOID *) (UINTN) FfsBuffer,
-  // Load the image to our new buffer
+             ImageSize,
-  //
+             NULL,
-  Status = PeCoffLoaderLoadImage (&ImageContext);
+             0
-  ASSERT_EFI_ERROR (Status);
+             );
-
-  //
-  // Relocate the image in our new buffer
-  //
-  Status = PeCoffLoaderRelocateImage (&ImageContext);
   ASSERT_EFI_ERROR (Status);
 
   //
@@ -127,16 +122,11 @@ RelocateImageUnder4GIfNeeded (
   //
   gBS->FreePool (Buffer);
 
-  //
+  DEBUG ((DEBUG_INFO, "Loading driver at 0x%08llx EntryPoint=0x%08x\n", FfsBuffer, UefiImageLoaderGetImageEntryPoint (&ImageContext)));
-  // Flush the instruction cache so the image data is written before we execute it
+  Status = ((EFI_IMAGE_ENTRY_POINT)(UefiImageLoaderGetImageEntryPoint (&ImageContext)))(NewImageHandle, gST);
-  //
-  InvalidateInstructionCacheRange ((VOID *)(UINTN)ImageContext.ImageAddress, (UINTN)ImageContext.ImageSize);
-
-  DEBUG ((DEBUG_INFO, "Loading driver at 0x%08x EntryPoint=0x%08x\n", (UINTN)ImageContext.ImageAddress, (UINTN)ImageContext.EntryPoint));
-  Status = ((EFI_IMAGE_ENTRY_POINT)(UINTN)(ImageContext.EntryPoint))(NewImageHandle, gST);
   if (EFI_ERROR (Status)) {
-    DEBUG ((DEBUG_ERROR, "Error: Image at 0x%08x start failed: %r\n", ImageContext.ImageAddress, Status));
+    DEBUG ((DEBUG_ERROR, "Error: Image at 0x%08llx start failed: %r\n", FfsBuffer, Status));
-    gBS->FreePages (FfsBuffer, Pages);
+    FreeAlignedPages ((VOID *)(UINTN)FfsBuffer, Pages);
   }
 
   //

IntelFsp2WrapperPkg/FspmWrapperPeim/FspmWrapperPeim.inf

@@ -38,8 +38,8 @@
   FspWrapperPlatformLib
   FspWrapperHobProcessLib
   CpuLib
-  PeCoffGetEntryPointLib
+  UefiCpuLib
-  PeCoffExtraActionLib
+  UefiImageExtraActionLib
   PerformanceLib
   TimerLib
   FspWrapperApiLib

IntelFsp2WrapperPkg/FspsWrapperPeim/FspsWrapperPeim.inf

@@ -39,8 +39,8 @@
   FspWrapperPlatformLib
   FspWrapperHobProcessLib
   CpuLib
-  PeCoffGetEntryPointLib
+  UefiCpuLib
-  PeCoffExtraActionLib
+  UefiImageExtraActionLib
   PerformanceLib
   FspWrapperApiLib
   FspWrapperApiTestLib

IntelFsp2WrapperPkg/IntelFsp2WrapperPkg.dsc

@@ -27,9 +27,9 @@
   DebugPrintErrorLevelLib|MdePkg/Library/BaseDebugPrintErrorLevelLib/BaseDebugPrintErrorLevelLib.inf
   PciCf8Lib|MdePkg/Library/BasePciCf8Lib/BasePciCf8Lib.inf
   IoLib|MdePkg/Library/BaseIoLibIntrinsic/BaseIoLibIntrinsic.inf
-  PeCoffGetEntryPointLib|MdePkg/Library/BasePeCoffGetEntryPointLib/BasePeCoffGetEntryPointLib.inf
+  UefiImageExtraActionLib|MdePkg/Library/BaseUefiImageExtraActionLibNull/BaseUefiImageExtraActionLibNull.inf
-  PeCoffExtraActionLib|MdePkg/Library/BasePeCoffExtraActionLibNull/BasePeCoffExtraActionLibNull.inf
+  PeCoffLib2|MdePkg/Library/BasePeCoffLib2/BasePeCoffLib2.inf
-  PeCoffLib|MdePkg/Library/BasePeCoffLib/BasePeCoffLib.inf
+  UefiImageLib|MdePkg/Library/BaseUefiImageLib/BaseUefiImageLibPeCoff.inf
   CacheMaintenanceLib|MdePkg/Library/BaseCacheMaintenanceLib/BaseCacheMaintenanceLib.inf
 
   PerformanceLib|MdePkg/Library/BasePerformanceLibNull/BasePerformanceLibNull.inf

MdeModulePkg/Bus/Pci/PciBusDxe/PciBus.h

@@ -40,7 +40,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/PcdLib.h>
 
 #include <IndustryStandard/Pci.h>
-#include <IndustryStandard/PeImage.h>
+#include <IndustryStandard/PeImage2.h>
 #include <IndustryStandard/Acpi.h>
 
 typedef struct _PCI_IO_DEVICE  PCI_IO_DEVICE;

MdeModulePkg/Core/Dxe/DxeMain.h

@@ -73,11 +73,11 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/ExtractGuidedSectionLib.h>
 #include <Library/CacheMaintenanceLib.h>
 #include <Library/BaseMemoryLib.h>
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
-#include <Library/PeCoffGetEntryPointLib.h>
+#include <Library/UefiImageExtraActionLib.h>
-#include <Library/PeCoffExtraActionLib.h>
 #include <Library/PcdLib.h>
 #include <Library/MemoryAllocationLib.h>
+#include <Library/MemoryAllocationLibEx.h>
 #include <Library/DevicePathLib.h>
 #include <Library/UefiBootServicesTableLib.h>
 #include <Library/ReportStatusCodeLib.h>
@@ -220,10 +220,10 @@ typedef struct {
   EFI_RUNTIME_IMAGE_ENTRY                 *RuntimeData;
   /// Pointer to Loaded Image Device Path Protocol
   EFI_DEVICE_PATH_PROTOCOL                *LoadedImageDevicePath;
-  /// PeCoffLoader ImageContext
-  PE_COFF_LOADER_IMAGE_CONTEXT            ImageContext;
   /// Status returned by LoadImage() service.
   EFI_STATUS                              LoadImageStatus;
+
+  VOID *HiiData;
 } LOADED_IMAGE_PRIVATE_DATA;
 
 #define LOADED_IMAGE_PRIVATE_DATA_FROM_THIS(a) \
@@ -390,7 +390,8 @@ CoreInitializeEventServices (
 **/
 EFI_STATUS
 CoreInitializeImageServices (
-  IN  VOID  *HobStart
+  IN  VOID                         *HobStart,
+  OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT *ImageContext
   );
 
 /**
@@ -2385,7 +2386,8 @@ VOID
 CoreNewDebugImageInfoEntry (
   IN  UINT32                     ImageInfoType,
   IN  EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage,
-  IN  EFI_HANDLE                 ImageHandle
+  IN  EFI_HANDLE                 ImageHandle,
+  IN  UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   );
 
 /**
@@ -2574,7 +2576,8 @@ VerifyFvHeaderChecksum (
 **/
 VOID
 MemoryProfileInit (
-  IN VOID  *HobStart
+  IN VOID                         *HobStart,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT *ImageContext
   );
 
 /**
@@ -2600,8 +2603,9 @@ MemoryProfileInstallProtocol (
 **/
 EFI_STATUS
 RegisterMemoryProfileImage (
-  IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,
+  IN EFI_DEVICE_PATH_PROTOCOL   *FilePath,
-  IN EFI_FV_FILETYPE            FileType
+  IN EFI_FV_FILETYPE            FileType,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT      *ImageContext
   );
 
 /**
@@ -2617,7 +2621,8 @@ RegisterMemoryProfileImage (
 **/
 EFI_STATUS
 UnregisterMemoryProfileImage (
-  IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry
+  EFI_DEVICE_PATH_PROTOCOL  *FilePath,
+  IN EFI_PHYSICAL_ADDRESS  ImageAddress
   );
 
 /**
@@ -2704,7 +2709,8 @@ InstallMemoryAttributesTableOnMemoryAllocation (
 **/
 VOID
 InsertImageRecord (
-  IN EFI_RUNTIME_IMAGE_ENTRY  *RuntimeImage
+  IN LOADED_IMAGE_PRIVATE_DATA   *Image,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT *ImageContext
   );
 
 /**
@@ -2725,8 +2731,8 @@ RemoveImageRecord (
 **/
 VOID
 ProtectUefiImage (
-  IN EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage,
+  IN LOADED_IMAGE_PRIVATE_DATA  *Image,
-  IN EFI_DEVICE_PATH_PROTOCOL   *LoadedImageDevicePath
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT *ImageContext
   );
 
 /**

MdeModulePkg/Core/Dxe/DxeMain.inf

@@ -82,9 +82,8 @@
   UefiLib
   DebugLib
   DxeCoreEntryPoint
-  PeCoffLib
+  UefiImageLib
-  PeCoffGetEntryPointLib
+  UefiImageExtraActionLib
-  PeCoffExtraActionLib
   ExtractGuidedSectionLib
   MemoryAllocationLib
   UefiBootServicesTableLib

MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c

@@ -236,15 +236,14 @@ DxeMain (
   IN  VOID  *HobStart
   )
 {
-  EFI_STATUS                    Status;
+  EFI_STATUS                      Status;
-  EFI_PHYSICAL_ADDRESS          MemoryBaseAddress;
+  EFI_PHYSICAL_ADDRESS            MemoryBaseAddress;
-  UINT64                        MemoryLength;
+  UINT64                          MemoryLength;
-  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT ImageContext;
-  UINTN                         Index;
+  UINTN                           Index;
-  EFI_HOB_GUID_TYPE             *GuidHob;
+  EFI_HOB_GUID_TYPE               *GuidHob;
-  EFI_VECTOR_HANDOFF_INFO       *VectorInfoList;
+  EFI_VECTOR_HANDOFF_INFO         *VectorInfoList;
-  EFI_VECTOR_HANDOFF_INFO       *VectorInfo;
+  EFI_VECTOR_HANDOFF_INFO         *VectorInfo;
-  VOID                          *EntryPoint;
 
   //
   // Setup the default exception handlers
@@ -276,8 +275,6 @@ DxeMain (
   //
   CoreInitializeMemoryServices (&HobStart, &MemoryBaseAddress, &MemoryLength);
 
-  MemoryProfileInit (HobStart);
-
   //
   // Start the Handle Services.
   //
@@ -287,9 +284,11 @@ DxeMain (
   //
   // Start the Image Services.
   //
-  Status = CoreInitializeImageServices (HobStart);
+  Status = CoreInitializeImageServices (HobStart, &ImageContext);
   ASSERT_EFI_ERROR (Status);
 
+  MemoryProfileInit (HobStart, &ImageContext);
+
   //
   // Initialize the Global Coherency Domain Services
   //
@@ -335,19 +334,9 @@ DxeMain (
 
   //
   // Report DXE Core image information to the PE/COFF Extra Action Library
+  // FIXME: This is done by DxeIpl, why is this needed here? Difference PEI/DXE?
   //
-  ZeroMem (&ImageContext, sizeof (ImageContext));
+  UefiImageLoaderRelocateImageExtraAction (&ImageContext);
-  ImageContext.ImageAddress  = (EFI_PHYSICAL_ADDRESS)(UINTN)gDxeCoreLoadedImage->ImageBase;
-  ImageContext.PdbPointer    = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)ImageContext.ImageAddress);
-  ImageContext.SizeOfHeaders = PeCoffGetSizeOfHeaders ((VOID *)(UINTN)ImageContext.ImageAddress);
-  Status                     = PeCoffLoaderGetEntryPoint ((VOID *)(UINTN)ImageContext.ImageAddress, &EntryPoint);
-  if (Status == EFI_SUCCESS) {
-    ImageContext.EntryPoint = (EFI_PHYSICAL_ADDRESS)(UINTN)EntryPoint;
-  }
-
-  ImageContext.Handle    = (VOID *)(UINTN)gDxeCoreLoadedImage->ImageBase;
-  ImageContext.ImageRead = PeCoffLoaderImageReadFromMemory;
-  PeCoffLoaderRelocateImageExtraAction (&ImageContext);
 
   //
   // Install the DXE Services Table into the EFI System Tables's Configuration Table
@@ -394,7 +383,8 @@ DxeMain (
   CoreNewDebugImageInfoEntry (
     EFI_DEBUG_IMAGE_INFO_TYPE_NORMAL,
     gDxeCoreLoadedImage,
-    gDxeCoreImageHandle
+    gDxeCoreImageHandle,
+    &ImageContext
     );
 
   DEBUG ((DEBUG_INFO | DEBUG_LOAD, "HOBLIST address in DXE = 0x%p\n", HobStart));

MdeModulePkg/Core/Dxe/Mem/HeapGuard.c

@@ -1533,7 +1533,7 @@ PromoteGuardedFreePages (
   OUT EFI_PHYSICAL_ADDRESS  *EndAddress
   )
 {
-  EFI_STATUS            Status;
+  EFI_STATUS              Status;
   UINTN                 AvailablePages;
   UINT64                Bitmap;
   EFI_PHYSICAL_ADDRESS  Start;

MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c

@@ -8,6 +8,10 @@
 
 #include "DxeMain.h"
 #include "Imem.h"
+#include "Library/UefiImageLib.h"
+#include "ProcessorBind.h"
+#include "Protocol/LoadedImage.h"
+#include "Uefi/UefiBaseType.h"
 
 #define IS_UEFI_MEMORY_PROFILE_ENABLED  ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 0)
 
@@ -104,8 +108,7 @@ EFIAPI
 ProfileProtocolRegisterImage (
   IN EDKII_MEMORY_PROFILE_PROTOCOL  *This,
   IN EFI_DEVICE_PATH_PROTOCOL       *FilePath,
-  IN PHYSICAL_ADDRESS               ImageBase,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT    *ImageContext,
-  IN UINT64                         ImageSize,
   IN EFI_FV_FILETYPE                FileType
   );
 
@@ -247,110 +250,6 @@ GetMemoryProfileContext (
   return mMemoryProfileContextPtr;
 }
 
-/**
-  Retrieves and returns the Subsystem of a PE/COFF image that has been loaded into system memory.
-  If Pe32Data is NULL, then ASSERT().
-
-  @param Pe32Data   The pointer to the PE/COFF image that is loaded in system memory.
-
-  @return The Subsystem of the PE/COFF image.
-
-**/
-UINT16
-InternalPeCoffGetSubsystem (
-  IN VOID  *Pe32Data
-  )
-{
-  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
-  EFI_IMAGE_DOS_HEADER                 *DosHdr;
-  UINT16                               Magic;
-
-  ASSERT (Pe32Data != NULL);
-
-  DosHdr = (EFI_IMAGE_DOS_HEADER *)Pe32Data;
-  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
-    //
-    // DOS image header is present, so read the PE header after the DOS image header.
-    //
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINTN)Pe32Data + (UINTN)((DosHdr->e_lfanew) & 0x0ffff));
-  } else {
-    //
-    // DOS image header is not present, so PE header is at the image base.
-    //
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)Pe32Data;
-  }
-
-  if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {
-    return Hdr.Te->Subsystem;
-  } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {
-    Magic = Hdr.Pe32->OptionalHeader.Magic;
-    if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
-      return Hdr.Pe32->OptionalHeader.Subsystem;
-    } else if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
-      return Hdr.Pe32Plus->OptionalHeader.Subsystem;
-    }
-  }
-
-  return 0x0000;
-}
-
-/**
-  Retrieves and returns a pointer to the entry point to a PE/COFF image that has been loaded
-  into system memory with the PE/COFF Loader Library functions.
-
-  Retrieves the entry point to the PE/COFF image specified by Pe32Data and returns this entry
-  point in EntryPoint.  If the entry point could not be retrieved from the PE/COFF image, then
-  return RETURN_INVALID_PARAMETER.  Otherwise return RETURN_SUCCESS.
-  If Pe32Data is NULL, then ASSERT().
-  If EntryPoint is NULL, then ASSERT().
-
-  @param  Pe32Data                  The pointer to the PE/COFF image that is loaded in system memory.
-  @param  EntryPoint                The pointer to entry point to the PE/COFF image to return.
-
-  @retval RETURN_SUCCESS            EntryPoint was returned.
-  @retval RETURN_INVALID_PARAMETER  The entry point could not be found in the PE/COFF image.
-
-**/
-RETURN_STATUS
-InternalPeCoffGetEntryPoint (
-  IN  VOID  *Pe32Data,
-  OUT VOID  **EntryPoint
-  )
-{
-  EFI_IMAGE_DOS_HEADER                 *DosHdr;
-  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
-
-  ASSERT (Pe32Data   != NULL);
-  ASSERT (EntryPoint != NULL);
-
-  DosHdr = (EFI_IMAGE_DOS_HEADER *)Pe32Data;
-  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
-    //
-    // DOS image header is present, so read the PE header after the DOS image header.
-    //
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINTN)Pe32Data + (UINTN)((DosHdr->e_lfanew) & 0x0ffff));
-  } else {
-    //
-    // DOS image header is not present, so PE header is at the image base.
-    //
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)Pe32Data;
-  }
-
-  //
-  // Calculate the entry point relative to the start of the image.
-  // AddressOfEntryPoint is common for PE32 & PE32+
-  //
-  if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {
-    *EntryPoint = (VOID *)((UINTN)Pe32Data + (UINTN)(Hdr.Te->AddressOfEntryPoint & 0x0ffffffff) + sizeof (EFI_TE_IMAGE_HEADER) - Hdr.Te->StrippedSize);
-    return RETURN_SUCCESS;
-  } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {
-    *EntryPoint = (VOID *)((UINTN)Pe32Data + (UINTN)(Hdr.Pe32->OptionalHeader.AddressOfEntryPoint & 0x0ffffffff));
-    return RETURN_SUCCESS;
-  }
-
-  return RETURN_UNSUPPORTED;
-}
-
 /**
   Build driver info.
 
@@ -367,32 +266,27 @@ InternalPeCoffGetEntryPoint (
 **/
 MEMORY_PROFILE_DRIVER_INFO_DATA *
 BuildDriverInfo (
-  IN MEMORY_PROFILE_CONTEXT_DATA  *ContextData,
+  IN MEMORY_PROFILE_CONTEXT_DATA      *ContextData,
-  IN EFI_GUID                     *FileName,
+  IN EFI_GUID                         *FileName,
-  IN PHYSICAL_ADDRESS             ImageBase,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext,
-  IN UINT64                       ImageSize,
+  IN EFI_FV_FILETYPE                  FileType
-  IN PHYSICAL_ADDRESS             EntryPoint,
-  IN UINT16                       ImageSubsystem,
-  IN EFI_FV_FILETYPE              FileType
   )
 {
+  RETURN_STATUS                     PdbStatus;
   EFI_STATUS                       Status;
   MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
   MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
-  VOID                             *EntryPointInImage;
+  CONST CHAR8                       *PdbString;
-  CHAR8                            *PdbString;
+  UINT32                            PdbSize;
-  UINTN                            PdbSize;
   UINTN                            PdbOccupiedSize;
 
-  PdbSize         = 0;
   PdbOccupiedSize = 0;
-  PdbString       = NULL;
+
-  if (ImageBase != 0) {
+  ASSERT (UefiImageLoaderGetImageAddress (ImageContext) != 0);
-    PdbString = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)ImageBase);
+
-    if (PdbString != NULL) {
+  PdbStatus = UefiImageGetSymbolsPath (ImageContext, &PdbString, &PdbSize);
-      PdbSize         = AsciiStrSize (PdbString);
+  if (!EFI_ERROR (PdbStatus)) {
-      PdbOccupiedSize = GET_OCCUPIED_SIZE (PdbSize, sizeof (UINT64));
+    PdbOccupiedSize = GET_OCCUPIED_SIZE (PdbSize, sizeof (UINT64));
-    }
   }
 
   //
@@ -420,19 +314,10 @@ BuildDriverInfo (
     CopyMem (&DriverInfo->FileName, FileName, sizeof (EFI_GUID));
   }
 
-  DriverInfo->ImageBase      = ImageBase;
+  DriverInfo->ImageBase      = UefiImageLoaderGetImageAddress (ImageContext);
-  DriverInfo->ImageSize      = ImageSize;
+  DriverInfo->ImageSize      = UefiImageGetImageSize (ImageContext);
-  DriverInfo->EntryPoint     = EntryPoint;
+  DriverInfo->EntryPoint     = UefiImageLoaderGetImageEntryPoint (ImageContext);
-  DriverInfo->ImageSubsystem = ImageSubsystem;
+  DriverInfo->ImageSubsystem = UefiImageGetSubsystem (ImageContext);
-  if ((EntryPoint != 0) && ((EntryPoint < ImageBase) || (EntryPoint >= (ImageBase + ImageSize)))) {
-    //
-    // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.
-    // So patch ImageBuffer here to align the EntryPoint.
-    //
-    Status = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)ImageBase, &EntryPointInImage);
-    ASSERT_EFI_ERROR (Status);
-    DriverInfo->ImageBase = ImageBase + EntryPoint - (PHYSICAL_ADDRESS)(UINTN)EntryPointInImage;
-  }
 
   DriverInfo->FileType          = FileType;
   DriverInfoData->AllocInfoList = (LIST_ENTRY *)(DriverInfoData + 1);
@@ -440,7 +325,7 @@ BuildDriverInfo (
   DriverInfo->CurrentUsage     = 0;
   DriverInfo->PeakUsage        = 0;
   DriverInfo->AllocRecordCount = 0;
-  if (PdbSize != 0) {
+  if (!RETURN_ERROR (PdbStatus)) {
     DriverInfo->PdbStringOffset = (UINT16)sizeof (MEMORY_PROFILE_DRIVER_INFO);
     DriverInfoData->PdbString   = (CHAR8 *)(DriverInfoData->AllocInfoList + 1);
     CopyMem (DriverInfoData->PdbString, PdbString, PdbSize);
@@ -528,13 +413,13 @@ NeedRecordThisDriver (
 **/
 BOOLEAN
 RegisterDxeCore (
-  IN VOID                         *HobStart,
+  IN VOID                             *HobStart,
-  IN MEMORY_PROFILE_CONTEXT_DATA  *ContextData
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext,
+  IN MEMORY_PROFILE_CONTEXT_DATA      *ContextData
   )
 {
   EFI_PEI_HOB_POINTERS               DxeCoreHob;
   MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData;
-  PHYSICAL_ADDRESS                   ImageBase;
   UINT8                              TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];
   MEDIA_FW_VOL_FILEPATH_DEVICE_PATH  *FilePath;
 
@@ -565,14 +450,10 @@ RegisterDxeCore (
     return FALSE;
   }
 
-  ImageBase      = DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryBaseAddress;
   DriverInfoData = BuildDriverInfo (
                      ContextData,
                      &DxeCoreHob.MemoryAllocationModule->ModuleName,
-                     ImageBase,
+                     ImageContext,
-                     DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryLength,
-                     DxeCoreHob.MemoryAllocationModule->EntryPoint,
-                     InternalPeCoffGetSubsystem ((VOID *)(UINTN)ImageBase),
                      EFI_FV_FILETYPE_DXE_CORE
                      );
   if (DriverInfoData == NULL) {
@@ -590,7 +471,8 @@ RegisterDxeCore (
 **/
 VOID
 MemoryProfileInit (
-  IN VOID  *HobStart
+  IN VOID                            *HobStart,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT *ImageContext
   )
 {
   MEMORY_PROFILE_CONTEXT_DATA  *ContextData;
@@ -615,7 +497,7 @@ MemoryProfileInit (
   mMemoryProfileDriverPath     = AllocateCopyPool (mMemoryProfileDriverPathSize, PcdGetPtr (PcdMemoryProfileDriverPath));
   mMemoryProfileContextPtr     = &mMemoryProfileContext;
 
-  RegisterDxeCore (HobStart, &mMemoryProfileContext);
+  RegisterDxeCore (HobStart, ImageContext, &mMemoryProfileContext);
 
   DEBUG ((DEBUG_INFO, "MemoryProfileInit MemoryProfileContext - 0x%x\n", &mMemoryProfileContext));
 }
@@ -692,8 +574,9 @@ GetFileNameFromFilePath (
 **/
 EFI_STATUS
 RegisterMemoryProfileImage (
-  IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,
+  IN EFI_DEVICE_PATH_PROTOCOL         *FilePath,
-  IN EFI_FV_FILETYPE            FileType
+  IN EFI_FV_FILETYPE                  FileType,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   )
 {
   MEMORY_PROFILE_CONTEXT_DATA      *ContextData;
@@ -703,7 +586,7 @@ RegisterMemoryProfileImage (
     return EFI_UNSUPPORTED;
   }
 
-  if (!NeedRecordThisDriver (DriverEntry->Info.FilePath)) {
+  if (!NeedRecordThisDriver (FilePath)) {
     return EFI_UNSUPPORTED;
   }
 
@@ -714,11 +597,8 @@ RegisterMemoryProfileImage (
 
   DriverInfoData = BuildDriverInfo (
                      ContextData,
-                     GetFileNameFromFilePath (DriverEntry->Info.FilePath),
+                     GetFileNameFromFilePath (FilePath),
-                     DriverEntry->ImageContext.ImageAddress,
+                     ImageContext,
-                     DriverEntry->ImageContext.ImageSize,
-                     DriverEntry->ImageContext.EntryPoint,
-                     DriverEntry->ImageContext.ImageType,
                      FileType
                      );
   if (DriverInfoData == NULL) {
@@ -831,21 +711,21 @@ GetMemoryProfileDriverInfoFromAddress (
 **/
 EFI_STATUS
 UnregisterMemoryProfileImage (
-  IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry
+  EFI_DEVICE_PATH_PROTOCOL  *FilePath,
+  IN EFI_PHYSICAL_ADDRESS  ImageAddress
   )
 {
-  EFI_STATUS                       Status;
+  //EFI_STATUS                        Status;
   MEMORY_PROFILE_CONTEXT_DATA      *ContextData;
   MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
   EFI_GUID                         *FileName;
-  PHYSICAL_ADDRESS                 ImageAddress;
+  //VOID                              *EntryPointInImage;
-  VOID                             *EntryPointInImage;
 
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {
     return EFI_UNSUPPORTED;
   }
 
-  if (!NeedRecordThisDriver (DriverEntry->Info.FilePath)) {
+  if (!NeedRecordThisDriver (FilePath)) {
     return EFI_UNSUPPORTED;
   }
 
@@ -855,17 +735,7 @@ UnregisterMemoryProfileImage (
   }
 
   DriverInfoData = NULL;
-  FileName       = GetFileNameFromFilePath (DriverEntry->Info.FilePath);
+  FileName = GetFileNameFromFilePath (FilePath);
-  ImageAddress   = DriverEntry->ImageContext.ImageAddress;
-  if ((DriverEntry->ImageContext.EntryPoint < ImageAddress) || (DriverEntry->ImageContext.EntryPoint >= (ImageAddress + DriverEntry->ImageContext.ImageSize))) {
-    //
-    // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.
-    // So patch ImageAddress here to align the EntryPoint.
-    //
-    Status = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)ImageAddress, &EntryPointInImage);
-    ASSERT_EFI_ERROR (Status);
-    ImageAddress = ImageAddress + (UINTN)DriverEntry->ImageContext.EntryPoint - (UINTN)EntryPointInImage;
-  }
 
   if (FileName != NULL) {
     DriverInfoData = GetMemoryProfileDriverInfoByFileNameAndAddress (ContextData, FileName, ImageAddress);
@@ -1629,27 +1499,13 @@ ProfileProtocolGetData (
 EFI_STATUS
 EFIAPI
 ProfileProtocolRegisterImage (
-  IN EDKII_MEMORY_PROFILE_PROTOCOL  *This,
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,
-  IN EFI_DEVICE_PATH_PROTOCOL       *FilePath,
+  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,
-  IN PHYSICAL_ADDRESS               ImageBase,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT    *ImageContext,
-  IN UINT64                         ImageSize,
+  IN EFI_FV_FILETYPE                    FileType
-  IN EFI_FV_FILETYPE                FileType
   )
 {
-  EFI_STATUS                 Status;
+  return RegisterMemoryProfileImage (FilePath, FileType, ImageContext);
-  LOADED_IMAGE_PRIVATE_DATA  DriverEntry;
-  VOID                       *EntryPointInImage;
-
-  ZeroMem (&DriverEntry, sizeof (DriverEntry));
-  DriverEntry.Info.FilePath             = FilePath;
-  DriverEntry.ImageContext.ImageAddress = ImageBase;
-  DriverEntry.ImageContext.ImageSize    = ImageSize;
-  Status                                = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)ImageBase, &EntryPointInImage);
-  ASSERT_EFI_ERROR (Status);
-  DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS)(UINTN)EntryPointInImage;
-  DriverEntry.ImageContext.ImageType  = InternalPeCoffGetSubsystem ((VOID *)(UINTN)ImageBase);
-
-  return RegisterMemoryProfileImage (&DriverEntry, FileType);
 }
 
 /**
@@ -1675,19 +1531,7 @@ ProfileProtocolUnregisterImage (
   IN UINT64                         ImageSize
   )
 {
-  EFI_STATUS                 Status;
+  return UnregisterMemoryProfileImage (FilePath, ImageBase);
-  LOADED_IMAGE_PRIVATE_DATA  DriverEntry;
-  VOID                       *EntryPointInImage;
-
-  ZeroMem (&DriverEntry, sizeof (DriverEntry));
-  DriverEntry.Info.FilePath             = FilePath;
-  DriverEntry.ImageContext.ImageAddress = ImageBase;
-  DriverEntry.ImageContext.ImageSize    = ImageSize;
-  Status                                = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)ImageBase, &EntryPointInImage);
-  ASSERT_EFI_ERROR (Status);
-  DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS)(UINTN)EntryPointInImage;
-
-  return UnregisterMemoryProfileImage (&DriverEntry);
 }
 
 /**

MdeModulePkg/Core/Dxe/Misc/DebugImageInfo.c

@@ -160,15 +160,20 @@ CoreUpdateDebugTableCrc32 (
 **/
 VOID
 CoreNewDebugImageInfoEntry (
-  IN  UINT32                     ImageInfoType,
+  IN     UINT32                           ImageInfoType,
-  IN  EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage,
+  IN     EFI_LOADED_IMAGE_PROTOCOL        *LoadedImage,
-  IN  EFI_HANDLE                 ImageHandle
+  IN     EFI_HANDLE                       ImageHandle,
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   )
 {
-  EFI_DEBUG_IMAGE_INFO  *Table;
+  EFI_DEBUG_IMAGE_INFO        *Table;
-  EFI_DEBUG_IMAGE_INFO  *NewTable;
+  EFI_DEBUG_IMAGE_INFO        *NewTable;
-  UINTN                 Index;
+  UINTN                       Index;
-  UINTN                 TableSize;
+  UINTN                       TableSize;
+  EFI_DEBUG_IMAGE_INFO_NORMAL *NormalImage;
+  RETURN_STATUS               Status;
+  CONST CHAR8                 *PdbPath;
+  UINT32                      PdbPathSize;
 
   //
   // Set the flag indicating that we're in the process of updating the table.
@@ -205,6 +210,15 @@ CoreNewDebugImageInfoEntry (
     // Copy the old table into the new one
     //
     CopyMem (NewTable, Table, TableSize);
+    mDebugInfoTableHeader.EfiDebugImageInfoTable = NewTable;
+    //
+    // Set the first empty entry index to be the original max table entries.
+    //
+    Index = mMaxTableEntries;
+    //
+    // Enlarge the max table entries.
+    //
+    mMaxTableEntries += EFI_PAGE_SIZE / EFI_DEBUG_TABLE_ENTRY_SIZE;
     //
     // Free the old table
     //
@@ -212,32 +226,31 @@ CoreNewDebugImageInfoEntry (
     //
     // Update the table header
     //
-    Table                                        = NewTable;
+    Table = NewTable;
-    mDebugInfoTableHeader.EfiDebugImageInfoTable = NewTable;
-    //
-    // Enlarge the max table entries and set the first empty entry index to
-    // be the original max table entries.
-    //
-    Index             = mMaxTableEntries;
-    mMaxTableEntries += EFI_PAGE_SIZE / EFI_DEBUG_TABLE_ENTRY_SIZE;
   }
 
   //
   // Allocate data for new entry
   //
-  Table[Index].NormalImage = AllocateZeroPool (sizeof (EFI_DEBUG_IMAGE_INFO_NORMAL));
+  NormalImage = AllocateZeroPool (sizeof (EFI_DEBUG_IMAGE_INFO_NORMAL));
-  if (Table[Index].NormalImage != NULL) {
+  if (NormalImage != NULL) {
     //
     // Update the entry
     //
-    Table[Index].NormalImage->ImageInfoType               = (UINT32)ImageInfoType;
+    NormalImage->ImageInfoType               = (UINT32)ImageInfoType;
-    Table[Index].NormalImage->LoadedImageProtocolInstance = LoadedImage;
+    NormalImage->LoadedImageProtocolInstance = LoadedImage;
-    Table[Index].NormalImage->ImageHandle                 = ImageHandle;
+    NormalImage->ImageHandle                 = ImageHandle;
+
+    Status = UefiImageGetSymbolsPath (ImageContext, &PdbPath, &PdbPathSize);
+    if (!RETURN_ERROR (Status)) {
+      NormalImage->PdbPath = AllocateCopyPool (PdbPathSize, PdbPath);
+    }
     //
     // Increase the number of EFI_DEBUG_IMAGE_INFO elements and set the mDebugInfoTable in modified status.
     //
-    mDebugInfoTableHeader.TableSize++;
     mDebugInfoTableHeader.UpdateStatus |= EFI_DEBUG_IMAGE_INFO_TABLE_MODIFIED;
+    Table[Index].NormalImage = NormalImage;
+    mDebugInfoTableHeader.TableSize++;
   }
 
   mDebugInfoTableHeader.UpdateStatus &= ~EFI_DEBUG_IMAGE_INFO_UPDATE_IN_PROGRESS;
@@ -254,8 +267,9 @@ CoreRemoveDebugImageInfoEntry (
   EFI_HANDLE  ImageHandle
   )
 {
-  EFI_DEBUG_IMAGE_INFO  *Table;
+  EFI_DEBUG_IMAGE_INFO        *Table;
-  UINTN                 Index;
+  UINTN                       Index;
+  EFI_DEBUG_IMAGE_INFO_NORMAL *NormalImage;
 
   mDebugInfoTableHeader.UpdateStatus |= EFI_DEBUG_IMAGE_INFO_UPDATE_IN_PROGRESS;
 
@@ -267,13 +281,19 @@ CoreRemoveDebugImageInfoEntry (
       // Found a match. Free up the record, then NULL the pointer to indicate the slot
       // is free.
       //
-      CoreFreePool (Table[Index].NormalImage);
+      NormalImage = Table[Index].NormalImage;
-      Table[Index].NormalImage = NULL;
       //
       // Decrease the number of EFI_DEBUG_IMAGE_INFO elements and set the mDebugInfoTable in modified status.
       //
-      mDebugInfoTableHeader.TableSize--;
       mDebugInfoTableHeader.UpdateStatus |= EFI_DEBUG_IMAGE_INFO_TABLE_MODIFIED;
+      mDebugInfoTableHeader.TableSize--;
+      Table[Index].NormalImage = NULL;
+
+      if (NormalImage->PdbPath != NULL) {
+        FreePool (NormalImage->PdbPath);
+      }
+
+      CoreFreePool (NormalImage);
       break;
     }
   }

MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c

@@ -22,6 +22,8 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 
 #include "DxeMain.h"
 #include "HeapGuard.h"
+#include "IndustryStandard/PeImage2.h"
+#include "ProcessorBind.h"
 
 /**
   This function for GetMemoryMap() with properties table capability.
@@ -73,7 +75,7 @@ CoreGetMemoryMapWithSeparatedImageSection (
 typedef struct {
   UINT32        Signature;
   UINTN         ImageRecordCount;
-  UINTN         CodeSegmentCountMax;
+  UINTN         SectionCountMax;
   LIST_ENTRY    ImageRecordList;
 } IMAGE_PROPERTIES_PRIVATE_DATA;
 
@@ -201,7 +203,7 @@ InstallMemoryAttributesTable (
       case EfiRuntimeServicesCode:
       case EfiRuntimeServicesData:
         CopyMem (MemoryAttributesEntry, MemoryMap, DescriptorSize);
-        MemoryAttributesEntry->Attribute &= (EFI_MEMORY_RO|EFI_MEMORY_XP|EFI_MEMORY_RUNTIME);
+        MemoryAttributesEntry->Attribute &= EFI_MEMORY_ACCESS_MASK | EFI_MEMORY_RUNTIME;
         DEBUG ((DEBUG_VERBOSE, "Entry (0x%x)\n", MemoryAttributesEntry));
         DEBUG ((DEBUG_VERBOSE, "  Type              - 0x%x\n", MemoryAttributesEntry->Type));
         DEBUG ((DEBUG_VERBOSE, "  PhysicalStart     - 0x%016lx\n", MemoryAttributesEntry->PhysicalStart));
@@ -284,15 +286,6 @@ InstallMemoryAttributesTableOnEndOfDxe (
 {
   mMemoryAttributesTableEndOfDxe = TRUE;
   InstallMemoryAttributesTable ();
-
-  DEBUG_CODE_BEGIN ();
-  if ( mImagePropertiesPrivateData.ImageRecordCount > 0) {
-    DEBUG ((DEBUG_INFO, "DXE - Total Runtime Image Count: 0x%x\n", mImagePropertiesPrivateData.ImageRecordCount));
-    DEBUG ((DEBUG_INFO, "DXE - Dump Runtime Image Records:\n"));
-    DumpImageRecords (&mImagePropertiesPrivateData.ImageRecordList);
-  }
-
-  DEBUG_CODE_END ();
 }
 
 /**
@@ -534,7 +527,11 @@ CoreGetMemoryMapWithSeparatedImageSection (
 
   CoreAcquiremMemoryAttributesTableLock ();
 
-  AdditionalRecordCount = (2 * mImagePropertiesPrivateData.CodeSegmentCountMax + 3) * mImagePropertiesPrivateData.ImageRecordCount;
+  //
+  // Per image, there may be one additional trailer. There may be prefixed data
+  // (counted as the original entry).
+  //
+  AdditionalRecordCount = (mImagePropertiesPrivateData.SectionCountMax + 1) * mImagePropertiesPrivateData.ImageRecordCount;
 
   OldMemoryMapSize = *MemoryMapSize;
   Status           = CoreGetMemoryMap (MemoryMapSize, MemoryMap, MapKey, DescriptorSize, DescriptorVersion);
@@ -574,6 +571,68 @@ CoreGetMemoryMapWithSeparatedImageSection (
 // Below functions are for ImageRecord
 //
 
+/**
+  Set MemoryAttributesTable according to PE/COFF image section alignment.
+
+  @param  SectionAlignment    PE/COFF section alignment
+**/
+STATIC
+VOID
+SetMemoryAttributesTableSectionAlignment (
+  IN UINT32  SectionAlignment
+  )
+{
+  if (((SectionAlignment & (RUNTIME_PAGE_ALLOCATION_GRANULARITY - 1)) != 0) &&
+      mMemoryAttributesTableEnable)
+  {
+    DEBUG ((DEBUG_VERBOSE, "SetMemoryAttributesTableSectionAlignment - Clear\n"));
+    mMemoryAttributesTableEnable = FALSE;
+  }
+}
+
+/**
+  Sort image record based upon the ImageBase from low to high.
+**/
+STATIC
+VOID
+InsertSortImageRecord (
+  IN UEFI_IMAGE_RECORD  *NewImageRecord
+  )
+{
+  UEFI_IMAGE_RECORD  *ImageRecord;
+  LIST_ENTRY         *PrevImageRecordLink;
+  LIST_ENTRY         *ImageRecordLink;
+  LIST_ENTRY         *ImageRecordList;
+
+  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
+
+  PrevImageRecordLink = ImageRecordList;
+  for (
+    ImageRecordLink = GetFirstNode (ImageRecordList);
+    !IsNull (ImageRecordLink, ImageRecordList);
+    ImageRecordLink = GetNextNode (ImageRecordList, PrevImageRecordLink)
+    ) {
+    ImageRecord = CR (
+                    ImageRecordLink,
+                    UEFI_IMAGE_RECORD,
+                    Link,
+                    UEFI_IMAGE_RECORD_SIGNATURE
+                    );
+    if (NewImageRecord->StartAddress < ImageRecord->StartAddress) {
+      break;
+    }
+
+    PrevImageRecordLink = ImageRecordLink;
+  }
+
+  InsertHeadList (PrevImageRecordLink, &NewImageRecord->Link);
+  mImagePropertiesPrivateData.ImageRecordCount++;
+
+  if (mImagePropertiesPrivateData.SectionCountMax < NewImageRecord->NumSegments) {
+    mImagePropertiesPrivateData.SectionCountMax = NewImageRecord->NumSegments;
+  }
+}
+
 /**
   Insert image record.
 
@@ -581,13 +640,18 @@ CoreGetMemoryMapWithSeparatedImageSection (
 **/
 VOID
 InsertImageRecord (
-  IN EFI_RUNTIME_IMAGE_ENTRY  *RuntimeImage
+  IN     LOADED_IMAGE_PRIVATE_DATA        *Image,
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   )
 {
-  EFI_STATUS               Status;
+  RETURN_STATUS            PdbStatus;
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
+  EFI_RUNTIME_IMAGE_ENTRY  *RuntimeImage;
-  CHAR8                    *PdbPointer;
+  UINT32                   SectionAlignment;
-  UINT32                   RequiredAlignment;
+  UEFI_IMAGE_RECORD        *ImageRecord;
+  CONST CHAR8              *PdbPointer;
+  UINT32                   PdbSize;
+
+  RuntimeImage = Image->RuntimeData;
 
   DEBUG ((DEBUG_VERBOSE, "InsertImageRecord - 0x%x\n", RuntimeImage));
 
@@ -596,71 +660,85 @@ InsertImageRecord (
     return;
   }
 
-  ImageRecord = AllocatePool (sizeof (*ImageRecord));
+  DEBUG ((DEBUG_VERBOSE, "ImageRecordCount - 0x%x\n", mImagePropertiesPrivateData.ImageRecordCount));
-  if (ImageRecord == NULL) {
-    return;
-  }
 
-  InitializeListHead (&ImageRecord->Link);
+  PdbStatus = UefiImageGetSymbolsPath (ImageContext, &PdbPointer, &PdbSize);
-  InitializeListHead (&ImageRecord->CodeSegmentList);
+  if (!EFI_ERROR (PdbStatus)) {
-
-  PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)RuntimeImage->ImageBase);
-  if (PdbPointer != NULL) {
     DEBUG ((DEBUG_VERBOSE, "  Image - %a\n", PdbPointer));
   }
 
-  RequiredAlignment = RUNTIME_PAGE_ALLOCATION_GRANULARITY;
+  //
-  Status            = CreateImagePropertiesRecord (
+  // Get SectionAlignment
-                        RuntimeImage->ImageBase,
+  //
-                        RuntimeImage->ImageSize,
+  SectionAlignment  = UefiImageGetSegmentAlignment (ImageContext);
-                        &RequiredAlignment,
-                        ImageRecord
-                        );
 
-  if (EFI_ERROR (Status)) {
+  SetMemoryAttributesTableSectionAlignment (SectionAlignment);
-    if (Status == EFI_ABORTED) {
+  if ((SectionAlignment & (RUNTIME_PAGE_ALLOCATION_GRANULARITY - 1)) != 0) {
-      mMemoryAttributesTableEnable = FALSE;
+    DEBUG ((
+      DEBUG_WARN,
+      "!!!!!!!!  InsertImageRecord - Section Alignment(0x%x) is not %dK  !!!!!!!!\n",
+      SectionAlignment, RUNTIME_PAGE_ALLOCATION_GRANULARITY >> 10));
+    if (!EFI_ERROR (PdbStatus)) {
+      DEBUG ((DEBUG_WARN, "!!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));
     }
 
-    Status = EFI_ABORTED;
+    return;
-    goto Finish;
   }
 
-  if (ImageRecord->CodeSegmentCount == 0) {
+  ImageRecord = UefiImageLoaderGetImageRecord (ImageContext);
-    mMemoryAttributesTableEnable = FALSE;
+  if (ImageRecord == NULL) {
-    DEBUG ((DEBUG_ERROR, "!!!!!!!!  InsertImageRecord - CodeSegmentCount is 0  !!!!!!!!\n"));
+    return ;
-    if (PdbPointer != NULL) {
+  }
-      DEBUG ((DEBUG_ERROR, "!!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));
+
+  UefiImageDebugPrintSegments (ImageContext);
+  UefiImageDebugPrintImageRecord (ImageRecord);
+
+  //
+  // Section order is guaranteed by the PE specification.
+  // Section validity (e.g. no overlap) is guaranteed by the PE specification.
+  //
+
+  InsertSortImageRecord (ImageRecord);
+}
+
+/**
+  Find image record according to image base and size.
+
+  @param  ImageBase    Base of PE image
+  @param  ImageSize    Size of PE image
+
+  @return image record
+**/
+STATIC
+UEFI_IMAGE_RECORD *
+FindImageRecord (
+  IN EFI_PHYSICAL_ADDRESS  ImageBase
+  )
+{
+  UEFI_IMAGE_RECORD  *ImageRecord;
+  LIST_ENTRY         *ImageRecordLink;
+  LIST_ENTRY         *ImageRecordList;
+
+  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
+
+  for (ImageRecordLink = ImageRecordList->ForwardLink;
+       ImageRecordLink != ImageRecordList;
+       ImageRecordLink = ImageRecordLink->ForwardLink)
+  {
+    ImageRecord = CR (
+                    ImageRecordLink,
+                    UEFI_IMAGE_RECORD,
+                    Link,
+                    UEFI_IMAGE_RECORD_SIGNATURE
+                    );
+
+    if (ImageBase == ImageRecord->StartAddress)
+    {
+      return ImageRecord;
     }
-
-    Status = EFI_ABORTED;
-    goto Finish;
   }
 
-  //
+  return NULL;
-  // Check overlap all section in ImageBase/Size
-  //
-  if (!IsImageRecordCodeSectionValid (ImageRecord)) {
-    DEBUG ((DEBUG_ERROR, "IsImageRecordCodeSectionValid - FAIL\n"));
-    Status = EFI_ABORTED;
-    goto Finish;
-  }
-
-  InsertTailList (&mImagePropertiesPrivateData.ImageRecordList, &ImageRecord->Link);
-  mImagePropertiesPrivateData.ImageRecordCount++;
-
-  if (mImagePropertiesPrivateData.CodeSegmentCountMax < ImageRecord->CodeSegmentCount) {
-    mImagePropertiesPrivateData.CodeSegmentCountMax = ImageRecord->CodeSegmentCount;
-  }
-
-  SortImageRecord (&mImagePropertiesPrivateData.ImageRecordList);
-
-Finish:
-  if (EFI_ERROR (Status) && (ImageRecord != NULL)) {
-    DeleteImagePropertiesRecord (ImageRecord);
-  }
-
-  return;
 }
 
 /**
@@ -673,7 +751,7 @@ RemoveImageRecord (
   IN EFI_RUNTIME_IMAGE_ENTRY  *RuntimeImage
   )
 {
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
+  UEFI_IMAGE_RECORD  *ImageRecord;
 
   DEBUG ((DEBUG_VERBOSE, "RemoveImageRecord - 0x%x\n", RuntimeImage));
   DEBUG ((DEBUG_VERBOSE, "RemoveImageRecord - 0x%016lx - 0x%016lx\n", (EFI_PHYSICAL_ADDRESS)(UINTN)RuntimeImage->ImageBase, RuntimeImage->ImageSize));
@@ -683,7 +761,7 @@ RemoveImageRecord (
     return;
   }
 
-  ImageRecord = FindImageRecord ((EFI_PHYSICAL_ADDRESS)(UINTN)RuntimeImage->ImageBase, RuntimeImage->ImageSize, &mImagePropertiesPrivateData.ImageRecordList);
+  ImageRecord = FindImageRecord ((EFI_PHYSICAL_ADDRESS)(UINTN)RuntimeImage->ImageBase);
   if (ImageRecord == NULL) {
     DEBUG ((DEBUG_ERROR, "!!!!!!!! ImageRecord not found !!!!!!!!\n"));
     return;

MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c

@@ -40,8 +40,12 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Protocol/FirmwareVolume2.h>
 #include <Protocol/SimpleFileSystem.h>
 
+#include "Base.h"
 #include "DxeMain.h"
+#include "Library/UefiImageLib.h"
 #include "Mem/HeapGuard.h"
+#include "ProcessorBind.h"
+#include "Uefi/UefiMultiPhase.h"
 
 //
 // Image type definitions
@@ -213,113 +217,73 @@ SetUefiImageMemoryAttributes (
 **/
 VOID
 SetUefiImageProtectionAttributes (
-  IN IMAGE_PROPERTIES_RECORD  *ImageRecord
+  IN UEFI_IMAGE_RECORD  *ImageRecord
   )
 {
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;
+  UEFI_IMAGE_RECORD_SEGMENT  *ImageRecordSegment;
-  LIST_ENTRY                            *ImageRecordCodeSectionLink;
+  UINTN                      SectionAddress;
-  LIST_ENTRY                            *ImageRecordCodeSectionEndLink;
+  UINT32                     Index;
-  LIST_ENTRY                            *ImageRecordCodeSectionList;
-  UINT64                                CurrentBase;
-  UINT64                                ImageEnd;
 
-  ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;
+  SectionAddress = ImageRecord->StartAddress;
-
+  for (Index = 0; Index < ImageRecord->NumSegments; Index++) {
-  CurrentBase = ImageRecord->ImageBase;
+    ImageRecordSegment = &ImageRecord->Segments[Index];
-  ImageEnd    = ImageRecord->ImageBase + ImageRecord->ImageSize;
-
-  ImageRecordCodeSectionLink    = ImageRecordCodeSectionList->ForwardLink;
-  ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;
-  while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
-    ImageRecordCodeSection = CR (
-                               ImageRecordCodeSectionLink,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,
-                               Link,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
-                               );
-    ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
-
-    ASSERT (CurrentBase <= ImageRecordCodeSection->CodeSegmentBase);
-    if (CurrentBase < ImageRecordCodeSection->CodeSegmentBase) {
-      //
-      // DATA
-      //
-      SetUefiImageMemoryAttributes (
-        CurrentBase,
-        ImageRecordCodeSection->CodeSegmentBase - CurrentBase,
-        EFI_MEMORY_XP
-        );
-    }
-
-    //
-    // CODE
-    //
     SetUefiImageMemoryAttributes (
-      ImageRecordCodeSection->CodeSegmentBase,
+      SectionAddress,
-      ImageRecordCodeSection->CodeSegmentSize,
+      ImageRecordSegment->Size,
-      EFI_MEMORY_RO
+      ImageRecordSegment->Attributes
       );
-    CurrentBase = ImageRecordCodeSection->CodeSegmentBase + ImageRecordCodeSection->CodeSegmentSize;
-  }
 
-  //
+    SectionAddress += ImageRecordSegment->Size;
-  // Last DATA
-  //
-  ASSERT (CurrentBase <= ImageEnd);
-  if (CurrentBase < ImageEnd) {
-    //
-    // DATA
-    //
-    SetUefiImageMemoryAttributes (
-      CurrentBase,
-      ImageEnd - CurrentBase,
-      EFI_MEMORY_XP
-      );
   }
-
-  return;
 }
 
 /**
-  Return the section alignment requirement for the PE image section type.
+  Return if the PE image section is aligned.
 
-  @param[in]  MemoryType  PE/COFF image memory type
+  @param[in]  SectionAlignment    PE/COFF section alignment
-
+  @param[in]  MemoryType          PE/COFF image memory type
-  @retval     The required section alignment for this memory type
 
+  @retval TRUE  The PE image section is aligned.
+  @retval FALSE The PE image section is not aligned.
 **/
 STATIC
-UINT32
+BOOLEAN
-GetMemoryProtectionSectionAlignment (
+IsMemoryProtectionSectionAligned (
+  IN UINT32           SectionAlignment,
   IN EFI_MEMORY_TYPE  MemoryType
   )
 {
-  UINT32  SectionAlignment;
+  UINT32  PageAlignment;
 
   switch (MemoryType) {
     case EfiRuntimeServicesCode:
     case EfiACPIMemoryNVS:
     case EfiReservedMemoryType:
-      SectionAlignment = RUNTIME_PAGE_ALLOCATION_GRANULARITY;
+      PageAlignment = RUNTIME_PAGE_ALLOCATION_GRANULARITY;
       break;
     case EfiRuntimeServicesData:
       ASSERT (FALSE);
-      SectionAlignment = RUNTIME_PAGE_ALLOCATION_GRANULARITY;
+      PageAlignment = RUNTIME_PAGE_ALLOCATION_GRANULARITY;
       break;
     case EfiBootServicesCode:
     case EfiLoaderCode:
-      SectionAlignment = EFI_PAGE_SIZE;
+      PageAlignment = EFI_PAGE_SIZE;
       break;
     case EfiACPIReclaimMemory:
     default:
       ASSERT (FALSE);
-      SectionAlignment = EFI_PAGE_SIZE;
+      PageAlignment = EFI_PAGE_SIZE;
       break;
   }
 
-  return SectionAlignment;
+  if ((SectionAlignment & (PageAlignment - 1)) != 0) {
+    return FALSE;
+  } else {
+    return TRUE;
+  }
 }
 
+// FIXME: Deduplicate
 /**
   Protect UEFI PE/COFF image.
 
@@ -328,22 +292,26 @@ GetMemoryProtectionSectionAlignment (
 **/
 VOID
 ProtectUefiImage (
-  IN EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage,
+  IN LOADED_IMAGE_PRIVATE_DATA     *Image,
-  IN EFI_DEVICE_PATH_PROTOCOL   *LoadedImageDevicePath
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   )
 {
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
+  RETURN_STATUS              PdbStatus;
-  UINT32                   ProtectionPolicy;
+  EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage;
-  EFI_STATUS               Status;
+  EFI_DEVICE_PATH_PROTOCOL   *LoadedImageDevicePath;
-  UINT32                   RequiredAlignment;
+  UINT32                     SectionAlignment;
+  UEFI_IMAGE_RECORD          *ImageRecord;
+  CONST CHAR8                *PdbPointer;
+  UINT32                     PdbSize;
+  BOOLEAN                    IsAligned;
+  UINT32                     ProtectionPolicy;
+
+  LoadedImage = &Image->Info;
+  LoadedImageDevicePath = Image->LoadedImageDevicePath;
 
   DEBUG ((DEBUG_INFO, "ProtectUefiImageCommon - 0x%x\n", LoadedImage));
   DEBUG ((DEBUG_INFO, "  - 0x%016lx - 0x%016lx\n", (EFI_PHYSICAL_ADDRESS)(UINTN)LoadedImage->ImageBase, LoadedImage->ImageSize));
 
-  if (gCpu == NULL) {
-    return;
-  }
-
   ProtectionPolicy = GetUefiImageProtectionPolicy (LoadedImage, LoadedImageDevicePath);
   switch (ProtectionPolicy) {
     case DO_NOT_PROTECT:
@@ -355,36 +323,49 @@ ProtectUefiImage (
       return;
   }
 
-  ImageRecord = AllocateZeroPool (sizeof (*ImageRecord));
+  PdbStatus = UefiImageGetSymbolsPath (ImageContext, &PdbPointer, &PdbSize);
-  if (ImageRecord == NULL) {
+  if (!RETURN_ERROR (PdbStatus)) {
-    return;
+    DEBUG ((DEBUG_VERBOSE, "  Image - %a\n", PdbPointer));
   }
 
-  RequiredAlignment = GetMemoryProtectionSectionAlignment (LoadedImage->ImageCodeType);
+  //
+  // Get SectionAlignment
+  //
+  SectionAlignment = UefiImageGetSegmentAlignment (ImageContext);
 
-  Status = CreateImagePropertiesRecord (
+  IsAligned = IsMemoryProtectionSectionAligned (SectionAlignment, LoadedImage->ImageCodeType);
-             LoadedImage->ImageBase,
+  if (!IsAligned) {
-             LoadedImage->ImageSize,
+    DEBUG ((
-             &RequiredAlignment,
+      DEBUG_VERBOSE,
-             ImageRecord
+      "!!!!!!!!  ProtectUefiImageCommon - Section Alignment(0x%x) is incorrect  !!!!!!!!\n",
-             );
+      SectionAlignment));
+    if (!RETURN_ERROR (PdbStatus)) {
+      DEBUG ((DEBUG_VERBOSE, "!!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));
+    }
 
-  if (EFI_ERROR (Status)) {
-    DEBUG ((DEBUG_ERROR, "%a failed to create image properties record\n", __func__));
-    FreePool (ImageRecord);
     goto Finish;
   }
 
-  //
+  ImageRecord = UefiImageLoaderGetImageRecord (ImageContext);
-  // CPU ARCH present. Update memory attribute directly.
+  if (ImageRecord == NULL) {
-  //
+    return ;
-  SetUefiImageProtectionAttributes (ImageRecord);
+  }
+
+  UefiImageDebugPrintSegments (ImageContext);
+  UefiImageDebugPrintImageRecord (ImageRecord);
 
   //
   // Record the image record in the list so we can undo the protections later
   //
   InsertTailList (&mProtectedImageRecordList, &ImageRecord->Link);
 
+  if (gCpu != NULL) {
+    //
+    // CPU ARCH present. Update memory attribute directly.
+    //
+    SetUefiImageProtectionAttributes (ImageRecord);
+  }
+
 Finish:
   return;
 }
@@ -401,30 +382,31 @@ UnprotectUefiImage (
   IN EFI_DEVICE_PATH_PROTOCOL   *LoadedImageDevicePath
   )
 {
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
+  UEFI_IMAGE_RECORD  *ImageRecord;
-  LIST_ENTRY               *ImageRecordLink;
+  LIST_ENTRY         *ImageRecordLink;
 
-  if (PcdGet32 (PcdImageProtectionPolicy) != 0) {
+  for (ImageRecordLink = mProtectedImageRecordList.ForwardLink;
-    for (ImageRecordLink = mProtectedImageRecordList.ForwardLink;
+        ImageRecordLink != &mProtectedImageRecordList;
-         ImageRecordLink != &mProtectedImageRecordList;
+        ImageRecordLink = ImageRecordLink->ForwardLink) {
-         ImageRecordLink = ImageRecordLink->ForwardLink)
+    ImageRecord = CR (
-    {
+                    ImageRecordLink,
-      ImageRecord = CR (
+                    UEFI_IMAGE_RECORD,
-                      ImageRecordLink,
+                    Link,
-                      IMAGE_PROPERTIES_RECORD,
+                    UEFI_IMAGE_RECORD_SIGNATURE
-                      Link,
+                    );
-                      IMAGE_PROPERTIES_RECORD_SIGNATURE
-                      );
 
-      if (ImageRecord->ImageBase == (EFI_PHYSICAL_ADDRESS)(UINTN)LoadedImage->ImageBase) {
+    if (ImageRecord->StartAddress == (EFI_PHYSICAL_ADDRESS)(UINTN)LoadedImage->ImageBase) {
+      // TODO: Revise for removal (e.g. CpuDxe integration)
+      if (gCpu != NULL) {
         SetUefiImageMemoryAttributes (
-          ImageRecord->ImageBase,
+          ImageRecord->StartAddress,
-          ImageRecord->ImageSize,
+          ImageRecord->EndAddress - ImageRecord->StartAddress,
           0
           );
-        DeleteImagePropertiesRecord (ImageRecord);
-        return;
       }
+      RemoveEntryList (&ImageRecord->Link);
+      FreePool (ImageRecord);
+      return;
     }
   }
 }
@@ -567,7 +549,6 @@ MergeMemoryMapForProtectionPolicy (
   Remove exec permissions from all regions whose type is identified by
   PcdDxeNxMemoryProtectionPolicy.
 **/
-STATIC
 VOID
 InitializeDxeNxMemoryProtectionPolicy (
   VOID
@@ -775,12 +756,9 @@ MemoryProtectionCpuArchProtocolNotify (
   IN VOID       *Context
   )
 {
-  EFI_STATUS                 Status;
+  EFI_STATUS         Status;
-  EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage;
+  LIST_ENTRY         *ImageRecordLink;
-  EFI_DEVICE_PATH_PROTOCOL   *LoadedImageDevicePath;
+  UEFI_IMAGE_RECORD  *ImageRecord;
-  UINTN                      NoHandles;
-  EFI_HANDLE                 *HandleBuffer;
-  UINTN                      Index;
 
   DEBUG ((DEBUG_INFO, "MemoryProtectionCpuArchProtocolNotify:\n"));
   Status = CoreLocateProtocol (&gEfiCpuArchProtocolGuid, NULL, (VOID **)&gCpu);
@@ -804,41 +782,22 @@ MemoryProtectionCpuArchProtocolNotify (
     goto Done;
   }
 
-  Status = gBS->LocateHandleBuffer (
+  for (ImageRecordLink = mProtectedImageRecordList.ForwardLink;
-                  ByProtocol,
+        ImageRecordLink != &mProtectedImageRecordList;
-                  &gEfiLoadedImageProtocolGuid,
+        ImageRecordLink = ImageRecordLink->ForwardLink) {
-                  NULL,
+    ImageRecord = CR (
-                  &NoHandles,
+                    ImageRecordLink,
-                  &HandleBuffer
+                    UEFI_IMAGE_RECORD,
-                  );
+                    Link,
-  if (EFI_ERROR (Status) && (NoHandles == 0)) {
+                    UEFI_IMAGE_RECORD_SIGNATURE
-    goto Done;
-  }
-
-  for (Index = 0; Index < NoHandles; Index++) {
-    Status = gBS->HandleProtocol (
-                    HandleBuffer[Index],
-                    &gEfiLoadedImageProtocolGuid,
-                    (VOID **)&LoadedImage
                     );
-    if (EFI_ERROR (Status)) {
-      continue;
-    }
 
-    Status = gBS->HandleProtocol (
+    //
-                    HandleBuffer[Index],
+    // CPU ARCH present. Update memory attribute directly.
-                    &gEfiLoadedImageDevicePathProtocolGuid,
+    //
-                    (VOID **)&LoadedImageDevicePath
+    SetUefiImageProtectionAttributes (ImageRecord);
-                    );
-    if (EFI_ERROR (Status)) {
-      LoadedImageDevicePath = NULL;
-    }
-
-    ProtectUefiImage (LoadedImage, LoadedImageDevicePath);
   }
 
-  FreePool (HandleBuffer);
-
 Done:
   CoreCloseEvent (Event);
 }

MdeModulePkg/Core/Pei/Dispatcher/Dispatcher.c

@@ -9,6 +9,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 **/
 
 #include "PeiMain.h"
+#include "ProcessorBind.h"
 
 //
 // Utility global variables
@@ -1377,8 +1378,7 @@ MigratePeim (
   EFI_FFS_FILE_HEADER  *FileHeader;
   VOID                 *Pe32Data;
   VOID                 *ImageAddress;
-  CHAR8                *AsciiString;
+  UINT32                    ImageSize;
-  UINTN                Index;
 
   Status = EFI_SUCCESS;
 
@@ -1386,24 +1386,10 @@ MigratePeim (
   ASSERT (!IS_FFS_FILE2 (FileHeader));
 
   ImageAddress = NULL;
-  PeiGetPe32Data (MigratedFileHandle, &ImageAddress);
+  PeiGetPe32Data (MigratedFileHandle, &ImageAddress, &ImageSize);
   if (ImageAddress != NULL) {
-    DEBUG_CODE_BEGIN ();
-    AsciiString = PeCoffLoaderGetPdbPointer (ImageAddress);
-    for (Index = 0; AsciiString[Index] != 0; Index++) {
-      if ((AsciiString[Index] == '\\') || (AsciiString[Index] == '/')) {
-        AsciiString = AsciiString + Index + 1;
-        Index       = 0;
-      } else if (AsciiString[Index] == '.') {
-        AsciiString[Index] = 0;
-      }
-    }
-
-    DEBUG ((DEBUG_VERBOSE, "%a", AsciiString));
-    DEBUG_CODE_END ();
-
     Pe32Data = (VOID *)((UINTN)ImageAddress - (UINTN)MigratedFileHandle + (UINTN)FileHandle);
-    Status   = LoadAndRelocatePeCoffImageInPlace (Pe32Data, ImageAddress);
+    Status   = LoadAndRelocateUefiImageInPlace (Pe32Data, ImageAddress, ImageSize);
     ASSERT_EFI_ERROR (Status);
   }
 
@@ -1465,6 +1451,104 @@ ConvertStatusCodeCallbacks (
   }
 }
 
+/**
+  Migrates SEC modules in the given firmware volume.
+
+  Migrating SECURITY_CORE files requires special treatment since they are not tracked for PEI dispatch.
+
+  This functioun should be called after the FV has been copied to its post-memory location and the PEI Core FV list has
+  been updated.
+
+  @param Private          Pointer to the PeiCore's private data structure.
+  @param FvIndex          The firmware volume index to migrate.
+  @param OrgFvHandle      The handle to the firmware volume in temporary memory.
+
+  @retval   EFI_SUCCESS           SEC modules were migrated successfully
+  @retval   EFI_INVALID_PARAMETER The Private pointer is NULL or FvCount is invalid.
+  @retval   EFI_NOT_FOUND         Can't find valid FFS header.
+
+**/
+EFI_STATUS
+EFIAPI
+MigrateSecModulesInFv (
+  IN PEI_CORE_INSTANCE    *Private,
+  IN  UINTN               FvIndex,
+  IN  UINTN               OrgFvHandle
+  )
+{
+  EFI_STATUS                  Status;
+  EFI_STATUS                  FindFileStatus;
+  EFI_PEI_FILE_HANDLE         MigratedFileHandle;
+  EFI_PEI_FILE_HANDLE         FileHandle;
+  UINT32                      SectionAuthenticationStatus;
+  UINT32                      FileSize;
+  VOID                        *OrgPe32SectionData;
+  VOID                        *Pe32SectionData;
+  UINT32                      Pe32SectionDataSize;
+  EFI_FFS_FILE_HEADER         *FfsFileHeader;
+  EFI_COMMON_SECTION_HEADER   *Section;
+  BOOLEAN                     IsFfs3Fv;
+  UINTN                       SectionInstance;
+
+  if (Private == NULL || FvIndex >= Private->FvCount) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  do {
+    FindFileStatus =  PeiFfsFindNextFile (
+                        GetPeiServicesTablePointer (),
+                        EFI_FV_FILETYPE_SECURITY_CORE,
+                        Private->Fv[FvIndex].FvHandle,
+                        &MigratedFileHandle
+                        );
+    if (!EFI_ERROR (FindFileStatus ) && MigratedFileHandle != NULL) {
+      FileHandle = (EFI_PEI_FILE_HANDLE) ((UINTN) MigratedFileHandle - (UINTN) Private->Fv[FvIndex].FvHandle + OrgFvHandle);
+      FfsFileHeader = (EFI_FFS_FILE_HEADER *) MigratedFileHandle;
+
+      DEBUG ((DEBUG_VERBOSE, "    Migrating SEC_CORE MigratedFileHandle at 0x%x.\n", (UINTN) MigratedFileHandle));
+      DEBUG ((DEBUG_VERBOSE, "                       FileHandle at 0x%x.\n", (UINTN) FileHandle));
+
+      IsFfs3Fv = CompareGuid (&Private->Fv[FvIndex].FvHeader->FileSystemGuid, &gEfiFirmwareFileSystem3Guid);
+      if (IS_FFS_FILE2 (FfsFileHeader)) {
+        ASSERT (FFS_FILE2_SIZE (FfsFileHeader) > 0x00FFFFFF);
+        if (!IsFfs3Fv) {
+          DEBUG ((DEBUG_ERROR, "It is a FFS3 formatted file: %g in a non-FFS3 formatted FV.\n", &FfsFileHeader->Name));
+          return EFI_NOT_FOUND;
+        }
+        Section = (EFI_COMMON_SECTION_HEADER *) ((UINT8 *) FfsFileHeader + sizeof (EFI_FFS_FILE_HEADER2));
+        FileSize = FFS_FILE2_SIZE (FfsFileHeader) - sizeof (EFI_FFS_FILE_HEADER2);
+      } else {
+        Section = (EFI_COMMON_SECTION_HEADER *) ((UINT8 *) FfsFileHeader + sizeof (EFI_FFS_FILE_HEADER));
+        FileSize = FFS_FILE_SIZE (FfsFileHeader) - sizeof (EFI_FFS_FILE_HEADER);
+      }
+
+      SectionInstance = 1;
+      SectionAuthenticationStatus = 0;
+      Status = ProcessSection (
+                GetPeiServicesTablePointer (),
+                EFI_SECTION_PE32,
+                &SectionInstance,
+                Section,
+                FileSize,
+                &Pe32SectionData,
+                &Pe32SectionDataSize,
+                &SectionAuthenticationStatus,
+                IsFfs3Fv
+                );
+
+      if (!EFI_ERROR (Status)) {
+        OrgPe32SectionData = (VOID *) ((UINTN) Pe32SectionData - (UINTN) MigratedFileHandle + (UINTN) FileHandle);
+        DEBUG ((DEBUG_VERBOSE, "      PE32 section in migrated file at 0x%x.\n", (UINTN) Pe32SectionData));
+        DEBUG ((DEBUG_VERBOSE, "      PE32 section in original file at 0x%x.\n", (UINTN) OrgPe32SectionData));
+        Status = LoadAndRelocateUefiImageInPlace (OrgPe32SectionData, Pe32SectionData, Pe32SectionDataSize);
+        ASSERT_EFI_ERROR (Status);
+      }
+    }
+  } while (!EFI_ERROR (FindFileStatus));
+
+  return EFI_SUCCESS;
+}
+
 /**
   Migrates PEIMs in the given firmware volume.
 

MdeModulePkg/Core/Pei/FwVol/FwVol.c

@@ -35,7 +35,8 @@ PEI_FW_VOL_INSTANCE  mPeiFfs2FwVol = {
     PeiFfsFvPpiGetFileInfo2,
     PeiFfsFvPpiFindSectionByType2,
     EFI_PEI_FIRMWARE_VOLUME_PPI_SIGNATURE,
-    EFI_PEI_FIRMWARE_VOLUME_PPI_REVISION
+    EFI_PEI_FIRMWARE_VOLUME_PPI_REVISION,
+    PeiFfsFvPpiFindSectionByType3
   }
 };
 
@@ -52,7 +53,8 @@ PEI_FW_VOL_INSTANCE  mPeiFfs3FwVol = {
     PeiFfsFvPpiGetFileInfo2,
     PeiFfsFvPpiFindSectionByType2,
     EFI_PEI_FIRMWARE_VOLUME_PPI_SIGNATURE,
-    EFI_PEI_FIRMWARE_VOLUME_PPI_REVISION
+    EFI_PEI_FIRMWARE_VOLUME_PPI_REVISION,
+    PeiFfsFvPpiFindSectionByType3
   }
 };
 
@@ -771,6 +773,8 @@ VerifyGuidedSectionGuid (
   @param SectionSize       The file size to search.
   @param OutputBuffer      A pointer to the discovered section, if successful.
                            NULL if section not found
+  @param OutputSize        The size of the discovered section, if successful.
+                           0 if section not found
   @param AuthenticationStatus Updated upon return to point to the authentication status for this section.
   @param IsFfs3Fv          Indicates the FV format.
 
@@ -786,6 +790,7 @@ ProcessSection (
   IN EFI_COMMON_SECTION_HEADER  *Section,
   IN UINTN                      SectionSize,
   OUT VOID                      **OutputBuffer,
+  OUT UINT32                    *OutputSize,
   OUT UINT32                    *AuthenticationStatus,
   IN BOOLEAN                    IsFfs3Fv
   )
@@ -803,11 +808,13 @@ ProcessSection (
   EFI_GUID                               *SectionDefinitionGuid;
   BOOLEAN                                SectionCached;
   VOID                                   *TempOutputBuffer;
+  UINT32                                  TempOutputSize;
   UINT32                                 TempAuthenticationStatus;
   UINT16                                 GuidedSectionAttributes;
 
   PrivateData   = PEI_CORE_INSTANCE_FROM_PS_THIS (PeiServices);
   *OutputBuffer = NULL;
+  *OutputSize   = 0;
   ParsedLength  = 0;
   Index         = 0;
   Status        = EFI_NOT_FOUND;
@@ -844,10 +851,13 @@ ProcessSection (
         //
         // Got it!
         //
+        // FIXME: Size checks
         if (IS_SECTION2 (Section)) {
           *OutputBuffer = (VOID *)((UINT8 *)Section + sizeof (EFI_COMMON_SECTION_HEADER2));
+          *OutputSize   = SECTION2_SIZE (Section) - sizeof (EFI_COMMON_SECTION_HEADER2);
         } else {
           *OutputBuffer = (VOID *)((UINT8 *)Section + sizeof (EFI_COMMON_SECTION_HEADER));
+          *OutputSize   = SECTION_SIZE (Section) - sizeof (EFI_COMMON_SECTION_HEADER);
         }
 
         return EFI_SUCCESS;
@@ -894,11 +904,13 @@ ProcessSection (
                                        PpiOutput,
                                        PpiOutputSize,
                                        &TempOutputBuffer,
+                     &TempOutputSize,
                                        &TempAuthenticationStatus,
                                        IsFfs3Fv
                                        );
           if (!EFI_ERROR (Status)) {
             *OutputBuffer         = TempOutputBuffer;
+            *OutputSize   = TempOutputSize;
             *AuthenticationStatus = TempAuthenticationStatus | Authentication;
             return EFI_SUCCESS;
           }
@@ -981,11 +993,13 @@ ProcessSection (
                                        PpiOutput,
                                        PpiOutputSize,
                                        &TempOutputBuffer,
+                     &TempOutputSize,
                                        &TempAuthenticationStatus,
                                        IsFfs3Fv
                                        );
           if (!EFI_ERROR (Status)) {
             *OutputBuffer         = TempOutputBuffer;
+            *OutputSize   = TempOutputSize;
             *AuthenticationStatus = TempAuthenticationStatus | Authentication;
             return EFI_SUCCESS;
           }
@@ -1072,6 +1086,38 @@ PeiFfsFindSectionData3 (
   OUT VOID                    **SectionData,
   OUT UINT32                  *AuthenticationStatus
   )
+{
+  UINT32 SectionDataSize;
+
+  return PeiFfsFindSectionData4 (PeiServices, SectionType, SectionInstance, FileHandle, SectionData, &SectionDataSize, AuthenticationStatus);
+}
+
+/**
+  Searches for the next matching section within the specified file.
+
+  @param  PeiServices           An indirect pointer to the EFI_PEI_SERVICES table published by the PEI Foundation.
+  @param  SectionType           The value of the section type to find.
+  @param  SectionInstance       Section instance to find.
+  @param  FileHandle            Handle of the firmware file to search.
+  @param  SectionData           A pointer to the discovered section, if successful.
+  @param  SectionDataSize       The size of the discovered section, if successful.
+  @param  AuthenticationStatus  A pointer to the authentication status for this section.
+
+  @retval EFI_SUCCESS      The section was found.
+  @retval EFI_NOT_FOUND    The section was not found.
+
+**/
+EFI_STATUS
+EFIAPI
+PeiFfsFindSectionData4 (
+  IN CONST EFI_PEI_SERVICES    **PeiServices,
+  IN     EFI_SECTION_TYPE      SectionType,
+  IN     UINTN                 SectionInstance,
+  IN     EFI_PEI_FILE_HANDLE   FileHandle,
+  OUT VOID                     **SectionData,
+  OUT UINT32                   *SectionDataSize,
+  OUT UINT32                   *AuthenticationStatus
+  )
 {
   PEI_CORE_FV_HANDLE  *CoreFvHandle;
 
@@ -1083,7 +1129,7 @@ PeiFfsFindSectionData3 (
   if ((CoreFvHandle->FvPpi->Signature == EFI_PEI_FIRMWARE_VOLUME_PPI_SIGNATURE) &&
       (CoreFvHandle->FvPpi->Revision == EFI_PEI_FIRMWARE_VOLUME_PPI_REVISION))
   {
-    return CoreFvHandle->FvPpi->FindSectionByType2 (CoreFvHandle->FvPpi, SectionType, SectionInstance, FileHandle, SectionData, AuthenticationStatus);
+    return CoreFvHandle->FvPpi->FindSectionByType3 (CoreFvHandle->FvPpi, SectionType, SectionInstance, FileHandle, SectionData, SectionDataSize, AuthenticationStatus);
   }
 
   //
@@ -2040,6 +2086,57 @@ PeiFfsFvPpiFindSectionByType2 (
   OUT VOID                               **SectionData,
   OUT UINT32                             *AuthenticationStatus
   )
+{
+  UINT32 SectionDataSize;
+
+  // FIXME: Deprecate
+  return PeiFfsFvPpiFindSectionByType3 (
+           This,
+           SearchType,
+           SearchInstance,
+           FileHandle,
+           SectionData,
+           &SectionDataSize,
+           AuthenticationStatus
+           );
+}
+
+/**
+  Find the next matching section in the firmware file.
+
+  This service enables PEI modules to discover sections
+  of a given instance and type within a valid file.
+
+  @param This                   Points to this instance of the
+                                EFI_PEI_FIRMWARE_VOLUME_PPI.
+  @param SearchType             A filter to find only sections of this
+                                type.
+  @param SearchInstance         A filter to find the specific instance
+                                of sections.
+  @param FileHandle             Handle of firmware file in which to
+                                search.
+  @param SectionData            Updated upon return to point to the
+                                section found.
+  @param SectionDataSize        Updated upon return to point to the
+                                section size found.
+  @param AuthenticationStatus   Updated upon return to point to the
+                                authentication status for this section.
+
+  @retval EFI_SUCCESS     Section was found.
+  @retval EFI_NOT_FOUND   Section of the specified type was not
+                          found. SectionData contains NULL.
+**/
+EFI_STATUS
+EFIAPI
+PeiFfsFvPpiFindSectionByType3 (
+  IN  CONST EFI_PEI_FIRMWARE_VOLUME_PPI    *This,
+  IN        EFI_SECTION_TYPE               SearchType,
+  IN        UINTN                          SearchInstance,
+  IN        EFI_PEI_FILE_HANDLE            FileHandle,
+  OUT VOID                                 **SectionData,
+  OUT UINT32                               *SectionDataSize,
+  OUT UINT32                               *AuthenticationStatus
+  )
 {
   EFI_STATUS                 Status;
   EFI_FFS_FILE_HEADER        *FfsFileHeader;
@@ -2089,6 +2186,7 @@ PeiFfsFvPpiFindSectionByType2 (
                                     Section,
                                     FileSize,
                                     SectionData,
+             SectionDataSize,
                                     &ExtractedAuthenticationStatus,
                                     FwVolInstance->IsFfs3Fv
                                     );

MdeModulePkg/Core/Pei/FwVol/FwVol.h

@@ -185,6 +185,43 @@ PeiFfsFvPpiFindSectionByType2 (
   OUT UINT32                             *AuthenticationStatus
   );
 
+/**
+  Find the next matching section in the firmware file.
+
+  This service enables PEI modules to discover sections
+  of a given instance and type within a valid file.
+
+  @param This                   Points to this instance of the
+                                EFI_PEI_FIRMWARE_VOLUME_PPI.
+  @param SearchType             A filter to find only sections of this
+                                type.
+  @param SearchInstance         A filter to find the specific instance
+                                of sections.
+  @param FileHandle             Handle of firmware file in which to
+                                search.
+  @param SectionData            Updated upon return to point to the
+                                section found.
+  @param SectionDataSize        Updated upon return to point to the
+                                section size found.
+  @param AuthenticationStatus   Updated upon return to point to the
+                                authentication status for this section.
+
+  @retval EFI_SUCCESS     Section was found.
+  @retval EFI_NOT_FOUND   Section of the specified type was not
+                          found. SectionData contains NULL.
+**/
+EFI_STATUS
+EFIAPI
+PeiFfsFvPpiFindSectionByType3 (
+  IN  CONST EFI_PEI_FIRMWARE_VOLUME_PPI    *This,
+  IN        EFI_SECTION_TYPE               SearchType,
+  IN        UINTN                          SearchInstance,
+  IN        EFI_PEI_FILE_HANDLE            FileHandle,
+  OUT VOID                                 **SectionData,
+  OUT UINT32                               *SectionDataSize,
+  OUT UINT32                               *AuthenticationStatus
+  );
+
 /**
   Returns information about a specific file.
 

MdeModulePkg/Core/Pei/Image/Image.c

@@ -18,40 +18,6 @@ EFI_PEI_PPI_DESCRIPTOR  gPpiLoadFilePpiList = {
   &mPeiLoadImagePpi
 };
 
-/**
-
-  Support routine for the PE/COFF Loader that reads a buffer from a PE/COFF file.
-  The function is used for XIP code to have optimized memory copy.
-
-  @param FileHandle      - The handle to the PE/COFF file
-  @param FileOffset      - The offset, in bytes, into the file to read
-  @param ReadSize        - The number of bytes to read from the file starting at FileOffset
-  @param Buffer          - A pointer to the buffer to read the data into.
-
-  @return EFI_SUCCESS - ReadSize bytes of data were read into Buffer from the PE/COFF file starting at FileOffset
-
-**/
-EFI_STATUS
-EFIAPI
-PeiImageRead (
-  IN     VOID   *FileHandle,
-  IN     UINTN  FileOffset,
-  IN     UINTN  *ReadSize,
-  OUT    VOID   *Buffer
-  )
-{
-  CHAR8  *Destination8;
-  CHAR8  *Source8;
-
-  Destination8 = Buffer;
-  Source8      = (CHAR8 *)((UINTN)FileHandle + FileOffset);
-  if (Destination8 != Source8) {
-    CopyMem (Destination8, Source8, *ReadSize);
-  }
-
-  return EFI_SUCCESS;
-}
-
 /**
   To check memory usage bit map array to figure out if the memory range the image will be loaded in is available or not. If
   memory range is available, the function will mark the corresponding bits to 1 which indicates the memory range is used.
@@ -135,111 +101,36 @@ CheckAndMarkFixLoadingMemoryUsageBitMap (
 
 **/
 EFI_STATUS
-GetPeCoffImageFixLoadingAssignedAddress (
+GetUefiImageFixLoadingAssignedAddress (
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext,
+  OUT EFI_PHYSICAL_ADDRESS  *LoadAddress,
-  IN     PEI_CORE_INSTANCE             *Private
+  IN  UINT64                ValueInSectionHeader,
+  IN  UINT32                ImageDestSize,
+  IN  PEI_CORE_INSTANCE     *Private
   )
 {
-  UINTN                            SectionHeaderOffset;
+   EFI_STATUS           Status;
-  EFI_STATUS                       Status;
+   EFI_PHYSICAL_ADDRESS FixLoadingAddress;
-  EFI_IMAGE_SECTION_HEADER         SectionHeader;
-  EFI_IMAGE_OPTIONAL_HEADER_UNION  *ImgHdr;
-  EFI_PHYSICAL_ADDRESS             FixLoadingAddress;
-  UINT16                           Index;
-  UINTN                            Size;
-  UINT16                           NumberOfSections;
-  UINT64                           ValueInSectionHeader;
 
-  FixLoadingAddress = 0;
+  if ((INT64)PcdGet64(PcdLoadModuleAtFixAddressEnable) > 0) {
-  Status            = EFI_NOT_FOUND;
-
-  //
-  // Get PeHeader pointer
-  //
-  ImgHdr = (EFI_IMAGE_OPTIONAL_HEADER_UNION *)((CHAR8 *)ImageContext->Handle + ImageContext->PeCoffHeaderOffset);
-  if (ImageContext->IsTeImage) {
     //
-    // for TE image, the fix loading address is saved in first section header that doesn't point
+    // When LMFA feature is configured as Load Module at Fixed Absolute Address mode, PointerToRelocations & PointerToLineNumbers field
-    // to code section.
+    // hold the absolute address of image base running in memory
     //
-    SectionHeaderOffset = sizeof (EFI_TE_IMAGE_HEADER);
+    FixLoadingAddress = ValueInSectionHeader;
-    NumberOfSections    = ImgHdr->Te.NumberOfSections;
   } else {
-    SectionHeaderOffset = ImageContext->PeCoffHeaderOffset +
-                          sizeof (UINT32) +
-                          sizeof (EFI_IMAGE_FILE_HEADER) +
-                          ImgHdr->Pe32.FileHeader.SizeOfOptionalHeader;
-    NumberOfSections = ImgHdr->Pe32.FileHeader.NumberOfSections;
-  }
-
-  //
-  // Get base address from the first section header that doesn't point to code section.
-  //
-  for (Index = 0; Index < NumberOfSections; Index++) {
     //
-    // Read section header from file
+    // When LMFA feature is configured as Load Module at Fixed offset mode, PointerToRelocations & PointerToLineNumbers field
+    // hold the offset relative to a platform-specific top address.
     //
-    Size   = sizeof (EFI_IMAGE_SECTION_HEADER);
+    FixLoadingAddress = (EFI_PHYSICAL_ADDRESS)(Private->LoadModuleAtFixAddressTopAddress + ValueInSectionHeader);
-    Status = ImageContext->ImageRead (
-                             ImageContext->Handle,
-                             SectionHeaderOffset,
-                             &Size,
-                             &SectionHeader
-                             );
-    if (EFI_ERROR (Status)) {
-      return Status;
-    }
-
-    Status = EFI_NOT_FOUND;
-
-    if ((SectionHeader.Characteristics & EFI_IMAGE_SCN_CNT_CODE) == 0) {
-      //
-      // Build tool will save the address in PointerToRelocations & PointerToLineNumbers fields in the first section header
-      // that doesn't point to code section in image header, as well as ImageBase field of image header. A notable thing is
-      // that for PEIM, the value in ImageBase field may not be equal to the value in PointerToRelocations & PointerToLineNumbers because
-      // for XIP PEIM, ImageBase field holds the image base address running on the Flash. And PointerToRelocations & PointerToLineNumbers
-      // hold the image base address when it is shadow to the memory. And there is an assumption that when the feature is enabled, if a
-      // module is assigned a loading address by tools, PointerToRelocations & PointerToLineNumbers fields should NOT be Zero, or
-      // else, these 2 fields should be set to Zero
-      //
-      ValueInSectionHeader = ReadUnaligned64 ((UINT64 *)&SectionHeader.PointerToRelocations);
-      if (ValueInSectionHeader != 0) {
-        //
-        // Found first section header that doesn't point to code section.
-        //
-        if ((INT64)PcdGet64 (PcdLoadModuleAtFixAddressEnable) > 0) {
-          //
-          // When LMFA feature is configured as Load Module at Fixed Absolute Address mode, PointerToRelocations & PointerToLineNumbers field
-          // hold the absolute address of image base running in memory
-          //
-          FixLoadingAddress = ValueInSectionHeader;
-        } else {
-          //
-          // When LMFA feature is configured as Load Module at Fixed offset mode, PointerToRelocations & PointerToLineNumbers field
-          // hold the offset relative to a platform-specific top address.
-          //
-          FixLoadingAddress = (EFI_PHYSICAL_ADDRESS)(Private->LoadModuleAtFixAddressTopAddress + (INT64)ValueInSectionHeader);
-        }
-
-        //
-        // Check if the memory range is available.
-        //
-        Status = CheckAndMarkFixLoadingMemoryUsageBitMap (Private, FixLoadingAddress, (UINT32)ImageContext->ImageSize);
-        if (!EFI_ERROR (Status)) {
-          //
-          // The assigned address is valid. Return the specified loading address
-          //
-          ImageContext->ImageAddress = FixLoadingAddress;
-        }
-      }
-
-      break;
-    }
-
-    SectionHeaderOffset += sizeof (EFI_IMAGE_SECTION_HEADER);
   }
+  //
+  // Check if the memory range is available.
+  //
+  Status = CheckAndMarkFixLoadingMemoryUsageBitMap (Private, FixLoadingAddress, ImageDestSize);
+  *LoadAddress = FixLoadingAddress;
 
-  DEBUG ((DEBUG_INFO|DEBUG_LOAD, "LOADING MODULE FIXED INFO: Loading module at fixed address 0x%11p. Status= %r \n", (VOID *)(UINTN)FixLoadingAddress, Status));
+  DEBUG ((EFI_D_INFO|EFI_D_LOAD, "LOADING MODULE FIXED INFO: Loading module at fixed address 0x%11p. Status= %r \n", (VOID *)(UINTN)FixLoadingAddress, Status));
   return Status;
 }
 
@@ -262,38 +153,49 @@ GetPeCoffImageFixLoadingAssignedAddress (
 
 **/
 EFI_STATUS
-LoadAndRelocatePeCoffImage (
+LoadAndRelocateUefiImage (
   IN  EFI_PEI_FILE_HANDLE   FileHandle,
   IN  VOID                  *Pe32Data,
-  OUT EFI_PHYSICAL_ADDRESS  *ImageAddress,
+  IN  UINT32                                    Pe32DataSize,
-  OUT UINT64                *ImageSize,
+  OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT           *ImageContext,
-  OUT EFI_PHYSICAL_ADDRESS  *EntryPoint
+  OUT EFI_PHYSICAL_ADDRESS                      *ImageAddress
   )
 {
   EFI_STATUS                    Status;
-  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;
+  BOOLEAN                       Success;
   PEI_CORE_INSTANCE             *Private;
-  UINT64                        AlignImageSize;
+  UINT32                                ImageSize;
+  UINT32                                ImageAlignment;
+  UINT64                                ValueInSectionHeader;
   BOOLEAN                       IsXipImage;
   EFI_STATUS                    ReturnStatus;
   BOOLEAN                       IsS3Boot;
   BOOLEAN                       IsPeiModule;
   BOOLEAN                       IsRegisterForShadow;
   EFI_FV_FILE_INFO              FileInfo;
+  UINT32                                DestinationPages;
+  UINT32                                DestinationSize;
+  EFI_PHYSICAL_ADDRESS                  Destination;
+  UINT16                                Machine;
+  BOOLEAN                               LoadDynamically;
 
   Private = PEI_CORE_INSTANCE_FROM_PS_THIS (GetPeiServicesTablePointer ());
 
   ReturnStatus = EFI_SUCCESS;
-  IsXipImage   = FALSE;
-  ZeroMem (&ImageContext, sizeof (ImageContext));
-  ImageContext.Handle    = Pe32Data;
-  ImageContext.ImageRead = PeiImageRead;
 
-  Status = PeCoffLoaderGetImageInfo (&ImageContext);
+  Status = UefiImageInitializeContext (ImageContext, Pe32Data, Pe32DataSize);
   if (EFI_ERROR (Status)) {
     return Status;
   }
 
+  Machine = UefiImageGetMachine (ImageContext);
+
+  if (!EFI_IMAGE_MACHINE_TYPE_SUPPORTED (Machine)) {
+    if (!EFI_IMAGE_MACHINE_CROSS_TYPE_SUPPORTED (Machine)) {
+      return EFI_UNSUPPORTED;
+    }
+  }
+
   //
   // Initialize local IsS3Boot and IsRegisterForShadow variable
   //
@@ -312,9 +214,7 @@ LoadAndRelocatePeCoffImage (
   //
   // XIP image that ImageAddress is same to Image handle.
   //
-  if (ImageContext.ImageAddress == (EFI_PHYSICAL_ADDRESS)(UINTN)Pe32Data) {
+  IsXipImage = UefiImageImageIsInplace (ImageContext);
-    IsXipImage = TRUE;
-  }
 
   //
   // Get file type first
@@ -336,7 +236,7 @@ LoadAndRelocatePeCoffImage (
   //
   // When Image has no reloc section, it can't be relocated into memory.
   //
-  if (ImageContext.RelocationsStripped && (Private->PeiMemoryInstalled) &&
+  if (UefiImageGetRelocsStripped (ImageContext) && (Private->PeiMemoryInstalled) &&
       ((!IsPeiModule) || PcdGetBool (PcdMigrateTemporaryRamFirmwareVolumes) ||
        (!IsS3Boot && (PcdGetBool (PcdShadowPeimOnBoot) || IsRegisterForShadow)) ||
        (IsS3Boot && PcdGetBool (PcdShadowPeimOnS3Boot)))
@@ -345,125 +245,95 @@ LoadAndRelocatePeCoffImage (
     DEBUG ((DEBUG_INFO|DEBUG_LOAD, "The image at 0x%08x without reloc section can't be loaded into memory\n", (UINTN)Pe32Data));
   }
 
-  //
+  LoadDynamically  = FALSE;
-  // Set default base address to current image address.
+  ImageSize        = UefiImageGetImageSize (ImageContext);
-  //
+  DestinationPages = EFI_SIZE_TO_PAGES (ImageSize);
-  ImageContext.ImageAddress = (EFI_PHYSICAL_ADDRESS)(UINTN)Pe32Data;
+  DestinationSize  = EFI_PAGES_TO_SIZE (DestinationPages);
 
   //
   // Allocate Memory for the image when memory is ready, and image is relocatable.
   // On normal boot, PcdShadowPeimOnBoot decides whether load PEIM or PeiCore into memory.
   // On S3 boot, PcdShadowPeimOnS3Boot decides whether load PEIM or PeiCore into memory.
   //
-  if ((!ImageContext.RelocationsStripped) && (Private->PeiMemoryInstalled) &&
+  if ((!UefiImageGetRelocsStripped (ImageContext)) && (Private->PeiMemoryInstalled) &&
       ((!IsPeiModule) || PcdGetBool (PcdMigrateTemporaryRamFirmwareVolumes) ||
        (!IsS3Boot && (PcdGetBool (PcdShadowPeimOnBoot) || IsRegisterForShadow)) ||
        (IsS3Boot && PcdGetBool (PcdShadowPeimOnS3Boot)))
       )
   {
-    //
+    Success = FALSE;
-    // Allocate more buffer to avoid buffer overflow.
+
-    //
+    if (PcdGet64(PcdLoadModuleAtFixAddressEnable) != 0 && (Private->HobList.HandoffInformationTable->BootMode != BOOT_ON_S3_RESUME)) {
-    if (ImageContext.IsTeImage) {
+      Status = UefiImageGetFixedAddress (ImageContext, &ValueInSectionHeader);
-      AlignImageSize = ImageContext.ImageSize + ((EFI_TE_IMAGE_HEADER *)Pe32Data)->StrippedSize - sizeof (EFI_TE_IMAGE_HEADER);
+      if (!RETURN_ERROR (Status)) {
-    } else {
+        Status = GetUefiImageFixLoadingAssignedAddress(&Destination, ValueInSectionHeader, DestinationSize, Private);
-      AlignImageSize = ImageContext.ImageSize;
+      }
+
+      if (!EFI_ERROR (Status)){
+        Success = Destination == UefiImageGetPreferredAddress (ImageContext);
+
+        if (!Success) {
+          DEBUG ((DEBUG_INFO|DEBUG_LOAD, "LOADING MODULE FIXED ERROR: Loading module at fixed address failed since relocs have been stripped.\n"));
+        }
+      } else {
+        DEBUG ((EFI_D_INFO|EFI_D_LOAD, "LOADING MODULE FIXED ERROR: Failed to load module at fixed address. \n"));
+      }
     }
 
-    if (ImageContext.SectionAlignment > EFI_PAGE_SIZE) {
+    if (!Success) {
-      AlignImageSize += ImageContext.SectionAlignment;
+      //
+      // Allocate more buffer to avoid buffer overflow.
+      //
+      ImageAlignment = UefiImageGetSegmentAlignment (ImageContext);
+
+      Destination = (UINTN)AllocateAlignedCodePages (
+                             DestinationPages,
+                             ImageAlignment
+                             );
+      Success = Destination != 0;
     }
 
-    if ((PcdGet64 (PcdLoadModuleAtFixAddressEnable) != 0) && (Private->HobList.HandoffInformationTable->BootMode != BOOT_ON_S3_RESUME)) {
+    if (Success) {
-      Status = GetPeCoffImageFixLoadingAssignedAddress (&ImageContext, Private);
+      LoadDynamically = TRUE;
+      //
+      // Load the image to our new buffer
+      //
+      Status = UefiImageLoadImageForExecution (
+                ImageContext,
+                (VOID *) (UINTN)Destination,
+                DestinationSize,
+                NULL,
+                0
+                );
       if (EFI_ERROR (Status)) {
-        DEBUG ((DEBUG_INFO|DEBUG_LOAD, "LOADING MODULE FIXED ERROR: Failed to load module at fixed address. \n"));
+        return              Status;
-        //
-        // The PEIM is not assigned valid address, try to allocate page to load it.
-        //
-        Status = PeiServicesAllocatePages (
-                   EfiBootServicesCode,
-                   EFI_SIZE_TO_PAGES ((UINT32)AlignImageSize),
-                   &ImageContext.ImageAddress
-                   );
-      }
-    } else {
-      Status = PeiServicesAllocatePages (
-                 EfiBootServicesCode,
-                 EFI_SIZE_TO_PAGES ((UINT32)AlignImageSize),
-                 &ImageContext.ImageAddress
-                 );
-    }
-
-    if (!EFI_ERROR (Status)) {
-      //
-      // Adjust the Image Address to make sure it is section alignment.
-      //
-      if (ImageContext.SectionAlignment > EFI_PAGE_SIZE) {
-        ImageContext.ImageAddress =
-          (ImageContext.ImageAddress + ImageContext.SectionAlignment - 1) &
-          ~((UINTN)ImageContext.SectionAlignment - 1);
-      }
-
-      //
-      // Fix alignment requirement when Load IPF TeImage into memory.
-      // Skip the reserved space for the stripped PeHeader when load TeImage into memory.
-      //
-      if (ImageContext.IsTeImage) {
-        ImageContext.ImageAddress = ImageContext.ImageAddress +
-                                    ((EFI_TE_IMAGE_HEADER *)Pe32Data)->StrippedSize -
-                                    sizeof (EFI_TE_IMAGE_HEADER);
       }
     } else {
       //
       // No enough memory resource.
       //
-      if (IsXipImage) {
+      if (!IsXipImage) {
-        //
-        // XIP image can still be invoked.
-        //
-        ImageContext.ImageAddress = (EFI_PHYSICAL_ADDRESS)(UINTN)Pe32Data;
-        ReturnStatus              = EFI_WARN_BUFFER_TOO_SMALL;
-      } else {
         //
         // Non XIP image can't be loaded because no enough memory is allocated.
         //
         ASSERT (FALSE);
         return EFI_OUT_OF_RESOURCES;
       }
+      //
+      // XIP image can still be invoked.
+      //
+      ReturnStatus = EFI_WARN_BUFFER_TOO_SMALL;
     }
   }
 
-  //
+  if (!LoadDynamically) {
-  // Load the image to our new buffer
+    Status = UefiImageLoadImageInplace (ImageContext);
-  //
+    if (EFI_ERROR (Status)) {
-  Status = PeCoffLoaderLoadImage (&ImageContext);
+      return Status;
-  if (EFI_ERROR (Status)) {
-    if (ImageContext.ImageError == IMAGE_ERROR_INVALID_SECTION_ALIGNMENT) {
-      DEBUG ((DEBUG_ERROR, "PEIM Image Address 0x%11p doesn't meet with section alignment 0x%x.\n", (VOID *)(UINTN)ImageContext.ImageAddress, ImageContext.SectionAlignment));
     }
-
-    return Status;
   }
 
-  //
+  *ImageAddress = UefiImageLoaderGetImageAddress (ImageContext);
-  // Relocate the image in our new buffer
-  //
-  Status = PeCoffLoaderRelocateImage (&ImageContext);
-  if (EFI_ERROR (Status)) {
-    return Status;
-  }
-
-  //
-  // Flush the instruction cache so the image data is written before we execute it
-  //
-  if (ImageContext.ImageAddress != (EFI_PHYSICAL_ADDRESS)(UINTN)Pe32Data) {
-    InvalidateInstructionCacheRange ((VOID *)(UINTN)ImageContext.ImageAddress, (UINTN)ImageContext.ImageSize);
-  }
-
-  *ImageAddress = ImageContext.ImageAddress;
-  *ImageSize    = ImageContext.ImageSize;
-  *EntryPoint   = ImageContext.EntryPoint;
 
   return ReturnStatus;
 }
@@ -479,50 +349,30 @@ LoadAndRelocatePeCoffImage (
 
 **/
 EFI_STATUS
-LoadAndRelocatePeCoffImageInPlace (
+LoadAndRelocateUefiImageInPlace (
   IN  VOID  *Pe32Data,
-  IN  VOID  *ImageAddress
+  IN  VOID    *ImageAddress,
+  IN  UINT32  ImageSize
   )
 {
   EFI_STATUS                    Status;
-  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT ImageContext;
 
-  ZeroMem (&ImageContext, sizeof (ImageContext));
+  ASSERT (Pe32Data != ImageAddress);
-  ImageContext.Handle    = Pe32Data;
-  ImageContext.ImageRead = PeiImageRead;
 
-  Status = PeCoffLoaderGetImageInfo (&ImageContext);
+  CopyMem (ImageAddress, Pe32Data, ImageSize);
+
+  Status = UefiImageInitializeContext (&ImageContext, ImageAddress, ImageSize);
   if (EFI_ERROR (Status)) {
     ASSERT_EFI_ERROR (Status);
     return Status;
   }
 
-  ImageContext.ImageAddress = (PHYSICAL_ADDRESS)(UINTN)ImageAddress;
-
   //
   // Load the image in place
   //
-  Status = PeCoffLoaderLoadImage (&ImageContext);
+  Status = UefiImageRelocateImageInplaceForExecution (&ImageContext);
-  if (EFI_ERROR (Status)) {
+  ASSERT_EFI_ERROR (Status);
-    ASSERT_EFI_ERROR (Status);
-    return Status;
-  }
-
-  //
-  // Relocate the image in place
-  //
-  Status = PeCoffLoaderRelocateImage (&ImageContext);
-  if (EFI_ERROR (Status)) {
-    ASSERT_EFI_ERROR (Status);
-    return Status;
-  }
-
-  //
-  // Flush the instruction cache so the image data is written before we execute it
-  //
-  if (ImageContext.ImageAddress != (EFI_PHYSICAL_ADDRESS)(UINTN)Pe32Data) {
-    InvalidateInstructionCacheRange ((VOID *)(UINTN)ImageContext.ImageAddress, (UINTN)ImageContext.ImageSize);
-  }
 
   return Status;
 }
@@ -540,7 +390,8 @@ LoadAndRelocatePeCoffImageInPlace (
 EFI_STATUS
 PeiGetPe32Data (
   IN     EFI_PEI_FILE_HANDLE  FileHandle,
-  OUT    VOID                 **Pe32Data
+  OUT    VOID                 **Pe32Data,
+  OUT    UINT32               *Pe32DataSize
   )
 {
   EFI_STATUS        Status;
@@ -562,22 +413,24 @@ PeiGetPe32Data (
   // Try to find a first exe section (if PcdPeiCoreImageLoaderSearchTeSectionFirst
   // is true, TE will be searched first).
   //
-  Status = PeiServicesFfsFindSectionData3 (
+  Status = PeiServicesFfsFindSectionData4 (
              SearchType1,
              0,
              FileHandle,
              Pe32Data,
+             Pe32DataSize,
              &AuthenticationState
              );
   //
   // If we didn't find a first exe section, try to find the second exe section.
   //
   if (EFI_ERROR (Status)) {
-    Status = PeiServicesFfsFindSectionData3 (
+    Status = PeiServicesFfsFindSectionData4 (
                SearchType2,
                0,
                FileHandle,
                Pe32Data,
+               Pe32DataSize,
                &AuthenticationState
                );
   }
@@ -617,15 +470,13 @@ PeiLoadImageLoadImage (
 {
   EFI_STATUS            Status;
   VOID                  *Pe32Data;
+  UINT32                      Pe32DataSize;
   EFI_PHYSICAL_ADDRESS  ImageAddress;
-  UINT64                ImageSize;
-  EFI_PHYSICAL_ADDRESS  ImageEntryPoint;
-  UINT16                Machine;
   EFI_SECTION_TYPE      SearchType1;
   EFI_SECTION_TYPE      SearchType2;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT ImageContext;
 
   *EntryPoint          = 0;
-  ImageSize            = 0;
   *AuthenticationState = 0;
 
   if (FeaturePcdGet (PcdPeiCoreImageLoaderSearchTeSectionFirst)) {
@@ -640,22 +491,24 @@ PeiLoadImageLoadImage (
   // Try to find a first exe section (if PcdPeiCoreImageLoaderSearchTeSectionFirst
   // is true, TE will be searched first).
   //
-  Status = PeiServicesFfsFindSectionData3 (
+  Status = PeiServicesFfsFindSectionData4 (
              SearchType1,
              0,
              FileHandle,
              &Pe32Data,
+             &Pe32DataSize,
              AuthenticationState
              );
   //
   // If we didn't find a first exe section, try to find the second exe section.
   //
   if (EFI_ERROR (Status)) {
-    Status = PeiServicesFfsFindSectionData3 (
+    Status = PeiServicesFfsFindSectionData4 (
                SearchType2,
                0,
                FileHandle,
                &Pe32Data,
+               &Pe32DataSize,
                AuthenticationState
                );
     if (EFI_ERROR (Status)) {
@@ -672,12 +525,12 @@ PeiLoadImageLoadImage (
   //
   // If memory is installed, perform the shadow operations
   //
-  Status = LoadAndRelocatePeCoffImage (
+  Status = LoadAndRelocateUefiImage (
              FileHandle,
              Pe32Data,
-             &ImageAddress,
+    Pe32DataSize,
-             &ImageSize,
+    &ImageContext,
-             &ImageEntryPoint
+    &ImageAddress
              );
 
   if (EFI_ERROR (Status)) {
@@ -687,83 +540,47 @@ PeiLoadImageLoadImage (
   //
   // Got the entry point from the loaded Pe32Data
   //
-  Pe32Data    = (VOID *)((UINTN)ImageAddress);
+  *EntryPoint = UefiImageLoaderGetImageEntryPoint (&ImageContext);
-  *EntryPoint = ImageEntryPoint;
-
-  Machine = PeCoffLoaderGetMachineType (Pe32Data);
-
-  if (!EFI_IMAGE_MACHINE_TYPE_SUPPORTED (Machine)) {
-    if (!EFI_IMAGE_MACHINE_CROSS_TYPE_SUPPORTED (Machine)) {
-      return EFI_UNSUPPORTED;
-    }
-  }
 
   if (ImageAddressArg != NULL) {
     *ImageAddressArg = ImageAddress;
   }
 
   if (ImageSizeArg != NULL) {
-    *ImageSizeArg = ImageSize;
+    *ImageSizeArg =UefiImageGetImageSize (&ImageContext);
   }
 
   DEBUG_CODE_BEGIN ();
-  CHAR8  *AsciiString;
+    CHAR8  EfiFileName[512];
-  CHAR8  EfiFileName[512];
+    UINT16 Machine;
-  INT32  Index;
+
-  INT32  StartIndex;
+    Machine = UefiImageGetMachine (&ImageContext);
 
-  //
-  // Print debug message: Loading PEIM at 0x12345678 EntryPoint=0x12345688 Driver.efi
-  //
-  if (Machine != EFI_IMAGE_MACHINE_IA64) {
-    DEBUG ((DEBUG_INFO | DEBUG_LOAD, "Loading PEIM at 0x%11p EntryPoint=0x%11p ", (VOID *)(UINTN)ImageAddress, (VOID *)(UINTN)*EntryPoint));
-  } else {
     //
-    // For IPF Image, the real entry point should be print.
+    // Print debug message: Loading PEIM at 0x12345678 EntryPoint=0x12345688 Driver.efi
     //
-    DEBUG ((DEBUG_INFO | DEBUG_LOAD, "Loading PEIM at 0x%11p EntryPoint=0x%11p ", (VOID *)(UINTN)ImageAddress, (VOID *)(UINTN)(*(UINT64 *)(UINTN)*EntryPoint)));
+    if (Machine != EFI_IMAGE_MACHINE_IA64) {
-  }
+      DEBUG ((DEBUG_INFO | DEBUG_LOAD, "Loading PEIM at 0x%11p EntryPoint=0x%11p ", (VOID *)(UINTN)ImageAddress, (VOID *)(UINTN)*EntryPoint));
-
+    } else {
-  //
+      //
-  // Print Module Name by PeImage PDB file name.
+      // For IPF Image, the real entry point should be print.
-  //
+      //
-  AsciiString = PeCoffLoaderGetPdbPointer (Pe32Data);
+      DEBUG ((DEBUG_INFO | DEBUG_LOAD, "Loading PEIM at 0x%11p EntryPoint=0x%11p ", (VOID *)(UINTN)ImageAddress, (VOID *)(UINTN)(*(UINT64 *)(UINTN)*EntryPoint)));
-
-  if (AsciiString != NULL) {
-    StartIndex = 0;
-    for (Index = 0; AsciiString[Index] != 0; Index++) {
-      if ((AsciiString[Index] == '\\') || (AsciiString[Index] == '/')) {
-        StartIndex = Index + 1;
-      }
     }
 
     //
-    // Copy the PDB file name to our temporary string, and replace .pdb with .efi
+    // Print Module Name by PeImage PDB file name.
-    // The PDB file name is limited in the range of 0~511.
-    // If the length is bigger than 511, trim the redundant characters to avoid overflow in array boundary.
     //
-    for (Index = 0; Index < sizeof (EfiFileName) - 4; Index++) {
+    Status = UefiImageGetModuleNameFromSymbolsPath (
-      EfiFileName[Index] = AsciiString[Index + StartIndex];
+               &ImageContext,
-      if (EfiFileName[Index] == 0) {
+               EfiFileName,
-        EfiFileName[Index] = '.';
+               sizeof (EfiFileName)
-      }
+               );
 
-      if (EfiFileName[Index] == '.') {
+    if (!RETURN_ERROR (Status)) {
-        EfiFileName[Index + 1] = 'e';
+      DEBUG ((DEBUG_INFO | DEBUG_LOAD, "%a", EfiFileName));
-        EfiFileName[Index + 2] = 'f';
-        EfiFileName[Index + 3] = 'i';
-        EfiFileName[Index + 4] = 0;
-        break;
-      }
     }
 
-    if (Index == sizeof (EfiFileName) - 4) {
-      EfiFileName[Index] = 0;
-    }
-
-    DEBUG ((DEBUG_INFO | DEBUG_LOAD, "%a", EfiFileName));
-  }
-
   DEBUG_CODE_END ();
 
   DEBUG ((DEBUG_INFO | DEBUG_LOAD, "\n"));
@@ -814,7 +631,7 @@ PeiLoadImageLoadImageWrapper (
   @retval FALSE      Relocation is not stripped.
 
 **/
-BOOLEAN
+/*BOOLEAN
 RelocationIsStrip (
   IN VOID  *Pe32Data
   )
@@ -863,7 +680,7 @@ RelocationIsStrip (
   }
 
   return FALSE;
-}
+}*/
 
 /**
   Routine to load image file for subsequent execution by LoadFile Ppi.
@@ -896,9 +713,9 @@ PeiLoadImage (
   EFI_PEI_LOAD_FILE_PPI  *LoadFile;
   EFI_PHYSICAL_ADDRESS   ImageAddress;
   UINT64                 ImageSize;
-  BOOLEAN                IsStrip;
+  //BOOLEAN                 IsStrip;
 
-  IsStrip = FALSE;
+  //IsStrip = FALSE;
   //
   // If any instances of PEI_LOAD_FILE_PPI are installed, they are called.
   // one at a time, until one reports EFI_SUCCESS.
@@ -924,7 +741,9 @@ PeiLoadImage (
         //
         // The shadowed PEIM must be relocatable.
         //
-        if (PeimState == PEIM_STATE_REGISTER_FOR_SHADOW) {
+        // FIXME:
+        /*if (PeimState == PEIM_STATE_REGISTER_FOR_SHADOW) {
+          // FIXME: Assumes headers were loaded into the image memory
           IsStrip = RelocationIsStrip ((VOID *)(UINTN)ImageAddress);
           ASSERT (!IsStrip);
           if (IsStrip) {
@@ -935,10 +754,10 @@ PeiLoadImage (
         //
         // The image to be started must have the machine type supported by PeiCore.
         //
-        ASSERT (EFI_IMAGE_MACHINE_TYPE_SUPPORTED (PeCoffLoaderGetMachineType ((VOID *)(UINTN)ImageAddress)));
+        ASSERT (EFI_IMAGE_MACHINE_TYPE_SUPPORTED (UefiImageLoaderGetMachineType ((VOID *)(UINTN)ImageAddress)));
-        if (!EFI_IMAGE_MACHINE_TYPE_SUPPORTED (PeCoffLoaderGetMachineType ((VOID *)(UINTN)ImageAddress))) {
+        if (!EFI_IMAGE_MACHINE_TYPE_SUPPORTED (UefiImageLoaderGetMachineType ((VOID *)(UINTN)ImageAddress))) {
           return EFI_UNSUPPORTED;
-        }
+        }*/
 
         return EFI_SUCCESS;
       }

MdeModulePkg/Core/Pei/PeiMain.h

@@ -36,12 +36,11 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/PerformanceLib.h>
 #include <Library/PeiServicesLib.h>
 #include <Library/ReportStatusCodeLib.h>
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
-#include <Library/PeCoffGetEntryPointLib.h>
 #include <Library/BaseMemoryLib.h>
 #include <Library/CacheMaintenanceLib.h>
 #include <Library/PcdLib.h>
-#include <IndustryStandard/PeImage.h>
+#include <IndustryStandard/PeImage2.h>
 #include <Library/PeiServicesTablePointerLib.h>
 #include <Library/MemoryAllocationLib.h>
 #include <Library/TimerLib.h>
@@ -312,7 +311,7 @@ struct _PEI_CORE_INSTANCE {
   //
   // This field points to the shadowed image read function
   //
-  PE_COFF_LOADER_READ_FILE          ShadowedImageRead;
+  VOID          *ShadowedImageRead;
 
   UINTN                             TempPeimCount;
 
@@ -929,6 +928,8 @@ PeiFfsFindNextFile (
   @param SectionSize          The file size to search.
   @param OutputBuffer         A pointer to the discovered section, if successful.
                               NULL if section not found.
+  @param OutputSize        The size of the discovered section, if successful.
+                           0 if section not found
   @param AuthenticationStatus Updated upon return to point to the authentication status for this section.
   @param IsFfs3Fv             Indicates the FV format.
 
@@ -944,6 +945,7 @@ ProcessSection (
   IN EFI_COMMON_SECTION_HEADER  *Section,
   IN UINTN                      SectionSize,
   OUT VOID                      **OutputBuffer,
+  OUT UINT32                    *OutputSize,
   OUT UINT32                    *AuthenticationStatus,
   IN BOOLEAN                    IsFfs3Fv
   );
@@ -995,6 +997,33 @@ PeiFfsFindSectionData3 (
   OUT UINT32                  *AuthenticationStatus
   );
 
+/**
+  Searches for the next matching section within the specified file.
+
+  @param  PeiServices           An indirect pointer to the EFI_PEI_SERVICES table published by the PEI Foundation.
+  @param  SectionType           The value of the section type to find.
+  @param  SectionInstance       Section instance to find.
+  @param  FileHandle            Handle of the firmware file to search.
+  @param  SectionData           A pointer to the discovered section, if successful.
+  @param  SectionDataSize       The size of the discovered section, if successful.
+  @param  AuthenticationStatus  A pointer to the authentication status for this section.
+
+  @retval EFI_SUCCESS      The section was found.
+  @retval EFI_NOT_FOUND    The section was not found.
+
+**/
+EFI_STATUS
+EFIAPI
+PeiFfsFindSectionData4 (
+  IN CONST EFI_PEI_SERVICES    **PeiServices,
+  IN     EFI_SECTION_TYPE      SectionType,
+  IN     UINTN                 SectionInstance,
+  IN     EFI_PEI_FILE_HANDLE   FileHandle,
+  OUT VOID                     **SectionData,
+  OUT UINT32                   *SectionDataSize,
+  OUT UINT32                   *AuthenticationStatus
+  );
+
 /**
   Search the firmware volumes by index
 
@@ -1417,9 +1446,10 @@ InitializeImageServices (
 
 **/
 EFI_STATUS
-LoadAndRelocatePeCoffImageInPlace (
+LoadAndRelocateUefiImageInPlace (
   IN  VOID  *Pe32Data,
-  IN  VOID  *ImageAddress
+  IN  VOID    *ImageAddress,
+  IN  UINT32  ImageSize
   );
 
 /**
@@ -1435,7 +1465,8 @@ LoadAndRelocatePeCoffImageInPlace (
 EFI_STATUS
 PeiGetPe32Data (
   IN     EFI_PEI_FILE_HANDLE  FileHandle,
-  OUT    VOID                 **Pe32Data
+  OUT    VOID                         **Pe32Data,
+  OUT    UINT32                       *Pe32DataSize
   );
 
 /**

MdeModulePkg/Core/Pei/PeiMain.inf

@@ -53,7 +53,6 @@
 
 [LibraryClasses]
   BaseMemoryLib
-  PeCoffGetEntryPointLib
   ReportStatusCodeLib
   PeiServicesLib
   PerformanceLib
@@ -63,7 +62,7 @@
   DebugLib
   MemoryAllocationLib
   CacheMaintenanceLib
-  PeCoffLib
+  UefiImageLib
   PeiServicesTablePointerLib
   PcdLib
   TimerLib

MdeModulePkg/Core/Pei/PeiMain/PeiMain.c

@@ -65,6 +65,7 @@ EFI_PEI_SERVICES  gPs = {
   PeiFfsGetFileInfo2,
   PeiResetSystem2,
   PeiFreePages,
+  PeiFfsFindSectionData4
 };
 
 /**

MdeModulePkg/Core/Pei/Ppi/Ppi.c

@@ -1082,14 +1082,16 @@ ConvertPeiCorePpiPointers (
   IN  PEI_CORE_FV_HANDLE  *CoreFvHandle
   )
 {
-  EFI_FV_FILE_INFO      FileInfo;
+  EFI_FV_FILE_INFO                FileInfo;
-  EFI_PHYSICAL_ADDRESS  OrgImageBase;
+  EFI_PHYSICAL_ADDRESS            OrgImageBase;
-  EFI_PHYSICAL_ADDRESS  MigratedImageBase;
+  EFI_PHYSICAL_ADDRESS            MigratedImageBase;
-  UINTN                 PeiCoreModuleSize;
+  UINTN                           PeiCoreModuleSize;
-  EFI_PEI_FILE_HANDLE   PeiCoreFileHandle;
+  EFI_PEI_FILE_HANDLE             PeiCoreFileHandle;
-  VOID                  *PeiCoreImageBase;
+  VOID                            *PeiCoreImageBase;
-  VOID                  *PeiCoreEntryPoint;
+  UINT32                          PeiCoreImageSize;
-  EFI_STATUS            Status;
+  //VOID                           *PeiCoreEntryPoint;
+  EFI_STATUS                      Status;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT ImageContext;
 
   PeiCoreFileHandle = NULL;
 
@@ -1108,17 +1110,18 @@ ConvertPeiCorePpiPointers (
     Status = CoreFvHandle->FvPpi->GetFileInfo (CoreFvHandle->FvPpi, PeiCoreFileHandle, &FileInfo);
     ASSERT_EFI_ERROR (Status);
 
-    Status = PeiGetPe32Data (PeiCoreFileHandle, &PeiCoreImageBase);
+    Status = PeiGetPe32Data (PeiCoreFileHandle, &PeiCoreImageBase, &PeiCoreImageSize);
     ASSERT_EFI_ERROR (Status);
 
     //
     // Find PEI Core EntryPoint in the BFV in temporary memory.
     //
-    Status = PeCoffLoaderGetEntryPoint ((VOID *)(UINTN)PeiCoreImageBase, &PeiCoreEntryPoint);
+    // FIXME: "Assume" sanity and skip full initialisation?
+    Status = UefiImageInitializeContext (&ImageContext, (VOID *) (UINTN) PeiCoreImageBase, PeiCoreImageSize);
     ASSERT_EFI_ERROR (Status);
 
     OrgImageBase      = (UINTN)PeiCoreImageBase;
-    MigratedImageBase = (UINTN)_ModuleEntryPoint - ((UINTN)PeiCoreEntryPoint - (UINTN)PeiCoreImageBase);
+    MigratedImageBase = (UINTN)_ModuleEntryPoint - UefiImageGetEntryPointAddress (&ImageContext);
 
     //
     // Size of loaded PEI_CORE in permanent memory.

MdeModulePkg/Core/PiSmmCore/DebugImageInfo.c

@@ -0,0 +1,153 @@
+/** @file
+  Support functions for managing debug image info table when loading and unloading
+  images.
+
+Copyright (c) 2006 - 2018, Intel Corporation. All rights reserved.<BR>
+SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include "PiSmmCore.h"
+
+// FIXME: Unify with DXE
+
+EFI_DEBUG_IMAGE_INFO_TABLE_HEADER  mDebugInfoTableHeader = {
+  0,          // volatile UINT32                 UpdateStatus;
+  0,          // UINT32                          TableSize;
+  NULL        // EFI_DEBUG_IMAGE_INFO            *EfiDebugImageInfoTable;
+};
+
+UINTN mMaxTableEntries = 0;
+
+#define EFI_DEBUG_TABLE_ENTRY_SIZE       (sizeof (VOID *))
+
+/**
+  Creates and initializes the DebugImageInfo Table.  Also creates the configuration
+  table and registers it into the system table.
+
+**/
+VOID
+SmmInitializeDebugImageInfoTable (
+  VOID
+  )
+{
+  EFI_STATUS Status;
+
+  //
+  // Install the EFI_SYSTEM_TABLE_POINTER structure in the EFI System
+  // Configuration Table
+  //
+  Status = SmmInstallConfigurationTable (
+             gSmst,
+             &gEfiDebugImageInfoTableGuid,
+             &mDebugInfoTableHeader,
+             sizeof (mDebugInfoTableHeader)
+             );
+  ASSERT_EFI_ERROR (Status);
+}
+
+
+/**
+  Adds a new DebugImageInfo structure to the DebugImageInfo Table.  Re-Allocates
+  the table if it's not large enough to accomidate another entry.
+
+  @param  ImageInfoType  type of debug image information
+  @param  LoadedImage    pointer to the loaded image protocol for the image being
+                         loaded
+  @param  ImageHandle    image handle for the image being loaded
+
+**/
+VOID
+SmmNewDebugImageInfoEntry (
+  IN     UINT32                           ImageInfoType,
+  IN     EFI_LOADED_IMAGE_PROTOCOL        *LoadedImage,
+  IN     EFI_HANDLE                       ImageHandle,
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
+  )
+{
+  EFI_DEBUG_IMAGE_INFO        *Table;
+  EFI_DEBUG_IMAGE_INFO        *NewTable;
+  UINTN                       Index;
+  UINTN                       TableSize;
+  EFI_DEBUG_IMAGE_INFO_NORMAL *NormalImage;
+  RETURN_STATUS               Status;
+  CONST CHAR8                 *PdbPath;
+  UINT32                      PdbPathSize;
+
+  //
+  // Set the flag indicating that we're in the process of updating the table.
+  //
+  mDebugInfoTableHeader.UpdateStatus |= EFI_DEBUG_IMAGE_INFO_UPDATE_IN_PROGRESS;
+
+  Table = mDebugInfoTableHeader.EfiDebugImageInfoTable;
+
+  if (mDebugInfoTableHeader.TableSize < mMaxTableEntries) {
+    //
+    // We still have empty entires in the Table, find the first empty entry.
+    //
+    Index = 0;
+    while (Table[Index].NormalImage != NULL) {
+      Index++;
+    }
+    //
+    // There must be an empty entry in the in the table.
+    //
+    ASSERT (Index < mMaxTableEntries);
+  } else {
+    //
+    //  Table is full, so re-allocate another page for a larger table...
+    //
+    TableSize = mMaxTableEntries * EFI_DEBUG_TABLE_ENTRY_SIZE;
+    NewTable = AllocateZeroPool (TableSize + EFI_PAGE_SIZE);
+    if (NewTable == NULL) {
+      mDebugInfoTableHeader.UpdateStatus &= ~EFI_DEBUG_IMAGE_INFO_UPDATE_IN_PROGRESS;
+      return;
+    }
+    //
+    // Copy the old table into the new one
+    //
+    CopyMem (NewTable, Table, TableSize);
+    mDebugInfoTableHeader.EfiDebugImageInfoTable = NewTable;
+    //
+    // Enlarge the max table entries.
+    //
+    mMaxTableEntries += EFI_PAGE_SIZE / EFI_DEBUG_TABLE_ENTRY_SIZE;
+    //
+    // Free the old table
+    //
+    SmmFreePool (Table);
+    //
+    // Update the table header
+    //
+    Table = NewTable;
+    //
+    // Set the first empty entry index to be the original max table entries.
+    //
+    Index             = mMaxTableEntries;
+  }
+
+  //
+  // Allocate data for new entry
+  //
+  NormalImage = AllocateZeroPool (sizeof (EFI_DEBUG_IMAGE_INFO_NORMAL));
+  if (NormalImage != NULL) {
+    //
+    // Update the entry
+    //
+    NormalImage->ImageInfoType               = (UINT32) ImageInfoType;
+    NormalImage->LoadedImageProtocolInstance = LoadedImage;
+    NormalImage->ImageHandle                 = ImageHandle;
+
+    Status = UefiImageGetSymbolsPath (ImageContext, &PdbPath, &PdbPathSize);
+    if (!RETURN_ERROR (Status)) {
+      NormalImage->PdbPath = AllocateCopyPool (PdbPathSize, PdbPath);
+    }
+    //
+    // Increase the number of EFI_DEBUG_IMAGE_INFO elements and set the mDebugInfoTable in modified status.
+    //
+    mDebugInfoTableHeader.UpdateStatus |= EFI_DEBUG_IMAGE_INFO_TABLE_MODIFIED;
+    Table[Index].NormalImage = NormalImage;
+    mDebugInfoTableHeader.TableSize++;
+  }
+  mDebugInfoTableHeader.UpdateStatus &= ~EFI_DEBUG_IMAGE_INFO_UPDATE_IN_PROGRESS;
+}

MdeModulePkg/Core/PiSmmCore/Dispatcher.c

@@ -33,7 +33,9 @@
 
 **/
 
+#include "Library/UefiImageLib.h"
 #include "PiSmmCore.h"
+#include "Uefi/UefiBaseType.h"
 
 //
 // SMM Dispatcher Data structures
@@ -213,84 +215,18 @@ CheckAndMarkFixLoadingMemoryUsageBitMap (
 
 **/
 EFI_STATUS
-GetPeCoffImageFixLoadingAssignedAddress (
+GetUefiImageFixLoadingAssignedAddress (
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  OUT    EFI_PHYSICAL_ADDRESS          *LoadAddress,
+  IN     UINT64                        ValueInSectionHeader,
+  IN     UINT32                        ImageDestSize
   )
 {
-  UINTN                            SectionHeaderOffset;
+  RETURN_STATUS        Status;
-  EFI_STATUS                       Status;
+  EFI_PHYSICAL_ADDRESS FixLoadingAddress;
-  EFI_IMAGE_SECTION_HEADER         SectionHeader;
-  EFI_IMAGE_OPTIONAL_HEADER_UNION  *ImgHdr;
-  EFI_PHYSICAL_ADDRESS             FixLoadingAddress;
-  UINT16                           Index;
-  UINTN                            Size;
-  UINT16                           NumberOfSections;
-  UINT64                           ValueInSectionHeader;
 
-  FixLoadingAddress = 0;
+  FixLoadingAddress = (EFI_PHYSICAL_ADDRESS)(gLoadModuleAtFixAddressSmramBase + ValueInSectionHeader);
-  Status            = EFI_NOT_FOUND;
+  Status = CheckAndMarkFixLoadingMemoryUsageBitMap (FixLoadingAddress, ImageDestSize);
-
+  *LoadAddress = FixLoadingAddress;
-  //
-  // Get PeHeader pointer
-  //
-  ImgHdr              = (EFI_IMAGE_OPTIONAL_HEADER_UNION *)((CHAR8 *)ImageContext->Handle + ImageContext->PeCoffHeaderOffset);
-  SectionHeaderOffset = ImageContext->PeCoffHeaderOffset +
-                        sizeof (UINT32) +
-                        sizeof (EFI_IMAGE_FILE_HEADER) +
-                        ImgHdr->Pe32.FileHeader.SizeOfOptionalHeader;
-  NumberOfSections = ImgHdr->Pe32.FileHeader.NumberOfSections;
-
-  //
-  // Get base address from the first section header that doesn't point to code section.
-  //
-  for (Index = 0; Index < NumberOfSections; Index++) {
-    //
-    // Read section header from file
-    //
-    Size   = sizeof (EFI_IMAGE_SECTION_HEADER);
-    Status = ImageContext->ImageRead (
-                             ImageContext->Handle,
-                             SectionHeaderOffset,
-                             &Size,
-                             &SectionHeader
-                             );
-    if (EFI_ERROR (Status)) {
-      return Status;
-    }
-
-    Status = EFI_NOT_FOUND;
-
-    if ((SectionHeader.Characteristics & EFI_IMAGE_SCN_CNT_CODE) == 0) {
-      //
-      // Build tool will save the address in PointerToRelocations & PointerToLineNumbers fields in the first section header
-      // that doesn't point to code section in image header.So there is an assumption that when the feature is enabled,
-      // if a module with a loading address assigned by tools, the PointerToRelocations & PointerToLineNumbers fields
-      // should not be Zero, or else, these 2 fields should be set to Zero
-      //
-      ValueInSectionHeader = ReadUnaligned64 ((UINT64 *)&SectionHeader.PointerToRelocations);
-      if (ValueInSectionHeader != 0) {
-        //
-        // Found first section header that doesn't point to code section in which build tool saves the
-        // offset to SMRAM base as image base in PointerToRelocations & PointerToLineNumbers fields
-        //
-        FixLoadingAddress = (EFI_PHYSICAL_ADDRESS)(gLoadModuleAtFixAddressSmramBase + (INT64)ValueInSectionHeader);
-        //
-        // Check if the memory range is available.
-        //
-        Status = CheckAndMarkFixLoadingMemoryUsageBitMap (FixLoadingAddress, (UINTN)(ImageContext->ImageSize + ImageContext->SectionAlignment));
-        if (!EFI_ERROR (Status)) {
-          //
-          // The assigned address is valid. Return the specified loading address
-          //
-          ImageContext->ImageAddress = FixLoadingAddress;
-        }
-      }
-
-      break;
-    }
-
-    SectionHeaderOffset += sizeof (EFI_IMAGE_SECTION_HEADER);
-  }
 
   DEBUG ((DEBUG_INFO|DEBUG_LOAD, "LOADING MODULE FIXED INFO: Loading module at fixed address %x, Status = %r\n", FixLoadingAddress, Status));
   return Status;
@@ -307,24 +243,29 @@ GetPeCoffImageFixLoadingAssignedAddress (
 EFI_STATUS
 EFIAPI
 SmmLoadImage (
-  IN OUT EFI_SMM_DRIVER_ENTRY  *DriverEntry
+  IN OUT EFI_SMM_DRIVER_ENTRY            *DriverEntry,
+     OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT *ImageContext
   )
 {
   UINT32                         AuthenticationStatus;
   UINTN                          FilePathSize;
   VOID                           *Buffer;
   UINTN                          Size;
-  UINTN                          PageCount;
+  UINT32                         DstBufferPages;
   EFI_GUID                       *NameGuid;
   EFI_STATUS                     Status;
   EFI_STATUS                     SecurityStatus;
   EFI_HANDLE                     DeviceHandle;
-  EFI_PHYSICAL_ADDRESS           DstBuffer;
+  UINT32                         ImageSize;
+  UINT32                         ImageAlignment;
+  UINT64                         ValueInSectionHeader;
+  VOID                           *DstBuffer;
+  UINT32                         DstBufferSize;
   EFI_DEVICE_PATH_PROTOCOL       *FilePath;
   EFI_DEVICE_PATH_PROTOCOL       *OriginalFilePath;
   EFI_DEVICE_PATH_PROTOCOL       *HandleFilePath;
   EFI_FIRMWARE_VOLUME2_PROTOCOL  *Fv;
-  PE_COFF_LOADER_IMAGE_CONTEXT   ImageContext;
+  EFI_PHYSICAL_ADDRESS           LoadAddress;
 
   PERF_LOAD_IMAGE_BEGIN (DriverEntry->ImageHandle);
 
@@ -413,6 +354,18 @@ SmmLoadImage (
     return Status;
   }
 
+  //
+  // Get information about the image being loaded
+  //
+  Status = UefiImageInitializeContextPreHash (ImageContext, Buffer, (UINT32) Size);
+  if (EFI_ERROR (Status)) {
+    if (Buffer != NULL) {
+      gBS->FreePool (Buffer);
+    }
+    return Status;
+  }
+
+  // FIXME: Context?
   //
   // Verify File Authentication through the Security2 Architectural Protocol
   //
@@ -420,8 +373,8 @@ SmmLoadImage (
     SecurityStatus = mSecurity2->FileAuthentication (
                                    mSecurity2,
                                    OriginalFilePath,
-                                   Buffer,
+                                  ImageContext,
-                                   Size,
+                                  sizeof (*ImageContext),
                                    FALSE
                                    );
   }
@@ -444,24 +397,23 @@ SmmLoadImage (
     return Status;
   }
 
-  //
+  Status = UefiImageInitializeContextPostHash (ImageContext);
-  // Initialize ImageContext
+  if (RETURN_ERROR (Status)) {
-  //
-  ImageContext.Handle    = Buffer;
-  ImageContext.ImageRead = PeCoffLoaderImageReadFromMemory;
-
-  //
-  // Get information about the image being loaded
-  //
-  Status = PeCoffLoaderGetImageInfo (&ImageContext);
-  if (EFI_ERROR (Status)) {
-    if (Buffer != NULL) {
-      gBS->FreePool (Buffer);
-    }
-
     return Status;
   }
 
+  //
+  // Stripped relocations are not supported for both fixed-address and dynamic
+  // loading.
+  //
+  if (UefiImageGetRelocsStripped (ImageContext)) {
+    return EFI_UNSUPPORTED;
+  }
+
+  ImageSize      = UefiImageGetImageSize (ImageContext);
+  DstBufferPages = EFI_SIZE_TO_PAGES (ImageSize);
+  DstBufferSize  = EFI_PAGES_TO_SIZE (DstBufferPages);
+  ImageAlignment = UefiImageGetSegmentAlignment (ImageContext);
   //
   // if Loading module at Fixed Address feature is enabled, then  cut out a memory range started from TESG BASE
   // to hold the Smm driver code
@@ -470,102 +422,68 @@ SmmLoadImage (
     //
     // Get the fixed loading address assigned by Build tool
     //
-    Status = GetPeCoffImageFixLoadingAssignedAddress (&ImageContext);
+    Status = UefiImageGetFixedAddress (ImageContext, &ValueInSectionHeader);
+    if (!RETURN_ERROR (Status)) {
+      Status = GetUefiImageFixLoadingAssignedAddress (&LoadAddress, ValueInSectionHeader, DstBufferSize);
+    }
+
     if (!EFI_ERROR (Status)) {
       //
       // Since the memory range to load Smm core already been cut out, so no need to allocate and free this range
       // following statements is to bypass SmmFreePages
       //
-      PageCount = 0;
+      DstBufferPages = 0;
-      DstBuffer = (UINTN)gLoadModuleAtFixAddressSmramBase;
+      DstBuffer      = (VOID *)(UINTN)LoadAddress;
     } else {
       DEBUG ((DEBUG_INFO|DEBUG_LOAD, "LOADING MODULE FIXED ERROR: Failed to load module at fixed address. \n"));
       //
       // allocate the memory to load the SMM driver
       //
-      PageCount = (UINTN)EFI_SIZE_TO_PAGES ((UINTN)ImageContext.ImageSize + ImageContext.SectionAlignment);
+      DstBuffer = AllocateAlignedCodePages (DstBufferPages, ImageAlignment);
-      DstBuffer = (UINTN)(-1);
+      if (DstBuffer == NULL) {
-
-      Status = SmmAllocatePages (
-                 AllocateMaxAddress,
-                 EfiRuntimeServicesCode,
-                 PageCount,
-                 &DstBuffer
-                 );
-      if (EFI_ERROR (Status)) {
         if (Buffer != NULL) {
           gBS->FreePool (Buffer);
         }
 
-        return Status;
+         return EFI_OUT_OF_RESOURCES;
-      }
+       }
-
-      ImageContext.ImageAddress = (EFI_PHYSICAL_ADDRESS)DstBuffer;
     }
   } else {
-    PageCount = (UINTN)EFI_SIZE_TO_PAGES ((UINTN)ImageContext.ImageSize + ImageContext.SectionAlignment);
+    DstBuffer = AllocateAlignedCodePages (DstBufferPages, ImageAlignment);
-    DstBuffer = (UINTN)(-1);
+    if (DstBuffer == NULL) {
-
-    Status = SmmAllocatePages (
-               AllocateMaxAddress,
-               EfiRuntimeServicesCode,
-               PageCount,
-               &DstBuffer
-               );
-    if (EFI_ERROR (Status)) {
       if (Buffer != NULL) {
         gBS->FreePool (Buffer);
       }
 
-      return Status;
+      return EFI_OUT_OF_RESOURCES;
     }
-
-    ImageContext.ImageAddress = (EFI_PHYSICAL_ADDRESS)DstBuffer;
   }
 
-  //
-  // Align buffer on section boundary
-  //
-  ImageContext.ImageAddress += ImageContext.SectionAlignment - 1;
-  ImageContext.ImageAddress &= ~((EFI_PHYSICAL_ADDRESS)ImageContext.SectionAlignment - 1);
-
   //
   // Load the image to our new buffer
   //
-  Status = PeCoffLoaderLoadImage (&ImageContext);
+  Status = UefiImageLoadImageForExecution (
+    ImageContext,
+    DstBuffer,
+    DstBufferSize,
+    NULL,
+    0
+    );
   if (EFI_ERROR (Status)) {
     if (Buffer != NULL) {
       gBS->FreePool (Buffer);
     }
 
-    SmmFreePages (DstBuffer, PageCount);
+    FreeAlignedPages (DstBuffer, DstBufferPages);
     return Status;
   }
 
-  //
-  // Relocate the image in our new buffer
-  //
-  Status = PeCoffLoaderRelocateImage (&ImageContext);
-  if (EFI_ERROR (Status)) {
-    if (Buffer != NULL) {
-      gBS->FreePool (Buffer);
-    }
-
-    SmmFreePages (DstBuffer, PageCount);
-    return Status;
-  }
-
-  //
-  // Flush the instruction cache so the image data are written before we execute it
-  //
-  InvalidateInstructionCacheRange ((VOID *)(UINTN)ImageContext.ImageAddress, (UINTN)ImageContext.ImageSize);
-
   //
   // Save Image EntryPoint in DriverEntry
   //
-  DriverEntry->ImageEntryPoint = ImageContext.EntryPoint;
+  DriverEntry->ImageEntryPoint  = UefiImageLoaderGetImageEntryPoint (ImageContext);
-  DriverEntry->ImageBuffer     = DstBuffer;
+  DriverEntry->ImageBuffer      = (UINTN)DstBuffer;
-  DriverEntry->NumberOfPage    = PageCount;
+  DriverEntry->NumberOfPage    = DstBufferPages;
 
   //
   // Allocate a Loaded Image Protocol in EfiBootServicesData
@@ -576,7 +494,7 @@ SmmLoadImage (
       gBS->FreePool (Buffer);
     }
 
-    SmmFreePages (DstBuffer, PageCount);
+    FreeAlignedPages (DstBuffer, DstBufferPages);
     return Status;
   }
 
@@ -604,14 +522,14 @@ SmmLoadImage (
       gBS->FreePool (Buffer);
     }
 
-    SmmFreePages (DstBuffer, PageCount);
+    FreeAlignedPages (DstBuffer, DstBufferPages);
     return Status;
   }
 
   CopyMem (DriverEntry->LoadedImage->FilePath, FilePath, GetDevicePathSize (FilePath));
 
-  DriverEntry->LoadedImage->ImageBase     = (VOID *)(UINTN)ImageContext.ImageAddress;
+  DriverEntry->LoadedImage->ImageBase     = DstBuffer;
-  DriverEntry->LoadedImage->ImageSize     = ImageContext.ImageSize;
+  DriverEntry->LoadedImage->ImageSize     = ImageSize;
   DriverEntry->LoadedImage->ImageCodeType = EfiRuntimeServicesCode;
   DriverEntry->LoadedImage->ImageDataType = EfiRuntimeServicesData;
 
@@ -625,14 +543,14 @@ SmmLoadImage (
     }
 
     gBS->FreePool (DriverEntry->LoadedImage->FilePath);
-    SmmFreePages (DstBuffer, PageCount);
+    FreeAlignedPages (DstBuffer, DstBufferPages);
     return Status;
   }
 
-  CopyMem (DriverEntry->SmmLoadedImage.FilePath, FilePath, GetDevicePathSize (FilePath));
+  CopyMem (DriverEntry->SmmLoadedImage.FilePath, FilePath, GetDevicePathSize(FilePath));
 
-  DriverEntry->SmmLoadedImage.ImageBase     = (VOID *)(UINTN)ImageContext.ImageAddress;
+  DriverEntry->SmmLoadedImage.ImageBase = DstBuffer;
-  DriverEntry->SmmLoadedImage.ImageSize     = ImageContext.ImageSize;
+  DriverEntry->SmmLoadedImage.ImageSize = ImageSize;
   DriverEntry->SmmLoadedImage.ImageCodeType = EfiRuntimeServicesCode;
   DriverEntry->SmmLoadedImage.ImageDataType = EfiRuntimeServicesData;
 
@@ -660,60 +578,42 @@ SmmLoadImage (
 
   PERF_LOAD_IMAGE_END (DriverEntry->ImageHandle);
 
+  SmmInsertImageRecord (&DriverEntry->SmmLoadedImage, ImageContext);
+
+  //
+  // Register the image in the Debug Image Info Table if the attribute is set
+  //
+  SmmNewDebugImageInfoEntry (
+    EFI_DEBUG_IMAGE_INFO_TYPE_NORMAL,
+    &DriverEntry->SmmLoadedImage,
+    DriverEntry->SmmImageHandle,
+    ImageContext
+    );
+
   //
   // Print the load address and the PDB file name if it is available
   //
 
   DEBUG_CODE_BEGIN ();
 
-  UINTN  Index;
+  CHAR8 EfiFileName[256];
-  UINTN  StartIndex;
-  CHAR8  EfiFileName[256];
 
-  DEBUG ((
+  DEBUG ((DEBUG_INFO | DEBUG_LOAD,
-    DEBUG_INFO | DEBUG_LOAD,
+         "Loading SMM driver at 0x%11p EntryPoint=0x%11p ",
-    "Loading SMM driver at 0x%11p EntryPoint=0x%11p ",
+         DstBuffer,
-    (VOID *)(UINTN)ImageContext.ImageAddress,
+         FUNCTION_ENTRY_POINT (UefiImageLoaderGetImageEntryPoint (ImageContext))));
-    FUNCTION_ENTRY_POINT (ImageContext.EntryPoint)
-    ));
 
   //
   // Print Module Name by Pdb file path.
   // Windows and Unix style file path are all trimmed correctly.
   //
-  if (ImageContext.PdbPointer != NULL) {
+  Status = UefiImageGetModuleNameFromSymbolsPath (
-    StartIndex = 0;
+               ImageContext,
-    for (Index = 0; ImageContext.PdbPointer[Index] != 0; Index++) {
+               EfiFileName,
-      if ((ImageContext.PdbPointer[Index] == '\\') || (ImageContext.PdbPointer[Index] == '/')) {
+               sizeof (EfiFileName)
-        StartIndex = Index + 1;
+               );
-      }
+  if (!EFI_ERROR (Status)) {
-    }
+    DEBUG ((DEBUG_INFO | DEBUG_LOAD, "%a", EfiFileName));
-
-    //
-    // Copy the PDB file name to our temporary string, and replace .pdb with .efi
-    // The PDB file name is limited in the range of 0~255.
-    // If the length is bigger than 255, trim the redundant characters to avoid overflow in array boundary.
-    //
-    for (Index = 0; Index < sizeof (EfiFileName) - 4; Index++) {
-      EfiFileName[Index] = ImageContext.PdbPointer[Index + StartIndex];
-      if (EfiFileName[Index] == 0) {
-        EfiFileName[Index] = '.';
-      }
-
-      if (EfiFileName[Index] == '.') {
-        EfiFileName[Index + 1] = 'e';
-        EfiFileName[Index + 2] = 'f';
-        EfiFileName[Index + 3] = 'i';
-        EfiFileName[Index + 4] = 0;
-        break;
-      }
-    }
-
-    if (Index == sizeof (EfiFileName) - 4) {
-      EfiFileName[Index] = 0;
-    }
-
-    DEBUG ((DEBUG_INFO | DEBUG_LOAD, "%a", EfiFileName));   // &Image->ImageContext.PdbPointer[StartIndex]));
   }
 
   DEBUG ((DEBUG_INFO | DEBUG_LOAD, "\n"));
@@ -852,11 +752,12 @@ SmmDispatcher (
   VOID
   )
 {
-  EFI_STATUS            Status;
+  EFI_STATUS                      Status;
-  LIST_ENTRY            *Link;
+  LIST_ENTRY                      *Link;
-  EFI_SMM_DRIVER_ENTRY  *DriverEntry;
+  EFI_SMM_DRIVER_ENTRY            *DriverEntry;
-  BOOLEAN               ReadyToRun;
+  BOOLEAN                         ReadyToRun;
-  BOOLEAN               PreviousSmmEntryPointRegistered;
+  BOOLEAN                         PreviousSmmEntryPointRegistered;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT ImageContext;
 
   if (!gRequestDispatch) {
     return EFI_NOT_FOUND;
@@ -889,8 +790,7 @@ SmmDispatcher (
       // skip the LoadImage
       //
       if (DriverEntry->ImageHandle == NULL) {
-        Status = SmmLoadImage (DriverEntry);
+        Status = SmmLoadImage (DriverEntry, &ImageContext);
-
         //
         // Update the driver state to reflect that it's been loaded
         //
@@ -929,7 +829,7 @@ SmmDispatcher (
       //
       // For each SMM driver, pass NULL as ImageHandle
       //
-      RegisterSmramProfileImage (DriverEntry, TRUE);
+      RegisterSmramProfileImage (&DriverEntry->FileName, TRUE, &ImageContext);
       PERF_START_IMAGE_BEGIN (DriverEntry->ImageHandle);
       Status = ((EFI_IMAGE_ENTRY_POINT)(UINTN)DriverEntry->ImageEntryPoint)(DriverEntry->ImageHandle, gST);
       PERF_START_IMAGE_END (DriverEntry->ImageHandle);
@@ -940,7 +840,12 @@ SmmDispatcher (
           DriverEntry->SmmLoadedImage.ImageBase,
           Status
           ));
-        UnregisterSmramProfileImage (DriverEntry, TRUE);
+        UnregisterSmramProfileImage (
+          &DriverEntry->FileName,
+          (UINTN) DriverEntry->LoadedImage->ImageBase,
+          DriverEntry->LoadedImage->ImageSize,
+          TRUE
+          );
         SmmFreePages (DriverEntry->ImageBuffer, DriverEntry->NumberOfPage);
         //
         // Uninstall LoadedImage

MdeModulePkg/Core/PiSmmCore/MemoryAllocation.c

@@ -22,7 +22,8 @@
 #include <Library/UefiBootServicesTableLib.h>
 #include <Library/BaseMemoryLib.h>
 #include <Library/DebugLib.h>
-#include "PiSmmCoreMemoryAllocationServices.h"
+#include "PiSmmCore.h"
+#include "PiSmmCorePrivateData.h"
 
 #include <Library/MemoryProfileLib.h>
 

MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c

@@ -16,12 +16,12 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/PcdLib.h>
 #include <Library/ImagePropertiesRecordLib.h>
 
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
-#include <Library/PeCoffGetEntryPointLib.h>
 
 #include <Guid/PiSmmMemoryAttributesTable.h>
 
 #include "PiSmmCore.h"
+#include "ProcessorBind.h"
 
 #define PREVIOUS_MEMORY_DESCRIPTOR(MemoryDescriptor, Size) \
   ((EFI_MEMORY_DESCRIPTOR *)((UINT8 *)(MemoryDescriptor) - (Size)))
@@ -31,7 +31,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 typedef struct {
   UINT32        Signature;
   UINTN         ImageRecordCount;
-  UINTN         CodeSegmentCountMax;
+  UINT32        NumberOfSectionsMax;
   LIST_ENTRY    ImageRecordList;
 } IMAGE_PROPERTIES_PRIVATE_DATA;
 
@@ -212,7 +212,10 @@ SmmCoreGetMemoryMapMemoryAttributesTable (
     return EFI_INVALID_PARAMETER;
   }
 
-  AdditionalRecordCount = (2 * mImagePropertiesPrivateData.CodeSegmentCountMax + 3) * mImagePropertiesPrivateData.ImageRecordCount;
+  //
+  // Per image, they may be one trailer. There may be prefixed data.
+  //
+  AdditionalRecordCount = (mImagePropertiesPrivateData.NumberOfSectionsMax + 1) * mImagePropertiesPrivateData.ImageRecordCount + 1;
 
   OldMemoryMapSize = *MemoryMapSize;
   Status           = SmmCoreGetMemoryMap (MemoryMapSize, MemoryMap, MapKey, DescriptorSize, DescriptorVersion);
@@ -251,6 +254,98 @@ SmmCoreGetMemoryMapMemoryAttributesTable (
 // Below functions are for ImageRecord
 //
 
+/**
+  Set MemoryProtectionAttribute according to PE/COFF image section alignment.
+
+  @param[in]  SectionAlignment    PE/COFF section alignment
+**/
+STATIC
+VOID
+SetMemoryAttributesTableSectionAlignment (
+  IN UINT32  SectionAlignment
+  )
+{
+  if (((SectionAlignment & (RUNTIME_PAGE_ALLOCATION_GRANULARITY - 1)) != 0) &&
+      ((mMemoryProtectionAttribute & EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA) != 0))
+  {
+    DEBUG ((DEBUG_VERBOSE, "SMM SetMemoryAttributesTableSectionAlignment - Clear\n"));
+    mMemoryProtectionAttribute &= ~((UINT64)EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA);
+  }
+}
+
+/**
+  Sort image record based upon the ImageBase from low to high.
+**/
+STATIC
+VOID
+InsertSortImageRecord (
+  IN UEFI_IMAGE_RECORD        *NewImageRecord
+  )
+{
+  UEFI_IMAGE_RECORD            *ImageRecord;
+  LIST_ENTRY                   *PrevImageRecordLink;
+  LIST_ENTRY                   *ImageRecordLink;
+  LIST_ENTRY                   *ImageRecordList;
+
+  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
+
+  PrevImageRecordLink = ImageRecordList;
+  for (
+    ImageRecordLink = GetFirstNode (ImageRecordList);
+    !IsNull (ImageRecordLink, ImageRecordList);
+    ImageRecordLink = GetNextNode (ImageRecordList, PrevImageRecordLink)
+    ) {
+    ImageRecord = CR (
+                    ImageRecordLink,
+                    UEFI_IMAGE_RECORD,
+                    Link,
+                    UEFI_IMAGE_RECORD_SIGNATURE
+                    );
+    if (NewImageRecord->StartAddress < ImageRecord->StartAddress) {
+      break;
+    }
+
+    PrevImageRecordLink = ImageRecordLink;
+  }
+
+  InsertHeadList (PrevImageRecordLink, &NewImageRecord->Link);
+  mImagePropertiesPrivateData.ImageRecordCount++;
+
+  if (mImagePropertiesPrivateData.NumberOfSectionsMax < NewImageRecord->NumSegments) {
+    mImagePropertiesPrivateData.NumberOfSectionsMax = NewImageRecord->NumSegments;
+  }
+}
+
+/**
+  Dump image record.
+**/
+STATIC
+VOID
+DumpImageRecord (
+  VOID
+  )
+{
+  UEFI_IMAGE_RECORD            *ImageRecord;
+  LIST_ENTRY               *ImageRecordLink;
+  LIST_ENTRY               *ImageRecordList;
+  UINTN                    Index;
+
+  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
+
+  for (ImageRecordLink = ImageRecordList->ForwardLink, Index = 0;
+       ImageRecordLink != ImageRecordList;
+       ImageRecordLink = ImageRecordLink->ForwardLink, Index++)
+  {
+    ImageRecord = CR (
+                    ImageRecordLink,
+                    UEFI_IMAGE_RECORD,
+                    Link,
+                    UEFI_IMAGE_RECORD_SIGNATURE
+                    );
+    DEBUG ((DEBUG_VERBOSE, "SMM  Image[%d]: 0x%016lx - 0x%016lx\n", Index, ImageRecord->StartAddress, ImageRecord->EndAddress - ImageRecord->StartAddress));
+  }
+}
+
 /**
   Insert image record.
 
@@ -258,82 +353,60 @@ SmmCoreGetMemoryMapMemoryAttributesTable (
 **/
 VOID
 SmmInsertImageRecord (
-  IN EFI_SMM_DRIVER_ENTRY  *DriverEntry
+  IN EFI_LOADED_IMAGE_PROTOCOL        *LoadedImage,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   )
 {
-  EFI_STATUS               Status;
+  RETURN_STATUS      PdbStatus;
-  IMAGE_PROPERTIES_RECORD  *ImageRecord;
+  PHYSICAL_ADDRESS   ImageBuffer;
-  CHAR8                    *PdbPointer;
+  UINTN              NumberOfPage;
-  UINT32                   RequiredAlignment;
+  UINT32             SectionAlignment;
+  UEFI_IMAGE_RECORD  *ImageRecord;
+  CONST CHAR8        *PdbPointer;
+  UINT32             PdbSize;
 
-  DEBUG ((DEBUG_VERBOSE, "SMM InsertImageRecord - 0x%x\n", DriverEntry));
+  ImageBuffer  = (UINTN)LoadedImage->ImageBase;
+  NumberOfPage = EFI_SIZE_TO_PAGES((UINTN)LoadedImage->ImageSize);
 
-  ImageRecord = AllocatePool (sizeof (*ImageRecord));
+  DEBUG ((DEBUG_VERBOSE, "SMM InsertImageRecord - 0x%016lx - 0x%08x\n", ImageBuffer, NumberOfPage));
+
+  DEBUG ((DEBUG_VERBOSE, "SMM ImageRecordCount - 0x%x\n", mImagePropertiesPrivateData.ImageRecordCount));
+
+  PdbStatus = UefiImageGetSymbolsPath (ImageContext, &PdbPointer, &PdbSize);
+  if (!RETURN_ERROR (PdbStatus)) {
+    DEBUG ((DEBUG_VERBOSE, "SMM   Image - %a\n", PdbPointer));
+  }
+
+  //
+  // Get SectionAlignment
+  //
+  SectionAlignment = UefiImageGetSegmentAlignment (ImageContext);
+
+  SetMemoryAttributesTableSectionAlignment (SectionAlignment);
+  if ((SectionAlignment & (RUNTIME_PAGE_ALLOCATION_GRANULARITY - 1)) != 0) {
+    DEBUG ((
+      DEBUG_WARN,
+      "SMM !!!!!!!!  InsertImageRecord - Section Alignment(0x%x) is not %dK  !!!!!!!!\n",
+      SectionAlignment, RUNTIME_PAGE_ALLOCATION_GRANULARITY >> 10));
+    if (!RETURN_ERROR (PdbStatus)) {
+      DEBUG ((DEBUG_WARN, "SMM !!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));
+    }
+
+    return;
+  }
+
+  //
+  // The image headers are not recorded among the sections, allocate one more.
+  //
+  ImageRecord = UefiImageLoaderGetImageRecord (ImageContext);
   if (ImageRecord == NULL) {
     return;
   }
 
-  InitializeListHead (&ImageRecord->Link);
+  UefiImageDebugPrintSegments (ImageContext);
-  InitializeListHead (&ImageRecord->CodeSegmentList);
+  UefiImageDebugPrintImageRecord (ImageRecord);
 
-  PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)DriverEntry->ImageBuffer);
+  InsertSortImageRecord (ImageRecord);
-  if (PdbPointer != NULL) {
-    DEBUG ((DEBUG_VERBOSE, "SMM   Image - %a\n", PdbPointer));
-  }
-
-  RequiredAlignment = RUNTIME_PAGE_ALLOCATION_GRANULARITY;
-  Status            = CreateImagePropertiesRecord (
-                        (VOID *)(UINTN)DriverEntry->ImageBuffer,
-                        LShiftU64 (DriverEntry->NumberOfPage, EFI_PAGE_SHIFT),
-                        &RequiredAlignment,
-                        ImageRecord
-                        );
-
-  if (EFI_ERROR (Status)) {
-    if (Status == EFI_ABORTED) {
-      mMemoryProtectionAttribute &=
-        ~((UINT64)EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA);
-    }
-
-    goto Finish;
-  }
-
-  if (ImageRecord->CodeSegmentCount == 0) {
-    mMemoryProtectionAttribute &=
-      ~((UINT64)EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA);
-    DEBUG ((DEBUG_ERROR, "SMM !!!!!!!!  InsertImageRecord - CodeSegmentCount is 0  !!!!!!!!\n"));
-    if (PdbPointer != NULL) {
-      DEBUG ((DEBUG_ERROR, "SMM !!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));
-    }
-
-    Status = EFI_ABORTED;
-    goto Finish;
-  }
-
-  //
-  // Check overlap all section in ImageBase/Size
-  //
-  if (!IsImageRecordCodeSectionValid (ImageRecord)) {
-    DEBUG ((DEBUG_ERROR, "SMM IsImageRecordCodeSectionValid - FAIL\n"));
-    Status = EFI_ABORTED;
-    goto Finish;
-  }
-
-  InsertTailList (&mImagePropertiesPrivateData.ImageRecordList, &ImageRecord->Link);
-  mImagePropertiesPrivateData.ImageRecordCount++;
-
-  if (mImagePropertiesPrivateData.CodeSegmentCountMax < ImageRecord->CodeSegmentCount) {
-    mImagePropertiesPrivateData.CodeSegmentCountMax = ImageRecord->CodeSegmentCount;
-  }
-
-  SortImageRecord (&mImagePropertiesPrivateData.ImageRecordList);
-
-Finish:
-  if (EFI_ERROR (Status) && (ImageRecord != NULL)) {
-    DeleteImagePropertiesRecord (ImageRecord);
-  }
-
-  return;
 }
 
 /**
@@ -418,60 +491,6 @@ PublishMemoryAttributesTable (
   ASSERT_EFI_ERROR (Status);
 }
 
-/**
-  This function installs all SMM image record information.
-**/
-VOID
-SmmInstallImageRecord (
-  VOID
-  )
-{
-  EFI_STATUS                 Status;
-  UINTN                      NoHandles;
-  EFI_HANDLE                 *HandleBuffer;
-  EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage;
-  UINTN                      Index;
-  EFI_SMM_DRIVER_ENTRY       DriverEntry;
-
-  Status = SmmLocateHandleBuffer (
-             ByProtocol,
-             &gEfiLoadedImageProtocolGuid,
-             NULL,
-             &NoHandles,
-             &HandleBuffer
-             );
-  if (EFI_ERROR (Status)) {
-    return;
-  }
-
-  for (Index = 0; Index < NoHandles; Index++) {
-    Status = gSmst->SmmHandleProtocol (
-                      HandleBuffer[Index],
-                      &gEfiLoadedImageProtocolGuid,
-                      (VOID **)&LoadedImage
-                      );
-    if (EFI_ERROR (Status)) {
-      continue;
-    }
-
-    DEBUG ((DEBUG_VERBOSE, "LoadedImage - 0x%x 0x%x ", LoadedImage->ImageBase, LoadedImage->ImageSize));
-    {
-      VOID  *PdbPointer;
-      PdbPointer = PeCoffLoaderGetPdbPointer (LoadedImage->ImageBase);
-      if (PdbPointer != NULL) {
-        DEBUG ((DEBUG_VERBOSE, "(%a) ", PdbPointer));
-      }
-    }
-    DEBUG ((DEBUG_VERBOSE, "\n"));
-    ZeroMem (&DriverEntry, sizeof (DriverEntry));
-    DriverEntry.ImageBuffer  = (UINTN)LoadedImage->ImageBase;
-    DriverEntry.NumberOfPage = EFI_SIZE_TO_PAGES ((UINTN)LoadedImage->ImageSize);
-    SmmInsertImageRecord (&DriverEntry);
-  }
-
-  FreePool (HandleBuffer);
-}
-
 /**
   Install MemoryAttributesTable.
 
@@ -489,8 +508,6 @@ SmmInstallMemoryAttributesTable (
   IN EFI_HANDLE      Handle
   )
 {
-  SmmInstallImageRecord ();
-
   DEBUG ((DEBUG_VERBOSE, "SMM MemoryProtectionAttribute - 0x%016lx\n", mMemoryProtectionAttribute));
   if ((mMemoryProtectionAttribute & EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA) == 0) {
     return EFI_SUCCESS;
@@ -500,7 +517,7 @@ SmmInstallMemoryAttributesTable (
   if ( mImagePropertiesPrivateData.ImageRecordCount > 0) {
     DEBUG ((DEBUG_INFO, "SMM - Total Runtime Image Count - 0x%x\n", mImagePropertiesPrivateData.ImageRecordCount));
     DEBUG ((DEBUG_INFO, "SMM - Dump Runtime Image Records:\n"));
-    DumpImageRecords (&mImagePropertiesPrivateData.ImageRecordList);
+    DumpImageRecord ();
   }
 
   DEBUG_CODE_END ();

MdeModulePkg/Core/PiSmmCore/PiSmmCore.c

@@ -815,7 +815,7 @@ SmmCoreInstallLoadedImage (
   mSmmCoreLoadedImage->SystemTable  = gST;
 
   mSmmCoreLoadedImage->ImageBase     = (VOID *)(UINTN)gSmmCorePrivate->PiSmmCoreImageBase;
-  mSmmCoreLoadedImage->ImageSize     = gSmmCorePrivate->PiSmmCoreImageSize;
+  mSmmCoreLoadedImage->ImageSize     = UefiImageGetImageSize (&gSmmCorePrivate->PiSmmCoreImageContext);
   mSmmCoreLoadedImage->ImageCodeType = EfiRuntimeServicesCode;
   mSmmCoreLoadedImage->ImageDataType = EfiRuntimeServicesData;
 
@@ -847,13 +847,13 @@ SmmCoreInstallLoadedImage (
   mSmmCoreDriverEntry->SmmLoadedImage.SystemTable  = gST;
 
   mSmmCoreDriverEntry->SmmLoadedImage.ImageBase     = (VOID *)(UINTN)gSmmCorePrivate->PiSmmCoreImageBase;
-  mSmmCoreDriverEntry->SmmLoadedImage.ImageSize     = gSmmCorePrivate->PiSmmCoreImageSize;
+  mSmmCoreDriverEntry->SmmLoadedImage.ImageSize     = UefiImageGetImageSize (&gSmmCorePrivate->PiSmmCoreImageContext);
   mSmmCoreDriverEntry->SmmLoadedImage.ImageCodeType = EfiRuntimeServicesCode;
   mSmmCoreDriverEntry->SmmLoadedImage.ImageDataType = EfiRuntimeServicesData;
 
-  mSmmCoreDriverEntry->ImageEntryPoint = gSmmCorePrivate->PiSmmCoreEntryPoint;
+  mSmmCoreDriverEntry->ImageEntryPoint = UefiImageLoaderGetImageEntryPoint (&gSmmCorePrivate->PiSmmCoreImageContext);
   mSmmCoreDriverEntry->ImageBuffer     = gSmmCorePrivate->PiSmmCoreImageBase;
-  mSmmCoreDriverEntry->NumberOfPage    = EFI_SIZE_TO_PAGES ((UINTN)gSmmCorePrivate->PiSmmCoreImageSize);
+  mSmmCoreDriverEntry->NumberOfPage    = EFI_SIZE_TO_PAGES ((UINTN)UefiImageGetImageSize (&gSmmCorePrivate->PiSmmCoreImageContext));
 
   //
   // Create a new image handle in the SMM handle database for the SMM Driver
@@ -867,6 +867,21 @@ SmmCoreInstallLoadedImage (
                                           );
   ASSERT_EFI_ERROR (Status);
 
+  SmmInsertImageRecord (&mSmmCoreDriverEntry->SmmLoadedImage, &gSmmCorePrivate->PiSmmCoreImageContext);
+
+  //
+  // Create the aligned system table pointer structure that is used by external
+  // debuggers to locate the system table...  Also, install debug image info
+  // configuration table.
+  //
+  SmmInitializeDebugImageInfoTable ();
+  SmmNewDebugImageInfoEntry (
+    EFI_DEBUG_IMAGE_INFO_TYPE_NORMAL,
+    &mSmmCoreDriverEntry->SmmLoadedImage,
+    mSmmCoreDriverEntry->SmmImageHandle,
+    &gSmmCorePrivate->PiSmmCoreImageContext
+    );
+
   return;
 }
 

MdeModulePkg/Core/PiSmmCore/PiSmmCore.h

@@ -37,11 +37,11 @@
 #include <Guid/SmiHandlerProfile.h>
 #include <Guid/EndOfS3Resume.h>
 #include <Guid/S3SmmInitDone.h>
+#include <Guid/DebugImageInfoTable.h>
 
 #include <Library/BaseLib.h>
 #include <Library/BaseMemoryLib.h>
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
-#include <Library/PeCoffGetEntryPointLib.h>
 #include <Library/CacheMaintenanceLib.h>
 #include <Library/DebugLib.h>
 #include <Library/ReportStatusCodeLib.h>
@@ -55,6 +55,7 @@
 #include <Library/HobLib.h>
 #include <Library/SmmMemLib.h>
 #include <Library/SafeIntLib.h>
+#include <Library/SmmServicesTableLib.h>
 
 #include "PiSmmCorePrivateData.h"
 #include "HeapGuard.h"
@@ -1020,8 +1021,9 @@ SmramProfileInstallProtocol (
 **/
 EFI_STATUS
 RegisterSmramProfileImage (
-  IN EFI_SMM_DRIVER_ENTRY  *DriverEntry,
+  IN EFI_GUID                     *FileName,
-  IN BOOLEAN               RegisterToDxe
+  IN BOOLEAN                      RegisterToDxe,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT *ImageContext
   );
 
 /**
@@ -1038,8 +1040,10 @@ RegisterSmramProfileImage (
 **/
 EFI_STATUS
 UnregisterSmramProfileImage (
-  IN EFI_SMM_DRIVER_ENTRY  *DriverEntry,
+  IN EFI_GUID                           *FileName,
-  IN BOOLEAN               UnregisterToDxe
+  IN PHYSICAL_ADDRESS                   ImageBase,
+  IN UINT64                             ImageSize,
+  IN BOOLEAN                            UnregisterFromDxe
   );
 
 /**
@@ -1352,4 +1356,50 @@ SmmEntryPointMemoryManagementHook (
   VOID
   );
 
+/**
+  Creates and initializes the DebugImageInfo Table.  Also creates the configuration
+  table and registers it into the system table.
+
+  Note:
+    This function allocates memory, frees it, and then allocates memory at an
+    address within the initial allocation. Since this function is called early
+    in DXE core initialization (before drivers are dispatched), this should not
+    be a problem.
+
+**/
+VOID
+SmmInitializeDebugImageInfoTable (
+  VOID
+  );
+
+
+/**
+  Adds a new DebugImageInfo structure to the DebugImageInfo Table.  Re-Allocates
+  the table if it's not large enough to accomidate another entry.
+
+  @param  ImageInfoType  type of debug image information
+  @param  LoadedImage    pointer to the loaded image protocol for the image being
+                         loaded
+  @param  ImageHandle    image handle for the image being loaded
+
+**/
+VOID
+SmmNewDebugImageInfoEntry (
+  IN  UINT32                      ImageInfoType,
+  IN  EFI_LOADED_IMAGE_PROTOCOL   *LoadedImage,
+  IN  EFI_HANDLE                  ImageHandle,
+  IN  UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
+  );
+
+/**
+  Insert image record.
+
+  @param[in]  DriverEntry    Driver information
+**/
+VOID
+SmmInsertImageRecord (
+  IN EFI_LOADED_IMAGE_PROTOCOL        *LoadedImage,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
+  );
+
 #endif

MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf

@@ -37,6 +37,8 @@
   SmiHandlerProfile.c
   HeapGuard.c
   HeapGuard.h
+  MemoryAllocation.c
+  DebugImageInfo.c
 
 [Packages]
   MdePkg/MdePkg.dec
@@ -46,8 +48,7 @@
   UefiDriverEntryPoint
   BaseLib
   BaseMemoryLib
-  PeCoffLib
+  UefiImageLib
-  PeCoffGetEntryPointLib
   CacheMaintenanceLib
   DebugLib
   ReportStatusCodeLib
@@ -99,6 +100,7 @@
   gEfiMdeModulePkgTokenSpaceGuid.PcdHeapGuardPoolType                   ## CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdHeapGuardPropertyMask               ## CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdAcpiS3Enable                        ## CONSUMES
+  gEfiMdeModulePkgTokenSpaceGuid.PcdMaxEfiSystemTablePointerAddress     ## CONSUMES
 
 [Guids]
   gAprioriGuid                                  ## SOMETIMES_CONSUMES   ## File
@@ -120,6 +122,7 @@
   gSmiHandlerProfileGuid
   gEdkiiEndOfS3ResumeGuid ## SOMETIMES_PRODUCES ## GUID # Install protocol
   gEdkiiS3SmmInitDoneGuid ## SOMETIMES_PRODUCES ## GUID # Install protocol
+  gEfiDebugImageInfoTableGuid                   ## PRODUCES             ## SystemTable
 
 [UserExtensions.TianoCore."ExtraFiles"]
   PiSmmCoreExtra.uni

MdeModulePkg/Core/PiSmmCore/PiSmmCorePrivateData.h

@@ -22,6 +22,7 @@
 /// the SMM Entry Point enabling the use of SMM Mode.  In this case, the SMM Core
 /// should be notified again to dispatch more SMM Drivers using SMM Mode.
 ///
+#include "Library/UefiImageLib.h"
 #define COMM_BUFFER_SMM_DISPATCH_ERROR    0x00
 #define COMM_BUFFER_SMM_DISPATCH_SUCCESS  0x01
 #define COMM_BUFFER_SMM_DISPATCH_RESTART  0x02
@@ -112,8 +113,7 @@ typedef struct {
   EFI_STATUS               ReturnStatus;
 
   EFI_PHYSICAL_ADDRESS     PiSmmCoreImageBase;
-  UINT64                   PiSmmCoreImageSize;
+  UEFI_IMAGE_LOADER_IMAGE_CONTEXT    PiSmmCoreImageContext;
-  EFI_PHYSICAL_ADDRESS     PiSmmCoreEntryPoint;
 } SMM_CORE_PRIVATE_DATA;
 
 #endif

MdeModulePkg/Core/PiSmmCore/PiSmmIpl.c

@@ -23,7 +23,7 @@
 
 #include <Library/BaseLib.h>
 #include <Library/BaseMemoryLib.h>
-#include <Library/PeCoffLib.h>
+#include <Library/UefiImageLib.h>
 #include <Library/CacheMaintenanceLib.h>
 #include <Library/MemoryAllocationLib.h>
 #include <Library/DebugLib.h>
@@ -36,6 +36,8 @@
 #include <Library/ReportStatusCodeLib.h>
 #include "PiSmmCorePrivateData.h"
 #include <Library/SafeIntLib.h>
+#include "ProcessorBind.h"
+#include "Uefi/UefiBaseType.h"
 
 #define SMRAM_CAPABILITIES  (EFI_MEMORY_WB | EFI_MEMORY_UC)
 
@@ -913,91 +915,41 @@ SmmIplSetVirtualAddressNotify (
   @retval EFI_NOT_FOUND             The image has no assigned fixed loading address.
 **/
 EFI_STATUS
-GetPeCoffImageFixLoadingAssignedAddress (
+GetUefiImageFixLoadingAssignedAddress (
-  IN OUT PE_COFF_LOADER_IMAGE_CONTEXT  *ImageContext
+  IN OUT UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext,
+  OUT    EFI_PHYSICAL_ADDRESS          *LoadAddress
   )
 {
-  UINTN                            SectionHeaderOffset;
+  EFI_STATUS           Status;
-  EFI_STATUS                       Status;
+  UINT64               ValueInSectionHeader;
-  EFI_IMAGE_SECTION_HEADER         SectionHeader;
+  EFI_PHYSICAL_ADDRESS FixLoadingAddress;
-  EFI_IMAGE_OPTIONAL_HEADER_UNION  *ImgHdr;
+  UINT32               SizeOfImage;
-  EFI_PHYSICAL_ADDRESS             FixLoadingAddress;
+  EFI_PHYSICAL_ADDRESS SmramBase;
-  UINT16                           Index;
+  UINT64               SmmCodeSize;
-  UINTN                            Size;
-  UINT16                           NumberOfSections;
-  EFI_PHYSICAL_ADDRESS             SmramBase;
-  UINT64                           SmmCodeSize;
-  UINT64                           ValueInSectionHeader;
 
+  Status = UefiImageGetFixedAddress (ImageContext, &ValueInSectionHeader);
+  if (RETURN_ERROR (Status)) {
+    return Status;
+  }
   //
   // Build tool will calculate the smm code size and then patch the PcdLoadFixAddressSmmCodePageNumber
   //
-  SmmCodeSize = EFI_PAGES_TO_SIZE (PcdGet32 (PcdLoadFixAddressSmmCodePageNumber));
+  SmmCodeSize = EFI_PAGES_TO_SIZE (PcdGet32(PcdLoadFixAddressSmmCodePageNumber));
+  SmramBase = mLMFAConfigurationTable->SmramBase;
 
-  FixLoadingAddress = 0;
+  FixLoadingAddress = SmramBase + ValueInSectionHeader;
-  Status            = EFI_NOT_FOUND;
+  SizeOfImage = UefiImageGetImageSize (ImageContext);
-  SmramBase         = mLMFAConfigurationTable->SmramBase;
-  //
-  // Get PeHeader pointer
-  //
-  ImgHdr              = (EFI_IMAGE_OPTIONAL_HEADER_UNION *)((CHAR8 *)ImageContext->Handle + ImageContext->PeCoffHeaderOffset);
-  SectionHeaderOffset = ImageContext->PeCoffHeaderOffset +
-                        sizeof (UINT32) +
-                        sizeof (EFI_IMAGE_FILE_HEADER) +
-                        ImgHdr->Pe32.FileHeader.SizeOfOptionalHeader;
-  NumberOfSections = ImgHdr->Pe32.FileHeader.NumberOfSections;
 
-  //
+  if (SmramBase + SmmCodeSize >= FixLoadingAddress + SizeOfImage
-  // Get base address from the first section header that doesn't point to code section.
+   && SmramBase <= FixLoadingAddress) {
-  //
-  for (Index = 0; Index < NumberOfSections; Index++) {
     //
-    // Read section header from file
+    // The assigned address is valid. Return the specified loading address
     //
-    Size   = sizeof (EFI_IMAGE_SECTION_HEADER);
+    *LoadAddress = FixLoadingAddress;
-    Status = ImageContext->ImageRead (
+    Status = EFI_SUCCESS;
-                             ImageContext->Handle,
-                             SectionHeaderOffset,
-                             &Size,
-                             &SectionHeader
-                             );
-    if (EFI_ERROR (Status)) {
-      return Status;
-    }
-
-    Status = EFI_NOT_FOUND;
-
-    if ((SectionHeader.Characteristics & EFI_IMAGE_SCN_CNT_CODE) == 0) {
-      //
-      // Build tool saves the offset to SMRAM base as image base in PointerToRelocations & PointerToLineNumbers fields in the
-      // first section header that doesn't point to code section in image header. And there is an assumption that when the
-      // feature is enabled, if a module is assigned a loading address by tools, PointerToRelocations & PointerToLineNumbers
-      // fields should NOT be Zero, or else, these 2 fields should be set to Zero
-      //
-      ValueInSectionHeader = ReadUnaligned64 ((UINT64 *)&SectionHeader.PointerToRelocations);
-      if (ValueInSectionHeader != 0) {
-        //
-        // Found first section header that doesn't point to code section in which build tool saves the
-        // offset to SMRAM base as image base in PointerToRelocations & PointerToLineNumbers fields
-        //
-        FixLoadingAddress = (EFI_PHYSICAL_ADDRESS)(SmramBase + (INT64)ValueInSectionHeader);
-
-        if ((SmramBase + SmmCodeSize > FixLoadingAddress) && (SmramBase <=  FixLoadingAddress)) {
-          //
-          // The assigned address is valid. Return the specified loading address
-          //
-          ImageContext->ImageAddress = FixLoadingAddress;
-          Status                     = EFI_SUCCESS;
-        }
-      }
-
-      break;
-    }
-
-    SectionHeaderOffset += sizeof (EFI_IMAGE_SECTION_HEADER);
   }
 
-  DEBUG ((DEBUG_INFO|DEBUG_LOAD, "LOADING MODULE FIXED INFO: Loading module at fixed address %x, Status = %r \n", FixLoadingAddress, Status));
+  DEBUG ((EFI_D_INFO|EFI_D_LOAD, "LOADING MODULE FIXED INFO: Loading module at fixed address %x, Status = %r \n", FixLoadingAddress, Status));
   return Status;
 }
 
@@ -1024,9 +976,14 @@ ExecuteSmmCoreFromSmram (
   EFI_STATUS                    Status;
   VOID                          *SourceBuffer;
   UINTN                         SourceSize;
-  PE_COFF_LOADER_IMAGE_CONTEXT  ImageContext;
+  UINT32                        ImageSize;
+  UINT32                        ImageAlignment;
+  UINT32                        DestinationPages;
+  UINT32                        DestinationSize;
+  UINT32                        AlignSubtrahend;
   UINTN                         PageCount;
   EFI_IMAGE_ENTRY_POINT         EntryPoint;
+  EFI_PHYSICAL_ADDRESS          LoadAddress;
 
   //
   // Search all Firmware Volumes for a PE/COFF image in a file of type SMM_CORE
@@ -1043,20 +1000,26 @@ ExecuteSmmCoreFromSmram (
     return Status;
   }
 
-  //
-  // Initialize ImageContext
-  //
-  ImageContext.Handle    = SourceBuffer;
-  ImageContext.ImageRead = PeCoffLoaderImageReadFromMemory;
-
   //
   // Get information about the image being loaded
   //
-  Status = PeCoffLoaderGetImageInfo (&ImageContext);
+  Status = UefiImageInitializeContext (&gSmmCorePrivate->PiSmmCoreImageContext, SourceBuffer, (UINT32) SourceSize);
   if (EFI_ERROR (Status)) {
     return Status;
   }
 
+  //
+  // Stripped relocations are not supported for both fixed-address and dynamic
+  // loading.
+  //
+  if (UefiImageGetRelocsStripped (&gSmmCorePrivate->PiSmmCoreImageContext)) {
+    return EFI_UNSUPPORTED;
+  }
+
+  ImageSize        = UefiImageGetImageSize (&gSmmCorePrivate->PiSmmCoreImageContext);
+  DestinationPages = EFI_SIZE_TO_PAGES (ImageSize);
+  DestinationSize  = EFI_PAGES_TO_SIZE (DestinationPages);
+  ImageAlignment   = UefiImageGetSegmentAlignment (&gSmmCorePrivate->PiSmmCoreImageContext);
   //
   // if Loading module at Fixed Address feature is enabled, the SMM core driver will be loaded to
   // the address assigned by build tool.
@@ -1065,7 +1028,7 @@ ExecuteSmmCoreFromSmram (
     //
     // Get the fixed loading address assigned by Build tool
     //
-    Status = GetPeCoffImageFixLoadingAssignedAddress (&ImageContext);
+    Status = GetUefiImageFixLoadingAssignedAddress (&gSmmCorePrivate->PiSmmCoreImageContext, &LoadAddress);
     if (!EFI_ERROR (Status)) {
       //
       // Since the memory range to load SMM CORE will be cut out in SMM core, so no need to allocate and free this range
@@ -1081,7 +1044,11 @@ ExecuteSmmCoreFromSmram (
       // Allocate memory for the image being loaded from the EFI_SRAM_DESCRIPTOR
       // specified by SmramRange
       //
-      PageCount = (UINTN)EFI_SIZE_TO_PAGES ((UINTN)ImageContext.ImageSize + ImageContext.SectionAlignment);
+      AlignSubtrahend = ALIGN_VALUE_SUBTRAHEND (
+                          SmramRange->CpuStart + SmramRange->PhysicalSize,
+                          ImageAlignment
+                          );
+      PageCount = (UINTN)DestinationPages + (UINTN)EFI_SIZE_TO_PAGES ((UINTN)AlignSubtrahend);
 
       ASSERT ((SmramRange->PhysicalSize & EFI_PAGE_MASK) == 0);
       ASSERT (SmramRange->PhysicalSize > EFI_PAGES_TO_SIZE (PageCount));
@@ -1095,14 +1062,18 @@ ExecuteSmmCoreFromSmram (
       //
       // Align buffer on section boundary
       //
-      ImageContext.ImageAddress = SmramRangeSmmCore->CpuStart;
+      LoadAddress = SmramRangeSmmCore->CpuStart;
     }
   } else {
     //
     // Allocate memory for the image being loaded from the EFI_SRAM_DESCRIPTOR
     // specified by SmramRange
     //
-    PageCount = (UINTN)EFI_SIZE_TO_PAGES ((UINTN)ImageContext.ImageSize + ImageContext.SectionAlignment);
+    AlignSubtrahend = ALIGN_VALUE_SUBTRAHEND (
+                        SmramRange->CpuStart + SmramRange->PhysicalSize,
+                        ImageAlignment
+                        );
+    PageCount = (UINTN)DestinationPages + (UINTN)EFI_SIZE_TO_PAGES ((UINTN)AlignSubtrahend);
 
     ASSERT ((SmramRange->PhysicalSize & EFI_PAGE_MASK) == 0);
     ASSERT (SmramRange->PhysicalSize > EFI_PAGES_TO_SIZE (PageCount));
@@ -1116,50 +1087,34 @@ ExecuteSmmCoreFromSmram (
     //
     // Align buffer on section boundary
     //
-    ImageContext.ImageAddress = SmramRangeSmmCore->CpuStart;
+    LoadAddress = SmramRangeSmmCore->CpuStart;
   }
 
-  ImageContext.ImageAddress += ImageContext.SectionAlignment - 1;
-  ImageContext.ImageAddress &= ~((EFI_PHYSICAL_ADDRESS)ImageContext.SectionAlignment - 1);
-
   //
   // Print debug message showing SMM Core load address.
   //
-  DEBUG ((DEBUG_INFO, "SMM IPL loading SMM Core at SMRAM address %p\n", (VOID *)(UINTN)ImageContext.ImageAddress));
+  DEBUG ((DEBUG_INFO, "SMM IPL loading SMM Core at SMRAM destination %p\n", (VOID *)(UINTN)LoadAddress));
 
   //
   // Load the image to our new buffer
   //
-  Status = PeCoffLoaderLoadImage (&ImageContext);
+  Status = UefiImageLoadImageForExecution (&gSmmCorePrivate->PiSmmCoreImageContext, (VOID *)(UINTN)LoadAddress, DestinationSize, NULL, 0);
   if (!EFI_ERROR (Status)) {
+    LoadAddress = UefiImageLoaderGetImageAddress (&gSmmCorePrivate->PiSmmCoreImageContext);
     //
-    // Relocate the image in our new buffer
+    // Print debug message showing SMM Core entry point address.
     //
-    Status = PeCoffLoaderRelocateImage (&ImageContext);
+    DEBUG ((DEBUG_INFO, "SMM IPL calling SMM Core at SMRAM address %p\n", (VOID *)(UINTN)(UefiImageLoaderGetImageEntryPoint (&gSmmCorePrivate->PiSmmCoreImageContext))));
-    if (!EFI_ERROR (Status)) {
-      //
-      // Flush the instruction cache so the image data are written before we execute it
-      //
-      InvalidateInstructionCacheRange ((VOID *)(UINTN)ImageContext.ImageAddress, (UINTN)ImageContext.ImageSize);
 
-      //
+    gSmmCorePrivate->PiSmmCoreImageBase = LoadAddress;
-      // Print debug message showing SMM Core entry point address.
+    DEBUG ((DEBUG_INFO, "PiSmmCoreImageBase - 0x%016lx\n", gSmmCorePrivate->PiSmmCoreImageBase));
-      //
+    DEBUG ((DEBUG_INFO, "PiSmmCoreImageSize - 0x%016lx\n", UefiImageGetImageSize (&gSmmCorePrivate->PiSmmCoreImageContext)));
-      DEBUG ((DEBUG_INFO, "SMM IPL calling SMM Core at SMRAM address %p\n", (VOID *)(UINTN)ImageContext.EntryPoint));
 
-      gSmmCorePrivate->PiSmmCoreImageBase = ImageContext.ImageAddress;
+    //
-      gSmmCorePrivate->PiSmmCoreImageSize = ImageContext.ImageSize;
+    // Execute image
-      DEBUG ((DEBUG_INFO, "PiSmmCoreImageBase - 0x%016lx\n", gSmmCorePrivate->PiSmmCoreImageBase));
+    //
-      DEBUG ((DEBUG_INFO, "PiSmmCoreImageSize - 0x%016lx\n", gSmmCorePrivate->PiSmmCoreImageSize));
+    EntryPoint = (EFI_IMAGE_ENTRY_POINT)(UINTN)(UefiImageLoaderGetImageEntryPoint (&gSmmCorePrivate->PiSmmCoreImageContext));
-
+    Status = EntryPoint ((EFI_HANDLE)Context, gST);
-      gSmmCorePrivate->PiSmmCoreEntryPoint = ImageContext.EntryPoint;
-
-      //
-      // Execute image
-      //
-      EntryPoint = (EFI_IMAGE_ENTRY_POINT)(UINTN)ImageContext.EntryPoint;
-      Status     = EntryPoint ((EFI_HANDLE)Context, gST);
-    }
   }
 
   //

MdeModulePkg/Core/PiSmmCore/PiSmmIpl.inf

@@ -35,7 +35,7 @@
   UefiDriverEntryPoint
   BaseLib
   BaseMemoryLib
-  PeCoffLib
+  UefiImageLib
   CacheMaintenanceLib
   MemoryAllocationLib
   DebugLib

MdeModulePkg/Core/PiSmmCore/SmiHandlerProfile.c

@@ -17,7 +17,6 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Library/PrintLib.h>
 #include <Library/UefiLib.h>
 #include <Library/DevicePathLib.h>
-#include <Library/PeCoffGetEntryPointLib.h>
 #include <Protocol/LoadedImage.h>
 #include <Protocol/SmmAccess2.h>
 #include <Protocol/SmmReadyToLock.h>
@@ -74,7 +73,7 @@ BuildSmiHandlerProfileDatabase (
 
 **/
 RETURN_STATUS
-InternalPeCoffGetEntryPoint (
+InternalUefiImageGetEntryPoint (
   IN  VOID  *Pe32Data,
   OUT VOID  **EntryPoint
   );
@@ -264,7 +263,6 @@ GetSmmLoadedImage (
   CHAR16                     *PathStr;
   EFI_SMM_DRIVER_ENTRY       *LoadedImagePrivate;
   PHYSICAL_ADDRESS           EntryPoint;
-  VOID                       *EntryPointInImage;
   EFI_GUID                   Guid;
   CHAR8                      *PdbString;
   PHYSICAL_ADDRESS           RealImageBase;
@@ -324,15 +322,6 @@ GetSmmLoadedImage (
     RealImageBase      = (UINTN)LoadedImage->ImageBase;
     if (LoadedImagePrivate->Signature == EFI_SMM_DRIVER_ENTRY_SIGNATURE) {
       EntryPoint = LoadedImagePrivate->ImageEntryPoint;
-      if ((EntryPoint != 0) && ((EntryPoint < (UINTN)LoadedImage->ImageBase) || (EntryPoint >= ((UINTN)LoadedImage->ImageBase + LoadedImage->ImageSize)))) {
-        //
-        // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.
-        // So patch ImageBuffer here to align the EntryPoint.
-        //
-        Status = InternalPeCoffGetEntryPoint (LoadedImage->ImageBase, &EntryPointInImage);
-        ASSERT_EFI_ERROR (Status);
-        RealImageBase = (UINTN)LoadedImage->ImageBase + EntryPoint - (UINTN)EntryPointInImage;
-      }
     }
 
     DEBUG ((DEBUG_INFO, "(0x%lx - 0x%lx", RealImageBase, LoadedImage->ImageSize));
@@ -342,10 +331,11 @@ GetSmmLoadedImage (
 
     DEBUG ((DEBUG_INFO, ")\n"));
 
-    if (RealImageBase != 0) {
+    // FIXME:
-      PdbString = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)RealImageBase);
+    /*if (RealImageBase != 0) {
+      PdbString = UefiImageLoaderGetPdbPointer ((VOID *)(UINTN)RealImageBase);
       DEBUG ((DEBUG_INFO, "       pdb - %a\n", PdbString));
-    } else {
+    } else*/ {
       PdbString = NULL;
     }
 

MdeModulePkg/Core/PiSmmCore/SmramProfileRecord.c

@@ -128,8 +128,7 @@ EFIAPI
 SmramProfileProtocolRegisterImage (
   IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
   IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,
-  IN PHYSICAL_ADDRESS                   ImageBase,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT    *ImageContext,
-  IN UINT64                             ImageSize,
   IN EFI_FV_FILETYPE                    FileType
   );
 
@@ -249,110 +248,6 @@ GetSmramProfileContext (
   return mSmramProfileContextPtr;
 }
 
-/**
-  Retrieves and returns the Subsystem of a PE/COFF image that has been loaded into system memory.
-  If Pe32Data is NULL, then ASSERT().
-
-  @param Pe32Data   The pointer to the PE/COFF image that is loaded in system memory.
-
-  @return The Subsystem of the PE/COFF image.
-
-**/
-UINT16
-InternalPeCoffGetSubsystem (
-  IN VOID  *Pe32Data
-  )
-{
-  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
-  EFI_IMAGE_DOS_HEADER                 *DosHdr;
-  UINT16                               Magic;
-
-  ASSERT (Pe32Data != NULL);
-
-  DosHdr = (EFI_IMAGE_DOS_HEADER *)Pe32Data;
-  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
-    //
-    // DOS image header is present, so read the PE header after the DOS image header.
-    //
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINTN)Pe32Data + (UINTN)((DosHdr->e_lfanew) & 0x0ffff));
-  } else {
-    //
-    // DOS image header is not present, so PE header is at the image base.
-    //
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)Pe32Data;
-  }
-
-  if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {
-    return Hdr.Te->Subsystem;
-  } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {
-    Magic = Hdr.Pe32->OptionalHeader.Magic;
-    if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
-      return Hdr.Pe32->OptionalHeader.Subsystem;
-    } else if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
-      return Hdr.Pe32Plus->OptionalHeader.Subsystem;
-    }
-  }
-
-  return 0x0000;
-}
-
-/**
-  Retrieves and returns a pointer to the entry point to a PE/COFF image that has been loaded
-  into system memory with the PE/COFF Loader Library functions.
-
-  Retrieves the entry point to the PE/COFF image specified by Pe32Data and returns this entry
-  point in EntryPoint.  If the entry point could not be retrieved from the PE/COFF image, then
-  return RETURN_INVALID_PARAMETER.  Otherwise return RETURN_SUCCESS.
-  If Pe32Data is NULL, then ASSERT().
-  If EntryPoint is NULL, then ASSERT().
-
-  @param  Pe32Data                  The pointer to the PE/COFF image that is loaded in system memory.
-  @param  EntryPoint                The pointer to entry point to the PE/COFF image to return.
-
-  @retval RETURN_SUCCESS            EntryPoint was returned.
-  @retval RETURN_INVALID_PARAMETER  The entry point could not be found in the PE/COFF image.
-
-**/
-RETURN_STATUS
-InternalPeCoffGetEntryPoint (
-  IN  VOID  *Pe32Data,
-  OUT VOID  **EntryPoint
-  )
-{
-  EFI_IMAGE_DOS_HEADER                 *DosHdr;
-  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
-
-  ASSERT (Pe32Data   != NULL);
-  ASSERT (EntryPoint != NULL);
-
-  DosHdr = (EFI_IMAGE_DOS_HEADER *)Pe32Data;
-  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
-    //
-    // DOS image header is present, so read the PE header after the DOS image header.
-    //
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINTN)Pe32Data + (UINTN)((DosHdr->e_lfanew) & 0x0ffff));
-  } else {
-    //
-    // DOS image header is not present, so PE header is at the image base.
-    //
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)Pe32Data;
-  }
-
-  //
-  // Calculate the entry point relative to the start of the image.
-  // AddressOfEntryPoint is common for PE32 & PE32+
-  //
-  if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {
-    *EntryPoint = (VOID *)((UINTN)Pe32Data + (UINTN)(Hdr.Te->AddressOfEntryPoint & 0x0ffffffff) + sizeof (EFI_TE_IMAGE_HEADER) - Hdr.Te->StrippedSize);
-    return RETURN_SUCCESS;
-  } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {
-    *EntryPoint = (VOID *)((UINTN)Pe32Data + (UINTN)(Hdr.Pe32->OptionalHeader.AddressOfEntryPoint & 0x0ffffffff));
-    return RETURN_SUCCESS;
-  }
-
-  return RETURN_UNSUPPORTED;
-}
-
 /**
   Build driver info.
 
@@ -371,30 +266,26 @@ MEMORY_PROFILE_DRIVER_INFO_DATA *
 BuildDriverInfo (
   IN MEMORY_PROFILE_CONTEXT_DATA  *ContextData,
   IN EFI_GUID                     *FileName,
-  IN PHYSICAL_ADDRESS             ImageBase,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT   *ImageContext,
-  IN UINT64                       ImageSize,
+  IN EFI_PHYSICAL_ADDRESS           LoadAddress,
-  IN PHYSICAL_ADDRESS             EntryPoint,
-  IN UINT16                       ImageSubsystem,
   IN EFI_FV_FILETYPE              FileType
   )
 {
+  RETURN_STATUS                     PdbStatus;
   EFI_STATUS                       Status;
   MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
-  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
+  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
-  VOID                             *EntryPointInImage;
+  CONST CHAR8                       *PdbString;
-  CHAR8                            *PdbString;
+  UINT32                            PdbSize;
-  UINTN                            PdbSize;
+  UINTN                             PdbOccupiedSize;
-  UINTN                            PdbOccupiedSize;
 
-  PdbSize         = 0;
   PdbOccupiedSize = 0;
-  PdbString       = NULL;
+
-  if (ImageBase != 0) {
+  ASSERT (UefiImageLoaderGetImageAddress (ImageContext) != 0);
-    PdbString = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)ImageBase);
+
-    if (PdbString != NULL) {
+  PdbStatus = UefiImageGetSymbolsPath (ImageContext, &PdbString, &PdbSize);
-      PdbSize         = AsciiStrSize (PdbString);
+  if (!EFI_ERROR (PdbStatus)) {
-      PdbOccupiedSize = GET_OCCUPIED_SIZE (PdbSize, sizeof (UINT64));
+    PdbOccupiedSize = GET_OCCUPIED_SIZE (PdbSize, sizeof (UINT64));
-    }
   }
 
   //
@@ -422,19 +313,10 @@ BuildDriverInfo (
     CopyMem (&DriverInfo->FileName, FileName, sizeof (EFI_GUID));
   }
 
-  DriverInfo->ImageBase      = ImageBase;
+  DriverInfo->ImageBase      = LoadAddress;
-  DriverInfo->ImageSize      = ImageSize;
+  DriverInfo->ImageSize      = UefiImageGetImageSize (ImageContext);
-  DriverInfo->EntryPoint     = EntryPoint;
+  DriverInfo->EntryPoint     = UefiImageLoaderGetImageEntryPoint (ImageContext);
-  DriverInfo->ImageSubsystem = ImageSubsystem;
+  DriverInfo->ImageSubsystem = UefiImageGetSubsystem (ImageContext);
-  if ((EntryPoint != 0) && ((EntryPoint < ImageBase) || (EntryPoint >= (ImageBase + ImageSize)))) {
-    //
-    // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.
-    // So patch ImageBuffer here to align the EntryPoint.
-    //
-    Status = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)ImageBase, &EntryPointInImage);
-    ASSERT_EFI_ERROR (Status);
-    DriverInfo->ImageBase = ImageBase + EntryPoint - (PHYSICAL_ADDRESS)(UINTN)EntryPointInImage;
-  }
 
   DriverInfo->FileType          = FileType;
   DriverInfoData->AllocInfoList = (LIST_ENTRY *)(DriverInfoData + 1);
@@ -442,7 +324,7 @@ BuildDriverInfo (
   DriverInfo->CurrentUsage     = 0;
   DriverInfo->PeakUsage        = 0;
   DriverInfo->AllocRecordCount = 0;
-  if (PdbSize != 0) {
+  if (!RETURN_ERROR (PdbStatus)) {
     DriverInfo->PdbStringOffset = (UINT16)sizeof (MEMORY_PROFILE_DRIVER_INFO);
     DriverInfoData->PdbString   = (CHAR8 *)(DriverInfoData->AllocInfoList + 1);
     CopyMem (DriverInfoData->PdbString, PdbString, PdbSize);
@@ -470,8 +352,7 @@ BuildDriverInfo (
 VOID
 RegisterImageToDxe (
   IN EFI_GUID          *FileName,
-  IN PHYSICAL_ADDRESS  ImageBase,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT   *ImageContext,
-  IN UINT64            ImageSize,
   IN EFI_FV_FILETYPE   FileType
   )
 {
@@ -490,8 +371,7 @@ RegisterImageToDxe (
       Status = ProfileProtocol->RegisterImage (
                                   ProfileProtocol,
                                   (EFI_DEVICE_PATH_PROTOCOL *)FilePath,
-                                  ImageBase,
+                                  ImageContext,
-                                  ImageSize,
                                   FileType
                                   );
     }
@@ -610,7 +490,6 @@ RegisterSmmCore (
   )
 {
   MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData;
-  PHYSICAL_ADDRESS                   ImageBase;
   UINT8                              TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];
   MEDIA_FW_VOL_FILEPATH_DEVICE_PATH  *FilePath;
 
@@ -622,14 +501,11 @@ RegisterSmmCore (
     return FALSE;
   }
 
-  ImageBase      = gSmmCorePrivate->PiSmmCoreImageBase;
   DriverInfoData = BuildDriverInfo (
                      ContextData,
                      &gEfiCallerIdGuid,
-                     ImageBase,
+                     &gSmmCorePrivate->PiSmmCoreImageContext,
-                     gSmmCorePrivate->PiSmmCoreImageSize,
+                     gSmmCorePrivate->PiSmmCoreImageBase,
-                     gSmmCorePrivate->PiSmmCoreEntryPoint,
-                     InternalPeCoffGetSubsystem ((VOID *)(UINTN)ImageBase),
                      EFI_FV_FILETYPE_SMM_CORE
                      );
   if (DriverInfoData == NULL) {
@@ -652,8 +528,7 @@ SmramProfileInit (
 
   RegisterImageToDxe (
     &gEfiCallerIdGuid,
-    gSmmCorePrivate->PiSmmCoreImageBase,
+    &gSmmCorePrivate->PiSmmCoreImageContext,
-    gSmmCorePrivate->PiSmmCoreImageSize,
     EFI_FV_FILETYPE_SMM_CORE
     );
 
@@ -754,8 +629,9 @@ GetFileNameFromFilePath (
 **/
 EFI_STATUS
 RegisterSmramProfileImage (
-  IN EFI_SMM_DRIVER_ENTRY  *DriverEntry,
+  IN EFI_GUID                     *FileName,
-  IN BOOLEAN               RegisterToDxe
+  IN BOOLEAN                      RegisterToDxe,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT *ImageContext
   )
 {
   MEMORY_PROFILE_CONTEXT_DATA        *ContextData;
@@ -765,9 +641,8 @@ RegisterSmramProfileImage (
 
   if (RegisterToDxe) {
     RegisterImageToDxe (
-      &DriverEntry->FileName,
+      FileName,
-      DriverEntry->ImageBuffer,
+      ImageContext,
-      EFI_PAGES_TO_SIZE (DriverEntry->NumberOfPage),
       EFI_FV_FILETYPE_SMM
       );
   }
@@ -777,7 +652,7 @@ RegisterSmramProfileImage (
   }
 
   FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;
-  EfiInitializeFwVolDevicepathNode (FilePath, &DriverEntry->FileName);
+  EfiInitializeFwVolDevicepathNode (FilePath, FileName);
   SetDevicePathEndNode (FilePath + 1);
 
   if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *)FilePath)) {
@@ -791,11 +666,9 @@ RegisterSmramProfileImage (
 
   DriverInfoData = BuildDriverInfo (
                      ContextData,
-                     &DriverEntry->FileName,
+                     FileName,
-                     DriverEntry->ImageBuffer,
+                     ImageContext,
-                     EFI_PAGES_TO_SIZE (DriverEntry->NumberOfPage),
+                     UefiImageLoaderGetImageAddress (ImageContext),
-                     DriverEntry->ImageEntryPoint,
-                     InternalPeCoffGetSubsystem ((VOID *)(UINTN)DriverEntry->ImageBuffer),
                      EFI_FV_FILETYPE_SMM
                      );
   if (DriverInfoData == NULL) {
@@ -909,24 +782,22 @@ GetMemoryProfileDriverInfoFromAddress (
 **/
 EFI_STATUS
 UnregisterSmramProfileImage (
-  IN EFI_SMM_DRIVER_ENTRY  *DriverEntry,
+  IN EFI_GUID                           *FileName,
-  IN BOOLEAN               UnregisterFromDxe
+  IN PHYSICAL_ADDRESS                   ImageBase,
+  IN UINT64                             ImageSize,
+  IN BOOLEAN                            UnregisterFromDxe
   )
 {
-  EFI_STATUS                         Status;
   MEMORY_PROFILE_CONTEXT_DATA        *ContextData;
   MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData;
-  EFI_GUID                           *FileName;
-  PHYSICAL_ADDRESS                   ImageAddress;
-  VOID                               *EntryPointInImage;
   UINT8                              TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];
   MEDIA_FW_VOL_FILEPATH_DEVICE_PATH  *FilePath;
 
   if (UnregisterFromDxe) {
     UnregisterImageFromDxe (
-      &DriverEntry->FileName,
+      FileName,
-      DriverEntry->ImageBuffer,
+      ImageBase,
-      EFI_PAGES_TO_SIZE (DriverEntry->NumberOfPage)
+      ImageSize
       );
   }
 
@@ -935,7 +806,7 @@ UnregisterSmramProfileImage (
   }
 
   FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;
-  EfiInitializeFwVolDevicepathNode (FilePath, &DriverEntry->FileName);
+  EfiInitializeFwVolDevicepathNode (FilePath, FileName);
   SetDevicePathEndNode (FilePath + 1);
 
   if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *)FilePath)) {
@@ -948,24 +819,13 @@ UnregisterSmramProfileImage (
   }
 
   DriverInfoData = NULL;
-  FileName       = &DriverEntry->FileName;
-  ImageAddress   = DriverEntry->ImageBuffer;
-  if ((DriverEntry->ImageEntryPoint < ImageAddress) || (DriverEntry->ImageEntryPoint >= (ImageAddress + EFI_PAGES_TO_SIZE (DriverEntry->NumberOfPage)))) {
-    //
-    // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.
-    // So patch ImageAddress here to align the EntryPoint.
-    //
-    Status = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)ImageAddress, &EntryPointInImage);
-    ASSERT_EFI_ERROR (Status);
-    ImageAddress = ImageAddress + (UINTN)DriverEntry->ImageEntryPoint - (UINTN)EntryPointInImage;
-  }
 
   if (FileName != NULL) {
-    DriverInfoData = GetMemoryProfileDriverInfoByFileNameAndAddress (ContextData, FileName, ImageAddress);
+    DriverInfoData = GetMemoryProfileDriverInfoByFileNameAndAddress (ContextData, FileName, ImageBase);
   }
 
   if (DriverInfoData == NULL) {
-    DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageAddress);
+    DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageBase);
   }
 
   if (DriverInfoData == NULL) {
@@ -1989,29 +1849,20 @@ EFIAPI
 SmramProfileProtocolRegisterImage (
   IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
   IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,
-  IN PHYSICAL_ADDRESS                   ImageBase,
+  IN UEFI_IMAGE_LOADER_IMAGE_CONTEXT    *ImageContext,
-  IN UINT64                             ImageSize,
   IN EFI_FV_FILETYPE                    FileType
   )
 {
-  EFI_STATUS            Status;
+  EFI_GUID *FileName;
-  EFI_SMM_DRIVER_ENTRY  DriverEntry;
+  EFI_GUID ZeroGuid;
-  VOID                  *EntryPointInImage;
-  EFI_GUID              *Name;
 
-  ZeroMem (&DriverEntry, sizeof (DriverEntry));
+  FileName = GetFileNameFromFilePath (FilePath);
-  Name = GetFileNameFromFilePath (FilePath);
+  if (FileName == NULL) {
-  if (Name != NULL) {
+    ZeroMem (&ZeroGuid, sizeof (ZeroGuid));
-    CopyMem (&DriverEntry.FileName, Name, sizeof (EFI_GUID));
+    FileName = &ZeroGuid;
   }
 
-  DriverEntry.ImageBuffer  = ImageBase;
+  return RegisterSmramProfileImage (FileName, FALSE, ImageContext);
-  DriverEntry.NumberOfPage = EFI_SIZE_TO_PAGES ((UINTN)ImageSize);
-  Status                   = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)DriverEntry.ImageBuffer, &EntryPointInImage);
-  ASSERT_EFI_ERROR (Status);
-  DriverEntry.ImageEntryPoint = (PHYSICAL_ADDRESS)(UINTN)EntryPointInImage;
-
-  return RegisterSmramProfileImage (&DriverEntry, FALSE);
 }
 
 /**
@@ -2037,24 +1888,16 @@ SmramProfileProtocolUnregisterImage (
   IN UINT64                             ImageSize
   )
 {
-  EFI_STATUS            Status;
+  EFI_GUID *FileName;
-  EFI_SMM_DRIVER_ENTRY  DriverEntry;
+  EFI_GUID ZeroGuid;
-  VOID                  *EntryPointInImage;
-  EFI_GUID              *Name;
 
-  ZeroMem (&DriverEntry, sizeof (DriverEntry));
+  FileName = GetFileNameFromFilePath (FilePath);
-  Name = GetFileNameFromFilePath (FilePath);
+  if (FileName == NULL) {
-  if (Name != NULL) {
+    ZeroMem (&ZeroGuid, sizeof (ZeroGuid));
-    CopyMem (&DriverEntry.FileName, Name, sizeof (EFI_GUID));
+    FileName = &ZeroGuid;
   }
 
-  DriverEntry.ImageBuffer  = ImageBase;
+  return UnregisterSmramProfileImage (FileName, ImageBase, ImageSize, FALSE);
-  DriverEntry.NumberOfPage = EFI_SIZE_TO_PAGES ((UINTN)ImageSize);
-  Status                   = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)DriverEntry.ImageBuffer, &EntryPointInImage);
-  ASSERT_EFI_ERROR (Status);
-  DriverEntry.ImageEntryPoint = (PHYSICAL_ADDRESS)(UINTN)EntryPointInImage;
-
-  return UnregisterSmramProfileImage (&DriverEntry, FALSE);
 }
 
 /**