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>

MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c

@@ -557,25 +557,6 @@ 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;
-  }
-}
-
 /**
   Insert image record.
 
@@ -586,20 +567,12 @@ InsertImageRecord (
   IN EFI_RUNTIME_IMAGE_ENTRY  *RuntimeImage
   )
 {
-  VOID                                  *ImageAddress;
+  EFI_STATUS               Status;
-  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;
+  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;
+  InitializeListHead (&ImageRecord->Link);
+  InitializeListHead (&ImageRecord->CodeSegmentList);
 
-  DEBUG ((DEBUG_VERBOSE, "ImageRecordCount - 0x%x\n", mImagePropertiesPrivateData.ImageRecordCount));
+  PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)RuntimeImage->ImageBase);
-
-  //
-  // Step 1: record whole region
-  //
-  ImageRecord->ImageBase = (EFI_PHYSICAL_ADDRESS)(UINTN)RuntimeImage->ImageBase;
-  ImageRecord->ImageSize = RuntimeImage->ImageSize;
-
-  ImageAddress = RuntimeImage->ImageBase;
-
-  PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)ImageAddress);
   if (PdbPointer != NULL) {
     DEBUG ((DEBUG_VERBOSE, "  Image - %a\n", PdbPointer));
   }
 
-  //
+  RequiredAlignment = RUNTIME_PAGE_ALLOCATION_GRANULARITY;
-  // Check PE/COFF image
+  Status            = CreateImagePropertiesRecord (
-  //
+                        RuntimeImage->ImageBase,
-  DosHdr             = (EFI_IMAGE_DOS_HEADER *)(UINTN)ImageAddress;
+                        RuntimeImage->ImageSize,
-  PeCoffHeaderOffset = 0;
+                        &RequiredAlignment,
-  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
+                        ImageRecord
-    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));
-    }
-
-    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) {
+  if (EFI_ERROR (Status)) {
-      DEBUG ((DEBUG_VERBOSE, "  VirtualSize          - 0x%08x\n", Section[Index].Misc.VirtualSize));
+    if (Status == EFI_ABORTED) {
-      DEBUG ((DEBUG_VERBOSE, "  VirtualAddress       - 0x%08x\n", Section[Index].VirtualAddress));
+      mMemoryAttributesTableEnable = FALSE;
-      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;
+    Status = EFI_ABORTED;
-
+    goto Finish;
-      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;
 }
 
@@ -771,8 +657,6 @@ RemoveImageRecord (
   )
 {
   IMAGE_PROPERTIES_RECORD  *ImageRecord;
-  LIST_ENTRY                            *CodeSegmentListHead;
-  IMAGE_PROPERTIES_RECORD_CODE_SECTION  *ImageRecordCodeSection;
 
   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;
+  DeleteImagePropertiesRecord (ImageRecord);
-  while (!IsListEmpty (CodeSegmentListHead)) {
-    ImageRecordCodeSection = CR (
-                               CodeSegmentListHead->ForwardLink,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,
-                               Link,
-                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
-                               );
-    RemoveEntryList (&ImageRecordCodeSection->Link);
-    FreePool (ImageRecordCodeSection);
-  }
 
-  RemoveEntryList (&ImageRecord->Link);
-  FreePool (ImageRecord);
   mImagePropertiesPrivateData.ImageRecordCount--;
 }

MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c

@@ -251,25 +251,6 @@ 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);
-  }
-}
-
 /**
   Insert image record.
 
@@ -280,158 +261,61 @@ SmmInsertImageRecord (
   IN EFI_SMM_DRIVER_ENTRY  *DriverEntry
   )
 {
-  VOID                                  *ImageAddress;
+  EFI_STATUS               Status;
-  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;
+  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;
+  InitializeListHead (&ImageRecord->Link);
+  InitializeListHead (&ImageRecord->CodeSegmentList);
 
-  DEBUG ((DEBUG_VERBOSE, "SMM ImageRecordCount - 0x%x\n", mImagePropertiesPrivateData.ImageRecordCount));
+  PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)DriverEntry->ImageBuffer);
-
-  //
-  // Step 1: record whole region
-  //
-  ImageRecord->ImageBase = DriverEntry->ImageBuffer;
-  ImageRecord->ImageSize = LShiftU64 (DriverEntry->NumberOfPage, EFI_PAGE_SHIFT);
-
-  ImageAddress = (VOID *)(UINTN)DriverEntry->ImageBuffer;
-
-  PdbPointer = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)ImageAddress);
   if (PdbPointer != NULL) {
     DEBUG ((DEBUG_VERBOSE, "SMM   Image - %a\n", PdbPointer));
   }
 
-  //
+  RequiredAlignment = RUNTIME_PAGE_ALLOCATION_GRANULARITY;
-  // Check PE/COFF image
+  Status            = CreateImagePropertiesRecord (
-  //
+                        (VOID *)(UINTN)DriverEntry->ImageBuffer,
-  DosHdr             = (EFI_IMAGE_DOS_HEADER *)(UINTN)ImageAddress;
+                        LShiftU64 (DriverEntry->NumberOfPage, EFI_PAGE_SHIFT),
-  PeCoffHeaderOffset = 0;
+                        &RequiredAlignment,
-  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
+                        ImageRecord
-    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;
-  }
-
-  //
-  // 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));
-    }
-
-    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) {
+  if (EFI_ERROR (Status)) {
-      DEBUG ((DEBUG_VERBOSE, "SMM   VirtualSize          - 0x%08x\n", Section[Index].Misc.VirtualSize));
+    if (Status == EFI_ABORTED) {
-      DEBUG ((DEBUG_VERBOSE, "SMM   VirtualAddress       - 0x%08x\n", Section[Index].VirtualAddress));
+      mMemoryProtectionAttribute &=
-      DEBUG ((DEBUG_VERBOSE, "SMM   SizeOfRawData        - 0x%08x\n", Section[Index].SizeOfRawData));
+        ~((UINT64)EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA);
-      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;
+    goto Finish;
-
-      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;
 }
 

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

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

MdeModulePkg/Library/ImagePropertiesRecordLib/ImagePropertiesRecordLib.inf

@@ -23,6 +23,7 @@
   BaseLib
   BaseMemoryLib
   DebugLib
+  MemoryAllocationLib
 
 [Packages]
   MdePkg/MdePkg.dec