MdeModulePkg: Add Logic to Create/Delete Image Properties Records

Add logic to create and delete image properties records. Where
applicable, redirect existing code to use the new library.

Cc: Jian J Wang <jian.j.wang@intel.com>
Cc: Liming Gao <gaoliming@byosoft.com.cn>
Cc: Dandan Bi <dandan.bi@intel.com>
Cc: Jiaxin Wu <jiaxin.wu@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Signed-off-by: Taylor Beebe <taylor.d.beebe@gmail.com>
Reviewed-by: Liming Gao <gaoliming@byosoft.com.cn>

5 files changed, 280 insertions, 294 deletions

diff --git a/MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c b/MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c
index af6c26244c..993db28106 100644
--- a/MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c
+++ b/MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c
@@ -558,25 +558,6 @@ CoreGetMemoryMapWithSeparatedImageSection (
 //

 

 /**

-  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;

-  }

-}

-

-/**

   Insert image record.

 

   @param  RuntimeImage    Runtime image information

@@ -586,20 +567,12 @@ InsertImageRecord (
   IN EFI_RUNTIME_IMAGE_ENTRY  *RuntimeImage

   )

 {

-  VOID                                  *ImageAddress;

-  EFI_IMAGE_DOS_HEADER                  *DosHdr;

-  UINT32                                PeCoffHeaderOffset;

-  UINT32                                SectionAlignment;

-  EFI_IMAGE_SECTION_HEADER              *Section;

-  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION   Hdr;

-  UINT8                                 *Name;

-  UINTN                                 Index;

-  IMAGE_PROPERTIES_RECORD               *ImageRecord;

-  CHAR8                                 *PdbPointer;

-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;

+  EFI_STATUS               Status;

+  IMAGE_PROPERTIES_RECORD  *ImageRecord;

+  CHAR8                    *PdbPointer;

+  UINT32                   RequiredAlignment;

 

   DEBUG ((DEBUG_VERBOSE, "InsertImageRecord - 0x%x\n", RuntimeImage));

-  DEBUG ((DEBUG_VERBOSE, "InsertImageRecord - 0x%016lx - 0x%016lx\n", (EFI_PHYSICAL_ADDRESS)(UINTN)RuntimeImage->ImageBase, RuntimeImage->ImageSize));

 

   if (mMemoryAttributesTableEndOfDxe) {

     DEBUG ((DEBUG_INFO, "Do not insert runtime image record after EndOfDxe\n"));

@@ -611,139 +584,48 @@ InsertImageRecord (
     return;

   }

 

-  ImageRecord->Signature = IMAGE_PROPERTIES_RECORD_SIGNATURE;

-

-  DEBUG ((DEBUG_VERBOSE, "ImageRecordCount - 0x%x\n", mImagePropertiesPrivateData.ImageRecordCount));

-

-  //

-  // Step 1: record whole region

-  //

-  ImageRecord->ImageBase = (EFI_PHYSICAL_ADDRESS)(UINTN)RuntimeImage->ImageBase;

-  ImageRecord->ImageSize = RuntimeImage->ImageSize;

-

-  ImageAddress = RuntimeImage->ImageBase;

+  InitializeListHead (&ImageRecord->Link);

+  InitializeListHead (&ImageRecord->CodeSegmentList);

 

-  PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)ImageAddress);

+  PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)RuntimeImage->ImageBase);

   if (PdbPointer != NULL) {

     DEBUG ((DEBUG_VERBOSE, "  Image - %a\n", PdbPointer));

   }

 

-  //

-  // Check PE/COFF image

-  //

-  DosHdr             = (EFI_IMAGE_DOS_HEADER *)(UINTN)ImageAddress;

-  PeCoffHeaderOffset = 0;

-  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {

-    PeCoffHeaderOffset = DosHdr->e_lfanew;

-  }

-

-  Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINT8 *)(UINTN)ImageAddress + PeCoffHeaderOffset);

-  if (Hdr.Pe32->Signature != EFI_IMAGE_NT_SIGNATURE) {

-    DEBUG ((DEBUG_VERBOSE, "Hdr.Pe32->Signature invalid - 0x%x\n", Hdr.Pe32->Signature));

-    // It might be image in SMM.

-    goto Finish;

-  }

-

-  //

-  // Get SectionAlignment

-  //

-  if (Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {

-    SectionAlignment = Hdr.Pe32->OptionalHeader.SectionAlignment;

-  } else {

-    SectionAlignment = Hdr.Pe32Plus->OptionalHeader.SectionAlignment;

-  }

-

-  SetMemoryAttributesTableSectionAlignment (SectionAlignment);

-  if ((SectionAlignment & (RUNTIME_PAGE_ALLOCATION_GRANULARITY - 1)) != 0) {

-    DEBUG ((

-      DEBUG_WARN,

-      "!!!!!!!!  InsertImageRecord - Section Alignment(0x%x) is not %dK  !!!!!!!!\n",

-      SectionAlignment,

-      RUNTIME_PAGE_ALLOCATION_GRANULARITY >> 10

-      ));

-    PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)ImageAddress);

-    if (PdbPointer != NULL) {

-      DEBUG ((DEBUG_WARN, "!!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));

+  RequiredAlignment = RUNTIME_PAGE_ALLOCATION_GRANULARITY;

+  Status            = CreateImagePropertiesRecord (

+                        RuntimeImage->ImageBase,

+                        RuntimeImage->ImageSize,

+                        &RequiredAlignment,

+                        ImageRecord

+                        );

+

+  if (EFI_ERROR (Status)) {

+    if (Status == EFI_ABORTED) {

+      mMemoryAttributesTableEnable = FALSE;

     }

 

+    Status = EFI_ABORTED;

     goto Finish;

   }

 

-  Section = (EFI_IMAGE_SECTION_HEADER *)(

-                                         (UINT8 *)(UINTN)ImageAddress +

-                                         PeCoffHeaderOffset +

-                                         sizeof (UINT32) +

-                                         sizeof (EFI_IMAGE_FILE_HEADER) +

-                                         Hdr.Pe32->FileHeader.SizeOfOptionalHeader

-                                         );

-  ImageRecord->CodeSegmentCount = 0;

-  InitializeListHead (&ImageRecord->CodeSegmentList);

-  for (Index = 0; Index < Hdr.Pe32->FileHeader.NumberOfSections; Index++) {

-    Name = Section[Index].Name;

-    DEBUG ((

-      DEBUG_VERBOSE,

-      "  Section - '%c%c%c%c%c%c%c%c'\n",

-      Name[0],

-      Name[1],

-      Name[2],

-      Name[3],

-      Name[4],

-      Name[5],

-      Name[6],

-      Name[7]

-      ));

-

-    if ((Section[Index].Characteristics & EFI_IMAGE_SCN_CNT_CODE) != 0) {

-      DEBUG ((DEBUG_VERBOSE, "  VirtualSize          - 0x%08x\n", Section[Index].Misc.VirtualSize));

-      DEBUG ((DEBUG_VERBOSE, "  VirtualAddress       - 0x%08x\n", Section[Index].VirtualAddress));

-      DEBUG ((DEBUG_VERBOSE, "  SizeOfRawData        - 0x%08x\n", Section[Index].SizeOfRawData));

-      DEBUG ((DEBUG_VERBOSE, "  PointerToRawData     - 0x%08x\n", Section[Index].PointerToRawData));

-      DEBUG ((DEBUG_VERBOSE, "  PointerToRelocations - 0x%08x\n", Section[Index].PointerToRelocations));

-      DEBUG ((DEBUG_VERBOSE, "  PointerToLinenumbers - 0x%08x\n", Section[Index].PointerToLinenumbers));

-      DEBUG ((DEBUG_VERBOSE, "  NumberOfRelocations  - 0x%08x\n", Section[Index].NumberOfRelocations));

-      DEBUG ((DEBUG_VERBOSE, "  NumberOfLinenumbers  - 0x%08x\n", Section[Index].NumberOfLinenumbers));

-      DEBUG ((DEBUG_VERBOSE, "  Characteristics      - 0x%08x\n", Section[Index].Characteristics));

-

-      //

-      // Step 2: record code section

-      //

-      ImageRecordCodeSection = AllocatePool (sizeof (*ImageRecordCodeSection));

-      if (ImageRecordCodeSection == NULL) {

-        return;

-      }

-

-      ImageRecordCodeSection->Signature = IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE;

-

-      ImageRecordCodeSection->CodeSegmentBase = (UINTN)ImageAddress + Section[Index].VirtualAddress;

-      ImageRecordCodeSection->CodeSegmentSize = Section[Index].SizeOfRawData;

-

-      DEBUG ((DEBUG_VERBOSE, "ImageCode: 0x%016lx - 0x%016lx\n", ImageRecordCodeSection->CodeSegmentBase, ImageRecordCodeSection->CodeSegmentSize));

-

-      InsertTailList (&ImageRecord->CodeSegmentList, &ImageRecordCodeSection->Link);

-      ImageRecord->CodeSegmentCount++;

-    }

-  }

-

   if (ImageRecord->CodeSegmentCount == 0) {

-    SetMemoryAttributesTableSectionAlignment (1);

+    mMemoryAttributesTableEnable = FALSE;

     DEBUG ((DEBUG_ERROR, "!!!!!!!!  InsertImageRecord - CodeSegmentCount is 0  !!!!!!!!\n"));

-    PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)ImageAddress);

     if (PdbPointer != NULL) {

       DEBUG ((DEBUG_ERROR, "!!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));

     }

 

+    Status = EFI_ABORTED;

     goto Finish;

   }

 

   //

-  // Final

-  //

-  SortImageRecordCodeSection (ImageRecord);

-  //

   // Check overlap all section in ImageBase/Size

   //

   if (!IsImageRecordCodeSectionValid (ImageRecord)) {

     DEBUG ((DEBUG_ERROR, "IsImageRecordCodeSectionValid - FAIL\n"));

+    Status = EFI_ABORTED;

     goto Finish;

   }

 

@@ -757,6 +639,10 @@ InsertImageRecord (
   SortImageRecord (&mImagePropertiesPrivateData.ImageRecordList);

 

 Finish:

+  if (EFI_ERROR (Status) && (ImageRecord != NULL)) {

+    DeleteImagePropertiesRecord (ImageRecord);

+  }

+

   return;

 }

 

@@ -770,9 +656,7 @@ RemoveImageRecord (
   IN EFI_RUNTIME_IMAGE_ENTRY  *RuntimeImage

   )

 {

-  IMAGE_PROPERTIES_RECORD               *ImageRecord;

-  LIST_ENTRY                            *CodeSegmentListHead;

-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;

+  IMAGE_PROPERTIES_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));

@@ -788,19 +672,7 @@ RemoveImageRecord (
     return;

   }

 

-  CodeSegmentListHead = &ImageRecord->CodeSegmentList;

-  while (!IsListEmpty (CodeSegmentListHead)) {

-    ImageRecordCodeSection = CR (

-                               CodeSegmentListHead->ForwardLink,

-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,

-                               Link,

-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE

-                               );

-    RemoveEntryList (&ImageRecordCodeSection->Link);

-    FreePool (ImageRecordCodeSection);

-  }

+  DeleteImagePropertiesRecord (ImageRecord);

 

-  RemoveEntryList (&ImageRecord->Link);

-  FreePool (ImageRecord);

   mImagePropertiesPrivateData.ImageRecordCount--;

 }

diff --git a/MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c b/MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c
index 2e4aaddef4..03de9b2c5f 100644
--- a/MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c
+++ b/MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c
@@ -252,25 +252,6 @@ SmmCoreGetMemoryMapMemoryAttributesTable (
 //

 

 /**

-  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);

-  }

-}

-

-/**

   Insert image record.

 

   @param[in]  DriverEntry    Driver information

@@ -280,158 +261,61 @@ SmmInsertImageRecord (
   IN EFI_SMM_DRIVER_ENTRY  *DriverEntry

   )

 {

-  VOID                                  *ImageAddress;

-  EFI_IMAGE_DOS_HEADER                  *DosHdr;

-  UINT32                                PeCoffHeaderOffset;

-  UINT32                                SectionAlignment;

-  EFI_IMAGE_SECTION_HEADER              *Section;

-  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION   Hdr;

-  UINT8                                 *Name;

-  UINTN                                 Index;

-  IMAGE_PROPERTIES_RECORD               *ImageRecord;

-  CHAR8                                 *PdbPointer;

-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;

+  EFI_STATUS               Status;

+  IMAGE_PROPERTIES_RECORD  *ImageRecord;

+  CHAR8                    *PdbPointer;

+  UINT32                   RequiredAlignment;

 

   DEBUG ((DEBUG_VERBOSE, "SMM InsertImageRecord - 0x%x\n", DriverEntry));

-  DEBUG ((DEBUG_VERBOSE, "SMM InsertImageRecord - 0x%016lx - 0x%08x\n", DriverEntry->ImageBuffer, DriverEntry->NumberOfPage));

 

   ImageRecord = AllocatePool (sizeof (*ImageRecord));

   if (ImageRecord == NULL) {

     return;

   }

 

-  ImageRecord->Signature = IMAGE_PROPERTIES_RECORD_SIGNATURE;

-

-  DEBUG ((DEBUG_VERBOSE, "SMM ImageRecordCount - 0x%x\n", mImagePropertiesPrivateData.ImageRecordCount));

-

-  //

-  // Step 1: record whole region

-  //

-  ImageRecord->ImageBase = DriverEntry->ImageBuffer;

-  ImageRecord->ImageSize = LShiftU64 (DriverEntry->NumberOfPage, EFI_PAGE_SHIFT);

-

-  ImageAddress = (VOID *)(UINTN)DriverEntry->ImageBuffer;

+  InitializeListHead (&ImageRecord->Link);

+  InitializeListHead (&ImageRecord->CodeSegmentList);

 

-  PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)ImageAddress);

+  PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)DriverEntry->ImageBuffer);

   if (PdbPointer != NULL) {

     DEBUG ((DEBUG_VERBOSE, "SMM   Image - %a\n", PdbPointer));

   }

 

-  //

-  // Check PE/COFF image

-  //

-  DosHdr             = (EFI_IMAGE_DOS_HEADER *)(UINTN)ImageAddress;

-  PeCoffHeaderOffset = 0;

-  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {

-    PeCoffHeaderOffset = DosHdr->e_lfanew;

-  }

-

-  Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINT8 *)(UINTN)ImageAddress + PeCoffHeaderOffset);

-  if (Hdr.Pe32->Signature != EFI_IMAGE_NT_SIGNATURE) {

-    DEBUG ((DEBUG_VERBOSE, "SMM Hdr.Pe32->Signature invalid - 0x%x\n", Hdr.Pe32->Signature));

-    goto Finish;

-  }

+  RequiredAlignment = RUNTIME_PAGE_ALLOCATION_GRANULARITY;

+  Status            = CreateImagePropertiesRecord (

+                        (VOID *)(UINTN)DriverEntry->ImageBuffer,

+                        LShiftU64 (DriverEntry->NumberOfPage, EFI_PAGE_SHIFT),

+                        &RequiredAlignment,

+                        ImageRecord

+                        );

 

-  //

-  // Get SectionAlignment

-  //

-  if (Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {

-    SectionAlignment = Hdr.Pe32->OptionalHeader.SectionAlignment;

-  } else {

-    SectionAlignment = Hdr.Pe32Plus->OptionalHeader.SectionAlignment;

-  }

-

-  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

-      ));

-    PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)ImageAddress);

-    if (PdbPointer != NULL) {

-      DEBUG ((DEBUG_WARN, "SMM !!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));

+  if (EFI_ERROR (Status)) {

+    if (Status == EFI_ABORTED) {

+      mMemoryProtectionAttribute &=

+        ~((UINT64)EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA);

     }

 

     goto Finish;

   }

 

-  Section = (EFI_IMAGE_SECTION_HEADER *)(

-                                         (UINT8 *)(UINTN)ImageAddress +

-                                         PeCoffHeaderOffset +

-                                         sizeof (UINT32) +

-                                         sizeof (EFI_IMAGE_FILE_HEADER) +

-                                         Hdr.Pe32->FileHeader.SizeOfOptionalHeader

-                                         );

-  ImageRecord->CodeSegmentCount = 0;

-  InitializeListHead (&ImageRecord->CodeSegmentList);

-  for (Index = 0; Index < Hdr.Pe32->FileHeader.NumberOfSections; Index++) {

-    Name = Section[Index].Name;

-    DEBUG ((

-      DEBUG_VERBOSE,

-      "SMM   Section - '%c%c%c%c%c%c%c%c'\n",

-      Name[0],

-      Name[1],

-      Name[2],

-      Name[3],

-      Name[4],

-      Name[5],

-      Name[6],

-      Name[7]

-      ));

-

-    if ((Section[Index].Characteristics & EFI_IMAGE_SCN_CNT_CODE) != 0) {

-      DEBUG ((DEBUG_VERBOSE, "SMM   VirtualSize          - 0x%08x\n", Section[Index].Misc.VirtualSize));

-      DEBUG ((DEBUG_VERBOSE, "SMM   VirtualAddress       - 0x%08x\n", Section[Index].VirtualAddress));

-      DEBUG ((DEBUG_VERBOSE, "SMM   SizeOfRawData        - 0x%08x\n", Section[Index].SizeOfRawData));

-      DEBUG ((DEBUG_VERBOSE, "SMM   PointerToRawData     - 0x%08x\n", Section[Index].PointerToRawData));

-      DEBUG ((DEBUG_VERBOSE, "SMM   PointerToRelocations - 0x%08x\n", Section[Index].PointerToRelocations));

-      DEBUG ((DEBUG_VERBOSE, "SMM   PointerToLinenumbers - 0x%08x\n", Section[Index].PointerToLinenumbers));

-      DEBUG ((DEBUG_VERBOSE, "SMM   NumberOfRelocations  - 0x%08x\n", Section[Index].NumberOfRelocations));

-      DEBUG ((DEBUG_VERBOSE, "SMM   NumberOfLinenumbers  - 0x%08x\n", Section[Index].NumberOfLinenumbers));

-      DEBUG ((DEBUG_VERBOSE, "SMM   Characteristics      - 0x%08x\n", Section[Index].Characteristics));

-

-      //

-      // Step 2: record code section

-      //

-      ImageRecordCodeSection = AllocatePool (sizeof (*ImageRecordCodeSection));

-      if (ImageRecordCodeSection == NULL) {

-        return;

-      }

-

-      ImageRecordCodeSection->Signature = IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE;

-

-      ImageRecordCodeSection->CodeSegmentBase = (UINTN)ImageAddress + Section[Index].VirtualAddress;

-      ImageRecordCodeSection->CodeSegmentSize = Section[Index].SizeOfRawData;

-

-      DEBUG ((DEBUG_VERBOSE, "SMM ImageCode: 0x%016lx - 0x%016lx\n", ImageRecordCodeSection->CodeSegmentBase, ImageRecordCodeSection->CodeSegmentSize));

-

-      InsertTailList (&ImageRecord->CodeSegmentList, &ImageRecordCodeSection->Link);

-      ImageRecord->CodeSegmentCount++;

-    }

-  }

-

   if (ImageRecord->CodeSegmentCount == 0) {

-    SetMemoryAttributesTableSectionAlignment (1);

+    mMemoryProtectionAttribute &=

+      ~((UINT64)EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA);

     DEBUG ((DEBUG_ERROR, "SMM !!!!!!!!  InsertImageRecord - CodeSegmentCount is 0  !!!!!!!!\n"));

-    PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)ImageAddress);

     if (PdbPointer != NULL) {

       DEBUG ((DEBUG_ERROR, "SMM !!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));

     }

 

+    Status = EFI_ABORTED;

     goto Finish;

   }

 

   //

-  // Final

-  //

-  SortImageRecordCodeSection (ImageRecord);

-  //

   // Check overlap all section in ImageBase/Size

   //

   if (!IsImageRecordCodeSectionValid (ImageRecord)) {

     DEBUG ((DEBUG_ERROR, "SMM IsImageRecordCodeSectionValid - FAIL\n"));

+    Status = EFI_ABORTED;

     goto Finish;

   }

 

@@ -445,6 +329,10 @@ SmmInsertImageRecord (
   SortImageRecord (&mImagePropertiesPrivateData.ImageRecordList);

 

 Finish:

+  if (EFI_ERROR (Status) && (ImageRecord != NULL)) {

+    DeleteImagePropertiesRecord (ImageRecord);

+  }

+

   return;

 }

 

diff --git a/MdeModulePkg/Include/Library/ImagePropertiesRecordLib.h b/MdeModulePkg/Include/Library/ImagePropertiesRecordLib.h
index e3f569ab03..5090a52153 100644
--- a/MdeModulePkg/Include/Library/ImagePropertiesRecordLib.h
+++ b/MdeModulePkg/Include/Library/ImagePropertiesRecordLib.h
@@ -192,4 +192,43 @@ DumpImageRecord (
   IN LIST_ENTRY  *ImageRecordList

   );

 

+/**

+  Creates an IMAGE_PROPERTIES_RECORD from a loaded PE image. The PE/COFF header will be found

+  and parsed to determine the number of code segments and their base addresses and sizes.

+

+  @param[in]      ImageBase               Base of the PE image

+  @param[in]      ImageSize               Size of the PE image

+  @param[in]      RequiredAlignment       If non-NULL, the alignment specified in the PE/COFF header

+                                          will be compared against this value.

+  @param[out]     ImageRecord             On out, a populated image properties record

+

+  @retval     EFI_INVALID_PARAMETER   This function ImageBase or ImageRecord was NULL, or the

+                                      image located at ImageBase was not a valid PE/COFF image

+  @retval     EFI_OUT_OF_RESOURCES    Failure to Allocate()

+  @retval     EFI_ABORTED             The input Alignment was non-NULL and did not match the

+                                      alignment specified in the PE/COFF header

+  @retval     EFI_SUCCESS             The image properties record was successfully created

+**/

+EFI_STATUS

+EFIAPI

+CreateImagePropertiesRecord (

+  IN  CONST   VOID                     *ImageBase,

+  IN  CONST   UINT64                   ImageSize,

+  IN  CONST   UINT32                   *Alignment OPTIONAL,

+  OUT         IMAGE_PROPERTIES_RECORD  *ImageRecord

+  );

+

+/**

+  Deleted an image properties record. The function will also call

+  RemoveEntryList() on each code segment and the input ImageRecord before

+  freeing each pool.

+

+  @param[in]      ImageRecord             The IMAGE_PROPERTIES_RECORD to delete

+**/

+VOID

+EFIAPI

+DeleteImagePropertiesRecord (

+  IN  IMAGE_PROPERTIES_RECORD  *ImageRecord

+  );

+

 #endif

diff --git a/MdeModulePkg/Library/ImagePropertiesRecordLib/ImagePropertiesRecordLib.c b/MdeModulePkg/Library/ImagePropertiesRecordLib/ImagePropertiesRecordLib.c
index 9b296aa457..6c5eb1dc31 100644
--- a/MdeModulePkg/Library/ImagePropertiesRecordLib/ImagePropertiesRecordLib.c
+++ b/MdeModulePkg/Library/ImagePropertiesRecordLib/ImagePropertiesRecordLib.c
@@ -13,6 +13,7 @@
 #include <Library/BaseLib.h>

 #include <Library/BaseMemoryLib.h>

 #include <Library/DebugLib.h>

+#include <Library/MemoryAllocationLib.h>

 #include <Library/ImagePropertiesRecordLib.h>

 

 #define PREVIOUS_MEMORY_DESCRIPTOR(MemoryDescriptor, Size) \

@@ -858,3 +859,188 @@ FindImageRecord (
 

   return NULL;

 }

+

+/**

+  Creates an IMAGE_PROPERTIES_RECORD from a loaded PE image. The PE/COFF header will be found

+  and parsed to determine the number of code segments and their base addresses and sizes.

+

+  @param[in]      ImageBase               Base of the PE image

+  @param[in]      ImageSize               Size of the PE image

+  @param[in]      RequiredAlignment       If non-NULL, the alignment specified in the PE/COFF header

+                                          will be compared against this value.

+  @param[out]     ImageRecord             On out, a populated image properties record

+

+  @retval     EFI_INVALID_PARAMETER   This function ImageBase or ImageRecord was NULL, or the

+                                      image located at ImageBase was not a valid PE/COFF image

+  @retval     EFI_OUT_OF_RESOURCES    Failure to Allocate()

+  @retval     EFI_ABORTED             The input Alignment was non-NULL and did not match the

+                                      alignment specified in the PE/COFF header

+  @retval     EFI_SUCCESS             The image properties record was successfully created

+**/

+EFI_STATUS

+EFIAPI

+CreateImagePropertiesRecord (

+  IN  CONST   VOID                     *ImageBase,

+  IN  CONST   UINT64                   ImageSize,

+  IN  CONST   UINT32                   *RequiredAlignment OPTIONAL,

+  OUT         IMAGE_PROPERTIES_RECORD  *ImageRecord

+  )

+{

+  EFI_STATUS                            Status;

+  EFI_IMAGE_DOS_HEADER                  *DosHdr;

+  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION   Hdr;

+  EFI_IMAGE_SECTION_HEADER              *Section;

+  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;

+  UINTN                                 Index;

+  UINT8                                 *Name;

+  UINT32                                SectionAlignment;

+  UINT32                                PeCoffHeaderOffset;

+

+  if ((ImageRecord == NULL) || (ImageBase == NULL)) {

+    return EFI_INVALID_PARAMETER;

+  }

+

+  DEBUG ((

+    DEBUG_VERBOSE,

+    "Creating Image Properties Record: 0x%016lx - 0x%016lx\n",

+    (EFI_PHYSICAL_ADDRESS)(UINTN)ImageBase,

+    ImageSize

+    ));

+

+  //

+  // Step 1: record whole region

+  //

+  Status                        = EFI_SUCCESS;

+  ImageRecord->Signature        = IMAGE_PROPERTIES_RECORD_SIGNATURE;

+  ImageRecord->ImageBase        = (EFI_PHYSICAL_ADDRESS)(UINTN)ImageBase;

+  ImageRecord->ImageSize        = ImageSize;

+  ImageRecord->CodeSegmentCount = 0;

+  InitializeListHead (&ImageRecord->Link);

+  InitializeListHead (&ImageRecord->CodeSegmentList);

+

+  // Check PE/COFF image

+  DosHdr             = (EFI_IMAGE_DOS_HEADER *)(UINTN)ImageBase;

+  PeCoffHeaderOffset = 0;

+  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {

+    PeCoffHeaderOffset = DosHdr->e_lfanew;

+  }

+

+  Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINT8 *)(UINTN)ImageBase + PeCoffHeaderOffset);

+  if (Hdr.Pe32->Signature != EFI_IMAGE_NT_SIGNATURE) {

+    DEBUG ((DEBUG_VERBOSE, "Hdr.Pe32->Signature invalid - 0x%x\n", Hdr.Pe32->Signature));

+    return EFI_INVALID_PARAMETER;

+  }

+

+  // Get SectionAlignment

+  if (Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {

+    SectionAlignment = Hdr.Pe32->OptionalHeader.SectionAlignment;

+  } else {

+    SectionAlignment = Hdr.Pe32Plus->OptionalHeader.SectionAlignment;

+  }

+

+  // Check RequiredAlignment

+  if ((RequiredAlignment != NULL) && ((SectionAlignment & (*RequiredAlignment - 1)) != 0)) {

+    DEBUG ((

+      DEBUG_WARN,

+      "!!!!!!!!  Image Section Alignment(0x%x) does not match Required Alignment (0x%x)  !!!!!!!!\n",

+      SectionAlignment,

+      *RequiredAlignment

+      ));

+

+    return EFI_ABORTED;

+  }

+

+  Section = (EFI_IMAGE_SECTION_HEADER *)(

+                                         (UINT8 *)(UINTN)ImageBase +

+                                         PeCoffHeaderOffset +

+                                         sizeof (UINT32) +

+                                         sizeof (EFI_IMAGE_FILE_HEADER) +

+                                         Hdr.Pe32->FileHeader.SizeOfOptionalHeader

+                                         );

+  for (Index = 0; Index < Hdr.Pe32->FileHeader.NumberOfSections; Index++) {

+    Name = Section[Index].Name;

+    DEBUG ((

+      DEBUG_VERBOSE,

+      "  Section - '%c%c%c%c%c%c%c%c'\n",

+      Name[0],

+      Name[1],

+      Name[2],

+      Name[3],

+      Name[4],

+      Name[5],

+      Name[6],

+      Name[7]

+      ));

+

+    if ((Section[Index].Characteristics & EFI_IMAGE_SCN_CNT_CODE) != 0) {

+      DEBUG ((DEBUG_VERBOSE, "  VirtualSize          - 0x%08x\n", Section[Index].Misc.VirtualSize));

+      DEBUG ((DEBUG_VERBOSE, "  VirtualAddress       - 0x%08x\n", Section[Index].VirtualAddress));

+      DEBUG ((DEBUG_VERBOSE, "  SizeOfRawData        - 0x%08x\n", Section[Index].SizeOfRawData));

+      DEBUG ((DEBUG_VERBOSE, "  PointerToRawData     - 0x%08x\n", Section[Index].PointerToRawData));

+      DEBUG ((DEBUG_VERBOSE, "  PointerToRelocations - 0x%08x\n", Section[Index].PointerToRelocations));

+      DEBUG ((DEBUG_VERBOSE, "  PointerToLinenumbers - 0x%08x\n", Section[Index].PointerToLinenumbers));

+      DEBUG ((DEBUG_VERBOSE, "  NumberOfRelocations  - 0x%08x\n", Section[Index].NumberOfRelocations));

+      DEBUG ((DEBUG_VERBOSE, "  NumberOfLinenumbers  - 0x%08x\n", Section[Index].NumberOfLinenumbers));

+      DEBUG ((DEBUG_VERBOSE, "  Characteristics      - 0x%08x\n", Section[Index].Characteristics));

+

+      // Record code section(s)

+      ImageRecordCodeSection = AllocatePool (sizeof (*ImageRecordCodeSection));

+      if (ImageRecordCodeSection == NULL) {

+        return EFI_OUT_OF_RESOURCES;

+      }

+

+      ImageRecordCodeSection->Signature = IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE;

+

+      ImageRecordCodeSection->CodeSegmentBase = (UINTN)ImageBase + Section[Index].VirtualAddress;

+      ImageRecordCodeSection->CodeSegmentSize = Section[Index].SizeOfRawData;

+

+      InsertTailList (&ImageRecord->CodeSegmentList, &ImageRecordCodeSection->Link);

+      ImageRecord->CodeSegmentCount++;

+    }

+  }

+

+  if (ImageRecord->CodeSegmentCount > 0) {

+    SortImageRecordCodeSection (ImageRecord);

+  }

+

+  return Status;

+}

+

+/**

+  Deleted an image properties record. The function will also call

+  RemoveEntryList() on each code segment and the input ImageRecord before

+  freeing each pool.

+

+  @param[in]      ImageRecord             The IMAGE_PROPERTIES_RECORD to delete

+**/

+VOID

+EFIAPI

+DeleteImagePropertiesRecord (

+  IN  IMAGE_PROPERTIES_RECORD  *ImageRecord

+  )

+{

+  LIST_ENTRY                            *CodeSegmentListHead;

+  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;

+

+  if (ImageRecord == NULL) {

+    return;

+  }

+

+  CodeSegmentListHead = &ImageRecord->CodeSegmentList;

+  while (!IsListEmpty (CodeSegmentListHead)) {

+    ImageRecordCodeSection = CR (

+                               CodeSegmentListHead->ForwardLink,

+                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,

+                               Link,

+                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE

+                               );

+    RemoveEntryList (&ImageRecordCodeSection->Link);

+    FreePool (ImageRecordCodeSection);

+  }

+

+  if (!IsListEmpty (&ImageRecord->Link)) {

+    RemoveEntryList (&ImageRecord->Link);

+  }

+

+  FreePool (ImageRecord);

+}

diff --git a/MdeModulePkg/Library/ImagePropertiesRecordLib/ImagePropertiesRecordLib.inf b/MdeModulePkg/Library/ImagePropertiesRecordLib/ImagePropertiesRecordLib.inf
index 4c1466fc33..cfe0c04b3b 100644
--- a/MdeModulePkg/Library/ImagePropertiesRecordLib/ImagePropertiesRecordLib.inf
+++ b/MdeModulePkg/Library/ImagePropertiesRecordLib/ImagePropertiesRecordLib.inf
@@ -23,6 +23,7 @@
   BaseLib

   BaseMemoryLib

   DebugLib

+  MemoryAllocationLib

 

 [Packages]

   MdePkg/MdePkg.dec