[mirror_edk2.git] / MdeModulePkg / Core / PiSmmCore / MemoryAttributesTable.c

/** @file\r
  PI SMM MemoryAttributes support\r
\r
Copyright (c) 2016, Intel Corporation. All rights reserved.<BR>\r
SPDX-License-Identifier: BSD-2-Clause-Patent\r
\r
**/\r
\r
#include <PiDxe.h>\r
#include <Library/BaseLib.h>\r
#include <Library/BaseMemoryLib.h>\r
#include <Library/MemoryAllocationLib.h>\r
#include <Library/UefiBootServicesTableLib.h>\r
#include <Library/SmmServicesTableLib.h>\r
#include <Library/DebugLib.h>\r
#include <Library/PcdLib.h>\r
\r
#include <Library/PeCoffLib.h>\r
#include <Library/PeCoffGetEntryPointLib.h>\r
\r
#include <Guid/PiSmmMemoryAttributesTable.h>\r
\r
#include "PiSmmCore.h"\r
\r
#define PREVIOUS_MEMORY_DESCRIPTOR(MemoryDescriptor, Size) \\r
  ((EFI_MEMORY_DESCRIPTOR *)((UINT8 *)(MemoryDescriptor) - (Size)))\r
\r
#define IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE SIGNATURE_32 ('I','P','R','C')\r
\r
typedef struct {\r
  UINT32                 Signature;\r
  LIST_ENTRY             Link;\r
  EFI_PHYSICAL_ADDRESS   CodeSegmentBase;\r
  UINT64                 CodeSegmentSize;\r
} IMAGE_PROPERTIES_RECORD_CODE_SECTION;\r
\r
#define IMAGE_PROPERTIES_RECORD_SIGNATURE SIGNATURE_32 ('I','P','R','D')\r
\r
typedef struct {\r
  UINT32                 Signature;\r
  LIST_ENTRY             Link;\r
  EFI_PHYSICAL_ADDRESS   ImageBase;\r
  UINT64                 ImageSize;\r
  UINTN                  CodeSegmentCount;\r
  LIST_ENTRY             CodeSegmentList;\r
} IMAGE_PROPERTIES_RECORD;\r
\r
#define IMAGE_PROPERTIES_PRIVATE_DATA_SIGNATURE SIGNATURE_32 ('I','P','P','D')\r
\r
typedef struct {\r
  UINT32                 Signature;\r
  UINTN                  ImageRecordCount;\r
  UINTN                  CodeSegmentCountMax;\r
  LIST_ENTRY             ImageRecordList;\r
} IMAGE_PROPERTIES_PRIVATE_DATA;\r
\r
IMAGE_PROPERTIES_PRIVATE_DATA  mImagePropertiesPrivateData = {\r
  IMAGE_PROPERTIES_PRIVATE_DATA_SIGNATURE,\r
  0,\r
  0,\r
  INITIALIZE_LIST_HEAD_VARIABLE (mImagePropertiesPrivateData.ImageRecordList)\r
};\r
\r
#define EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA  BIT0\r
\r
UINT64 mMemoryProtectionAttribute = EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA;\r
\r
//\r
// Below functions are for MemoryMap\r
//\r
\r
/**\r
  Converts a number of EFI_PAGEs to a size in bytes.\r
\r
  NOTE: Do not use EFI_PAGES_TO_SIZE because it handles UINTN only.\r
\r
  @param[in]  Pages     The number of EFI_PAGES.\r
\r
  @return  The number of bytes associated with the number of EFI_PAGEs specified\r
           by Pages.\r
**/\r
STATIC\r
UINT64\r
EfiPagesToSize (\r
  IN UINT64 Pages\r
  )\r
{\r
  return LShiftU64 (Pages, EFI_PAGE_SHIFT);\r
}\r
\r
/**\r
  Converts a size, in bytes, to a number of EFI_PAGESs.\r
\r
  NOTE: Do not use EFI_SIZE_TO_PAGES because it handles UINTN only.\r
\r
  @param[in]  Size      A size in bytes.\r
\r
  @return  The number of EFI_PAGESs associated with the number of bytes specified\r
           by Size.\r
\r
**/\r
STATIC\r
UINT64\r
EfiSizeToPages (\r
  IN UINT64 Size\r
  )\r
{\r
  return RShiftU64 (Size, EFI_PAGE_SHIFT) + ((((UINTN)Size) & EFI_PAGE_MASK) ? 1 : 0);\r
}\r
\r
\r
/**\r
  Sort memory map entries based upon PhysicalStart, from low to high.\r
\r
  @param[in,out]  MemoryMap         A pointer to the buffer in which firmware places\r
                                    the current memory map.\r
  @param[in]      MemoryMapSize     Size, in bytes, of the MemoryMap buffer.\r
  @param[in]      DescriptorSize    Size, in bytes, of an individual EFI_MEMORY_DESCRIPTOR.\r
**/\r
STATIC\r
VOID\r
SortMemoryMap (\r
  IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,\r
  IN UINTN                      MemoryMapSize,\r
  IN UINTN                      DescriptorSize\r
  )\r
{\r
  EFI_MEMORY_DESCRIPTOR       *MemoryMapEntry;\r
  EFI_MEMORY_DESCRIPTOR       *NextMemoryMapEntry;\r
  EFI_MEMORY_DESCRIPTOR       *MemoryMapEnd;\r
  EFI_MEMORY_DESCRIPTOR       TempMemoryMap;\r
\r
  MemoryMapEntry = MemoryMap;\r
  NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
  MemoryMapEnd = (EFI_MEMORY_DESCRIPTOR *) ((UINT8 *) MemoryMap + MemoryMapSize);\r
  while (MemoryMapEntry < MemoryMapEnd) {\r
    while (NextMemoryMapEntry < MemoryMapEnd) {\r
      if (MemoryMapEntry->PhysicalStart > NextMemoryMapEntry->PhysicalStart) {\r
        CopyMem (&TempMemoryMap, MemoryMapEntry, sizeof(EFI_MEMORY_DESCRIPTOR));\r
        CopyMem (MemoryMapEntry, NextMemoryMapEntry, sizeof(EFI_MEMORY_DESCRIPTOR));\r
        CopyMem (NextMemoryMapEntry, &TempMemoryMap, sizeof(EFI_MEMORY_DESCRIPTOR));\r
      }\r
\r
      NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (NextMemoryMapEntry, DescriptorSize);\r
    }\r
\r
    MemoryMapEntry      = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
    NextMemoryMapEntry  = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
  }\r
\r
  return ;\r
}\r
\r
/**\r
  Merge continuous memory map entries whose have same attributes.\r
\r
  @param[in, out]  MemoryMap              A pointer to the buffer in which firmware places\r
                                          the current memory map.\r
  @param[in, out]  MemoryMapSize          A pointer to the size, in bytes, of the\r
                                          MemoryMap buffer. On input, this is the size of\r
                                          the current memory map.  On output,\r
                                          it is the size of new memory map after merge.\r
  @param[in]       DescriptorSize         Size, in bytes, of an individual EFI_MEMORY_DESCRIPTOR.\r
**/\r
STATIC\r
VOID\r
MergeMemoryMap (\r
  IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,\r
  IN OUT UINTN                  *MemoryMapSize,\r
  IN UINTN                      DescriptorSize\r
  )\r
{\r
  EFI_MEMORY_DESCRIPTOR       *MemoryMapEntry;\r
  EFI_MEMORY_DESCRIPTOR       *MemoryMapEnd;\r
  UINT64                      MemoryBlockLength;\r
  EFI_MEMORY_DESCRIPTOR       *NewMemoryMapEntry;\r
  EFI_MEMORY_DESCRIPTOR       *NextMemoryMapEntry;\r
\r
  MemoryMapEntry = MemoryMap;\r
  NewMemoryMapEntry = MemoryMap;\r
  MemoryMapEnd = (EFI_MEMORY_DESCRIPTOR *) ((UINT8 *) MemoryMap + *MemoryMapSize);\r
  while ((UINTN)MemoryMapEntry < (UINTN)MemoryMapEnd) {\r
    CopyMem (NewMemoryMapEntry, MemoryMapEntry, sizeof(EFI_MEMORY_DESCRIPTOR));\r
    NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
\r
    do {\r
      MemoryBlockLength = (UINT64) (EfiPagesToSize (MemoryMapEntry->NumberOfPages));\r
      if (((UINTN)NextMemoryMapEntry < (UINTN)MemoryMapEnd) &&\r
          (MemoryMapEntry->Type == NextMemoryMapEntry->Type) &&\r
          (MemoryMapEntry->Attribute == NextMemoryMapEntry->Attribute) &&\r
          ((MemoryMapEntry->PhysicalStart + MemoryBlockLength) == NextMemoryMapEntry->PhysicalStart)) {\r
        MemoryMapEntry->NumberOfPages += NextMemoryMapEntry->NumberOfPages;\r
        if (NewMemoryMapEntry != MemoryMapEntry) {\r
          NewMemoryMapEntry->NumberOfPages += NextMemoryMapEntry->NumberOfPages;\r
        }\r
\r
        NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (NextMemoryMapEntry, DescriptorSize);\r
        continue;\r
      } else {\r
        MemoryMapEntry = PREVIOUS_MEMORY_DESCRIPTOR (NextMemoryMapEntry, DescriptorSize);\r
        break;\r
      }\r
    } while (TRUE);\r
\r
    MemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
    NewMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (NewMemoryMapEntry, DescriptorSize);\r
  }\r
\r
  *MemoryMapSize = (UINTN)NewMemoryMapEntry - (UINTN)MemoryMap;\r
\r
  return ;\r
}\r
\r
/**\r
  Enforce memory map attributes.\r
  This function will set EfiRuntimeServicesData/EfiMemoryMappedIO/EfiMemoryMappedIOPortSpace to be EFI_MEMORY_XP.\r
\r
  @param[in, out]  MemoryMap              A pointer to the buffer in which firmware places\r
                                          the current memory map.\r
  @param[in]       MemoryMapSize          Size, in bytes, of the MemoryMap buffer.\r
  @param[in]       DescriptorSize         Size, in bytes, of an individual EFI_MEMORY_DESCRIPTOR.\r
**/\r
STATIC\r
VOID\r
EnforceMemoryMapAttribute (\r
  IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,\r
  IN UINTN                      MemoryMapSize,\r
  IN UINTN                      DescriptorSize\r
  )\r
{\r
  EFI_MEMORY_DESCRIPTOR       *MemoryMapEntry;\r
  EFI_MEMORY_DESCRIPTOR       *MemoryMapEnd;\r
\r
  MemoryMapEntry = MemoryMap;\r
  MemoryMapEnd   = (EFI_MEMORY_DESCRIPTOR *) ((UINT8 *) MemoryMap + MemoryMapSize);\r
  while ((UINTN)MemoryMapEntry < (UINTN)MemoryMapEnd) {\r
    if (MemoryMapEntry->Attribute != 0) {\r
      // It is PE image, the attribute is already set.\r
    } else {\r
      switch (MemoryMapEntry->Type) {\r
      case EfiRuntimeServicesCode:\r
        MemoryMapEntry->Attribute = EFI_MEMORY_RO;\r
        break;\r
      case EfiRuntimeServicesData:\r
      default:\r
        MemoryMapEntry->Attribute |= EFI_MEMORY_XP;\r
        break;\r
      }\r
    }\r
    MemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
  }\r
\r
  return ;\r
}\r
\r
/**\r
  Return the first image record, whose [ImageBase, ImageSize] covered by [Buffer, Length].\r
\r
  @param[in] Buffer  Start Address\r
  @param[in] Length  Address length\r
\r
  @return first image record covered by [buffer, length]\r
**/\r
STATIC\r
IMAGE_PROPERTIES_RECORD *\r
GetImageRecordByAddress (\r
  IN EFI_PHYSICAL_ADDRESS  Buffer,\r
  IN UINT64                Length\r
  )\r
{\r
  IMAGE_PROPERTIES_RECORD    *ImageRecord;\r
  LIST_ENTRY                 *ImageRecordLink;\r
  LIST_ENTRY                 *ImageRecordList;\r
\r
  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;\r
\r
  for (ImageRecordLink = ImageRecordList->ForwardLink;\r
       ImageRecordLink != ImageRecordList;\r
       ImageRecordLink = ImageRecordLink->ForwardLink) {\r
    ImageRecord = CR (\r
                    ImageRecordLink,\r
                    IMAGE_PROPERTIES_RECORD,\r
                    Link,\r
                    IMAGE_PROPERTIES_RECORD_SIGNATURE\r
                    );\r
\r
    if ((Buffer <= ImageRecord->ImageBase) &&\r
        (Buffer + Length >= ImageRecord->ImageBase + ImageRecord->ImageSize)) {\r
      return ImageRecord;\r
    }\r
  }\r
\r
  return NULL;\r
}\r
\r
/**\r
  Set the memory map to new entries, according to one old entry,\r
  based upon PE code section and data section in image record\r
\r
  @param[in]       ImageRecord            An image record whose [ImageBase, ImageSize] covered\r
                                          by old memory map entry.\r
  @param[in, out]  NewRecord              A pointer to several new memory map entries.\r
                                          The caller guarantee the buffer size be 1 +\r
                                          (SplitRecordCount * DescriptorSize) calculated\r
                                          below.\r
  @param[in]       OldRecord              A pointer to one old memory map entry.\r
  @param[in]       DescriptorSize         Size, in bytes, of an individual EFI_MEMORY_DESCRIPTOR.\r
**/\r
STATIC\r
UINTN\r
SetNewRecord (\r
  IN IMAGE_PROPERTIES_RECORD       *ImageRecord,\r
  IN OUT EFI_MEMORY_DESCRIPTOR     *NewRecord,\r
  IN EFI_MEMORY_DESCRIPTOR         *OldRecord,\r
  IN UINTN                         DescriptorSize\r
  )\r
{\r
  EFI_MEMORY_DESCRIPTOR                     TempRecord;\r
  IMAGE_PROPERTIES_RECORD_CODE_SECTION      *ImageRecordCodeSection;\r
  LIST_ENTRY                                *ImageRecordCodeSectionLink;\r
  LIST_ENTRY                                *ImageRecordCodeSectionEndLink;\r
  LIST_ENTRY                                *ImageRecordCodeSectionList;\r
  UINTN                                     NewRecordCount;\r
  UINT64                                    PhysicalEnd;\r
  UINT64                                    ImageEnd;\r
\r
  CopyMem (&TempRecord, OldRecord, sizeof(EFI_MEMORY_DESCRIPTOR));\r
  PhysicalEnd = TempRecord.PhysicalStart + EfiPagesToSize(TempRecord.NumberOfPages);\r
  NewRecordCount = 0;\r
\r
  //\r
  // Always create a new entry for non-PE image record\r
  //\r
  if (ImageRecord->ImageBase > TempRecord.PhysicalStart) {\r
    NewRecord->Type = TempRecord.Type;\r
    NewRecord->PhysicalStart = TempRecord.PhysicalStart;\r
    NewRecord->VirtualStart  = 0;\r
    NewRecord->NumberOfPages = EfiSizeToPages(ImageRecord->ImageBase - TempRecord.PhysicalStart);\r
    NewRecord->Attribute     = TempRecord.Attribute;\r
    NewRecord = NEXT_MEMORY_DESCRIPTOR (NewRecord, DescriptorSize);\r
    NewRecordCount ++;\r
    TempRecord.PhysicalStart = ImageRecord->ImageBase;\r
    TempRecord.NumberOfPages = EfiSizeToPages(PhysicalEnd - TempRecord.PhysicalStart);\r
  }\r
\r
  ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;\r
\r
  ImageRecordCodeSectionLink = ImageRecordCodeSectionList->ForwardLink;\r
  ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;\r
  while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {\r
    ImageRecordCodeSection = CR (\r
                               ImageRecordCodeSectionLink,\r
                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,\r
                               Link,\r
                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE\r
                               );\r
    ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;\r
\r
    if (TempRecord.PhysicalStart <= ImageRecordCodeSection->CodeSegmentBase) {\r
      //\r
      // DATA\r
      //\r
      NewRecord->Type = EfiRuntimeServicesData;\r
      NewRecord->PhysicalStart = TempRecord.PhysicalStart;\r
      NewRecord->VirtualStart  = 0;\r
      NewRecord->NumberOfPages = EfiSizeToPages(ImageRecordCodeSection->CodeSegmentBase - NewRecord->PhysicalStart);\r
      NewRecord->Attribute     = TempRecord.Attribute | EFI_MEMORY_XP;\r
      if (NewRecord->NumberOfPages != 0) {\r
        NewRecord = NEXT_MEMORY_DESCRIPTOR (NewRecord, DescriptorSize);\r
        NewRecordCount ++;\r
      }\r
\r
      //\r
      // CODE\r
      //\r
      NewRecord->Type = EfiRuntimeServicesCode;\r
      NewRecord->PhysicalStart = ImageRecordCodeSection->CodeSegmentBase;\r
      NewRecord->VirtualStart  = 0;\r
      NewRecord->NumberOfPages = EfiSizeToPages(ImageRecordCodeSection->CodeSegmentSize);\r
      NewRecord->Attribute     = (TempRecord.Attribute & (~EFI_MEMORY_XP)) | EFI_MEMORY_RO;\r
      if (NewRecord->NumberOfPages != 0) {\r
        NewRecord = NEXT_MEMORY_DESCRIPTOR (NewRecord, DescriptorSize);\r
        NewRecordCount ++;\r
      }\r
\r
      TempRecord.PhysicalStart = ImageRecordCodeSection->CodeSegmentBase + EfiPagesToSize (EfiSizeToPages(ImageRecordCodeSection->CodeSegmentSize));\r
      TempRecord.NumberOfPages = EfiSizeToPages(PhysicalEnd - TempRecord.PhysicalStart);\r
      if (TempRecord.NumberOfPages == 0) {\r
        break;\r
      }\r
    }\r
  }\r
\r
  ImageEnd = ImageRecord->ImageBase + ImageRecord->ImageSize;\r
\r
  //\r
  // Final DATA\r
  //\r
  if (TempRecord.PhysicalStart < ImageEnd) {\r
    NewRecord->Type = EfiRuntimeServicesData;\r
    NewRecord->PhysicalStart = TempRecord.PhysicalStart;\r
    NewRecord->VirtualStart  = 0;\r
    NewRecord->NumberOfPages = EfiSizeToPages (ImageEnd - TempRecord.PhysicalStart);\r
    NewRecord->Attribute     = TempRecord.Attribute | EFI_MEMORY_XP;\r
    NewRecordCount ++;\r
  }\r
\r
  return NewRecordCount;\r
}\r
\r
/**\r
  Return the max number of new splitted entries, according to one old entry,\r
  based upon PE code section and data section.\r
\r
  @param[in]  OldRecord              A pointer to one old memory map entry.\r
\r
  @retval  0 no entry need to be splitted.\r
  @return  the max number of new splitted entries\r
**/\r
STATIC\r
UINTN\r
GetMaxSplitRecordCount (\r
  IN EFI_MEMORY_DESCRIPTOR *OldRecord\r
  )\r
{\r
  IMAGE_PROPERTIES_RECORD *ImageRecord;\r
  UINTN                   SplitRecordCount;\r
  UINT64                  PhysicalStart;\r
  UINT64                  PhysicalEnd;\r
\r
  SplitRecordCount = 0;\r
  PhysicalStart = OldRecord->PhysicalStart;\r
  PhysicalEnd = OldRecord->PhysicalStart + EfiPagesToSize(OldRecord->NumberOfPages);\r
\r
  do {\r
    ImageRecord = GetImageRecordByAddress (PhysicalStart, PhysicalEnd - PhysicalStart);\r
    if (ImageRecord == NULL) {\r
      break;\r
    }\r
    SplitRecordCount += (2 * ImageRecord->CodeSegmentCount + 2);\r
    PhysicalStart = ImageRecord->ImageBase + ImageRecord->ImageSize;\r
  } while ((ImageRecord != NULL) && (PhysicalStart < PhysicalEnd));\r
\r
  return SplitRecordCount;\r
}\r
\r
/**\r
  Split the memory map to new entries, according to one old entry,\r
  based upon PE code section and data section.\r
\r
  @param[in]       OldRecord              A pointer to one old memory map entry.\r
  @param[in, out]  NewRecord              A pointer to several new memory map entries.\r
                                          The caller guarantee the buffer size be 1 +\r
                                          (SplitRecordCount * DescriptorSize) calculated\r
                                          below.\r
  @param[in]       MaxSplitRecordCount    The max number of splitted entries\r
  @param[in]       DescriptorSize         Size, in bytes, of an individual EFI_MEMORY_DESCRIPTOR.\r
\r
  @retval  0 no entry is splitted.\r
  @return  the real number of splitted record.\r
**/\r
STATIC\r
UINTN\r
SplitRecord (\r
  IN EFI_MEMORY_DESCRIPTOR     *OldRecord,\r
  IN OUT EFI_MEMORY_DESCRIPTOR *NewRecord,\r
  IN UINTN                     MaxSplitRecordCount,\r
  IN UINTN                     DescriptorSize\r
  )\r
{\r
  EFI_MEMORY_DESCRIPTOR   TempRecord;\r
  IMAGE_PROPERTIES_RECORD *ImageRecord;\r
  IMAGE_PROPERTIES_RECORD *NewImageRecord;\r
  UINT64                  PhysicalStart;\r
  UINT64                  PhysicalEnd;\r
  UINTN                   NewRecordCount;\r
  UINTN                   TotalNewRecordCount;\r
\r
  if (MaxSplitRecordCount == 0) {\r
    CopyMem (NewRecord, OldRecord, DescriptorSize);\r
    return 0;\r
  }\r
\r
  TotalNewRecordCount = 0;\r
\r
  //\r
  // Override previous record\r
  //\r
  CopyMem (&TempRecord, OldRecord, sizeof(EFI_MEMORY_DESCRIPTOR));\r
  PhysicalStart = TempRecord.PhysicalStart;\r
  PhysicalEnd = TempRecord.PhysicalStart + EfiPagesToSize(TempRecord.NumberOfPages);\r
\r
  ImageRecord = NULL;\r
  do {\r
    NewImageRecord = GetImageRecordByAddress (PhysicalStart, PhysicalEnd - PhysicalStart);\r
    if (NewImageRecord == NULL) {\r
      //\r
      // No more image covered by this range, stop\r
      //\r
      if (PhysicalEnd > PhysicalStart) {\r
        //\r
        // Always create a new entry for non-PE image record\r
        //\r
        NewRecord->Type = TempRecord.Type;\r
        NewRecord->PhysicalStart = TempRecord.PhysicalStart;\r
        NewRecord->VirtualStart  = 0;\r
        NewRecord->NumberOfPages = TempRecord.NumberOfPages;\r
        NewRecord->Attribute     = TempRecord.Attribute;\r
        TotalNewRecordCount ++;\r
      }\r
      break;\r
    }\r
    ImageRecord = NewImageRecord;\r
\r
    //\r
    // Set new record\r
    //\r
    NewRecordCount = SetNewRecord (ImageRecord, NewRecord, &TempRecord, DescriptorSize);\r
    TotalNewRecordCount += NewRecordCount;\r
    NewRecord = (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)NewRecord + NewRecordCount * DescriptorSize);\r
\r
    //\r
    // Update PhysicalStart, in order to exclude the image buffer already splitted.\r
    //\r
    PhysicalStart = ImageRecord->ImageBase + ImageRecord->ImageSize;\r
    TempRecord.PhysicalStart = PhysicalStart;\r
    TempRecord.NumberOfPages = EfiSizeToPages (PhysicalEnd - PhysicalStart);\r
  } while ((ImageRecord != NULL) && (PhysicalStart < PhysicalEnd));\r
\r
  return TotalNewRecordCount - 1;\r
}\r
\r
/**\r
  Split the original memory map, and add more entries to describe PE code section and data section.\r
  This function will set EfiRuntimeServicesData to be EFI_MEMORY_XP.\r
  This function will merge entries with same attributes finally.\r
\r
  NOTE: It assumes PE code/data section are page aligned.\r
  NOTE: It assumes enough entry is prepared for new memory map.\r
\r
  Split table:\r
   +---------------+\r
   | Record X      |\r
   +---------------+\r
   | Record RtCode |\r
   +---------------+\r
   | Record Y      |\r
   +---------------+\r
   ==>\r
   +---------------+\r
   | Record X      |\r
   +---------------+\r
   | Record RtCode |\r
   +---------------+ ----\r
   | Record RtData |     |\r
   +---------------+     |\r
   | Record RtCode |     |-> PE/COFF1\r
   +---------------+     |\r
   | Record RtData |     |\r
   +---------------+ ----\r
   | Record RtCode |\r
   +---------------+ ----\r
   | Record RtData |     |\r
   +---------------+     |\r
   | Record RtCode |     |-> PE/COFF2\r
   +---------------+     |\r
   | Record RtData |     |\r
   +---------------+ ----\r
   | Record RtCode |\r
   +---------------+\r
   | Record Y      |\r
   +---------------+\r
\r
  @param[in, out]  MemoryMapSize          A pointer to the size, in bytes, of the\r
                                          MemoryMap buffer. On input, this is the size of\r
                                          old MemoryMap before split. The actual buffer\r
                                          size of MemoryMap is MemoryMapSize +\r
                                          (AdditionalRecordCount * DescriptorSize) calculated\r
                                          below. On output, it is the size of new MemoryMap\r
                                          after split.\r
  @param[in, out]  MemoryMap              A pointer to the buffer in which firmware places\r
                                          the current memory map.\r
  @param[in]       DescriptorSize         Size, in bytes, of an individual EFI_MEMORY_DESCRIPTOR.\r
**/\r
STATIC\r
VOID\r
SplitTable (\r
  IN OUT UINTN                  *MemoryMapSize,\r
  IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,\r
  IN UINTN                      DescriptorSize\r
  )\r
{\r
  INTN        IndexOld;\r
  INTN        IndexNew;\r
  UINTN       MaxSplitRecordCount;\r
  UINTN       RealSplitRecordCount;\r
  UINTN       TotalSplitRecordCount;\r
  UINTN       AdditionalRecordCount;\r
\r
  AdditionalRecordCount = (2 * mImagePropertiesPrivateData.CodeSegmentCountMax + 2) * mImagePropertiesPrivateData.ImageRecordCount;\r
\r
  TotalSplitRecordCount = 0;\r
  //\r
  // Let old record point to end of valid MemoryMap buffer.\r
  //\r
  IndexOld = ((*MemoryMapSize) / DescriptorSize) - 1;\r
  //\r
  // Let new record point to end of full MemoryMap buffer.\r
  //\r
  IndexNew = ((*MemoryMapSize) / DescriptorSize) - 1 + AdditionalRecordCount;\r
  for (; IndexOld >= 0; IndexOld--) {\r
    MaxSplitRecordCount = GetMaxSplitRecordCount ((EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + IndexOld * DescriptorSize));\r
    //\r
    // Split this MemoryMap record\r
    //\r
    IndexNew -= MaxSplitRecordCount;\r
    RealSplitRecordCount = SplitRecord (\r
                             (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + IndexOld * DescriptorSize),\r
                             (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + IndexNew * DescriptorSize),\r
                             MaxSplitRecordCount,\r
                             DescriptorSize\r
                             );\r
    //\r
    // Adjust IndexNew according to real split.\r
    //\r
    if (MaxSplitRecordCount != RealSplitRecordCount) {\r
      CopyMem (\r
        ((UINT8 *)MemoryMap + (IndexNew + MaxSplitRecordCount - RealSplitRecordCount) * DescriptorSize),\r
        ((UINT8 *)MemoryMap + IndexNew * DescriptorSize),\r
        (RealSplitRecordCount + 1) * DescriptorSize\r
        );\r
    }\r
    IndexNew = IndexNew + MaxSplitRecordCount - RealSplitRecordCount;\r
    TotalSplitRecordCount += RealSplitRecordCount;\r
    IndexNew --;\r
  }\r
  //\r
  // Move all records to the beginning.\r
  //\r
  CopyMem (\r
    MemoryMap,\r
    (UINT8 *)MemoryMap + (AdditionalRecordCount - TotalSplitRecordCount) * DescriptorSize,\r
    (*MemoryMapSize) + TotalSplitRecordCount * DescriptorSize\r
    );\r
\r
  *MemoryMapSize = (*MemoryMapSize) + DescriptorSize * TotalSplitRecordCount;\r
\r
  //\r
  // Sort from low to high (Just in case)\r
  //\r
  SortMemoryMap (MemoryMap, *MemoryMapSize, DescriptorSize);\r
\r
  //\r
  // Set RuntimeData to XP\r
  //\r
  EnforceMemoryMapAttribute (MemoryMap, *MemoryMapSize, DescriptorSize);\r
\r
  //\r
  // Merge same type to save entry size\r
  //\r
  MergeMemoryMap (MemoryMap, MemoryMapSize, DescriptorSize);\r
\r
  return ;\r
}\r
\r
/**\r
  This function for GetMemoryMap() with memory attributes table.\r
\r
  It calls original GetMemoryMap() to get the original memory map information. Then\r
  plus the additional memory map entries for PE Code/Data separation.\r
\r
  @param[in, out]  MemoryMapSize          A pointer to the size, in bytes, of the\r
                                          MemoryMap buffer. On input, this is the size of\r
                                          the buffer allocated by the caller.  On output,\r
                                          it is the size of the buffer returned by the\r
                                          firmware  if the buffer was large enough, or the\r
                                          size of the buffer needed  to contain the map if\r
                                          the buffer was too small.\r
  @param[in, out]  MemoryMap              A pointer to the buffer in which firmware places\r
                                          the current memory map.\r
  @param[out]      MapKey                 A pointer to the location in which firmware\r
                                          returns the key for the current memory map.\r
  @param[out]      DescriptorSize         A pointer to the location in which firmware\r
                                          returns the size, in bytes, of an individual\r
                                          EFI_MEMORY_DESCRIPTOR.\r
  @param[out]      DescriptorVersion      A pointer to the location in which firmware\r
                                          returns the version number associated with the\r
                                          EFI_MEMORY_DESCRIPTOR.\r
\r
  @retval EFI_SUCCESS            The memory map was returned in the MemoryMap\r
                                 buffer.\r
  @retval EFI_BUFFER_TOO_SMALL   The MemoryMap buffer was too small. The current\r
                                 buffer size needed to hold the memory map is\r
                                 returned in MemoryMapSize.\r
  @retval EFI_INVALID_PARAMETER  One of the parameters has an invalid value.\r
\r
**/\r
STATIC\r
EFI_STATUS\r
EFIAPI\r
SmmCoreGetMemoryMapMemoryAttributesTable (\r
  IN OUT UINTN                  *MemoryMapSize,\r
  IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,\r
  OUT UINTN                     *MapKey,\r
  OUT UINTN                     *DescriptorSize,\r
  OUT UINT32                    *DescriptorVersion\r
  )\r
{\r
  EFI_STATUS  Status;\r
  UINTN       OldMemoryMapSize;\r
  UINTN       AdditionalRecordCount;\r
\r
  //\r
  // If PE code/data is not aligned, just return.\r
  //\r
  if ((mMemoryProtectionAttribute & EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA) == 0) {\r
    return SmmCoreGetMemoryMap (MemoryMapSize, MemoryMap, MapKey, DescriptorSize, DescriptorVersion);\r
  }\r
\r
  if (MemoryMapSize == NULL) {\r
    return EFI_INVALID_PARAMETER;\r
  }\r
\r
  AdditionalRecordCount = (2 * mImagePropertiesPrivateData.CodeSegmentCountMax + 2) * mImagePropertiesPrivateData.ImageRecordCount;\r
\r
  OldMemoryMapSize = *MemoryMapSize;\r
  Status = SmmCoreGetMemoryMap (MemoryMapSize, MemoryMap, MapKey, DescriptorSize, DescriptorVersion);\r
  if (Status == EFI_BUFFER_TOO_SMALL) {\r
    *MemoryMapSize = *MemoryMapSize + (*DescriptorSize) * AdditionalRecordCount;\r
  } else if (Status == EFI_SUCCESS) {\r
    if (OldMemoryMapSize - *MemoryMapSize < (*DescriptorSize) * AdditionalRecordCount) {\r
      *MemoryMapSize = *MemoryMapSize + (*DescriptorSize) * AdditionalRecordCount;\r
      //\r
      // Need update status to buffer too small\r
      //\r
      Status = EFI_BUFFER_TOO_SMALL;\r
    } else {\r
      //\r
      // Split PE code/data\r
      //\r
      ASSERT(MemoryMap != NULL);\r
      SplitTable (MemoryMapSize, MemoryMap, *DescriptorSize);\r
    }\r
  }\r
\r
  return Status;\r
}\r
\r
//\r
// Below functions are for ImageRecord\r
//\r
\r
/**\r
  Set MemoryProtectionAttribute according to PE/COFF image section alignment.\r
\r
  @param[in]  SectionAlignment    PE/COFF section alignment\r
**/\r
STATIC\r
VOID\r
SetMemoryAttributesTableSectionAlignment (\r
  IN UINT32  SectionAlignment\r
  )\r
{\r
  if (((SectionAlignment & (RUNTIME_PAGE_ALLOCATION_GRANULARITY - 1)) != 0) &&\r
      ((mMemoryProtectionAttribute & EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA) != 0)) {\r
    DEBUG ((DEBUG_VERBOSE, "SMM SetMemoryAttributesTableSectionAlignment - Clear\n"));\r
    mMemoryProtectionAttribute &= ~((UINT64)EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA);\r
  }\r
}\r
\r
/**\r
  Swap two code sections in image record.\r
\r
  @param[in]  FirstImageRecordCodeSection    first code section in image record\r
  @param[in]  SecondImageRecordCodeSection   second code section in image record\r
**/\r
STATIC\r
VOID\r
SwapImageRecordCodeSection (\r
  IN IMAGE_PROPERTIES_RECORD_CODE_SECTION      *FirstImageRecordCodeSection,\r
  IN IMAGE_PROPERTIES_RECORD_CODE_SECTION      *SecondImageRecordCodeSection\r
  )\r
{\r
  IMAGE_PROPERTIES_RECORD_CODE_SECTION      TempImageRecordCodeSection;\r
\r
  TempImageRecordCodeSection.CodeSegmentBase = FirstImageRecordCodeSection->CodeSegmentBase;\r
  TempImageRecordCodeSection.CodeSegmentSize = FirstImageRecordCodeSection->CodeSegmentSize;\r
\r
  FirstImageRecordCodeSection->CodeSegmentBase = SecondImageRecordCodeSection->CodeSegmentBase;\r
  FirstImageRecordCodeSection->CodeSegmentSize = SecondImageRecordCodeSection->CodeSegmentSize;\r
\r
  SecondImageRecordCodeSection->CodeSegmentBase = TempImageRecordCodeSection.CodeSegmentBase;\r
  SecondImageRecordCodeSection->CodeSegmentSize = TempImageRecordCodeSection.CodeSegmentSize;\r
}\r
\r
/**\r
  Sort code section in image record, based upon CodeSegmentBase from low to high.\r
\r
  @param[in]  ImageRecord    image record to be sorted\r
**/\r
STATIC\r
VOID\r
SortImageRecordCodeSection (\r
  IN IMAGE_PROPERTIES_RECORD              *ImageRecord\r
  )\r
{\r
  IMAGE_PROPERTIES_RECORD_CODE_SECTION      *ImageRecordCodeSection;\r
  IMAGE_PROPERTIES_RECORD_CODE_SECTION      *NextImageRecordCodeSection;\r
  LIST_ENTRY                                *ImageRecordCodeSectionLink;\r
  LIST_ENTRY                                *NextImageRecordCodeSectionLink;\r
  LIST_ENTRY                                *ImageRecordCodeSectionEndLink;\r
  LIST_ENTRY                                *ImageRecordCodeSectionList;\r
\r
  ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;\r
\r
  ImageRecordCodeSectionLink = ImageRecordCodeSectionList->ForwardLink;\r
  NextImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;\r
  ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;\r
  while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {\r
    ImageRecordCodeSection = CR (\r
                               ImageRecordCodeSectionLink,\r
                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,\r
                               Link,\r
                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE\r
                               );\r
    while (NextImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {\r
      NextImageRecordCodeSection = CR (\r
                                     NextImageRecordCodeSectionLink,\r
                                     IMAGE_PROPERTIES_RECORD_CODE_SECTION,\r
                                     Link,\r
                                     IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE\r
                                     );\r
      if (ImageRecordCodeSection->CodeSegmentBase > NextImageRecordCodeSection->CodeSegmentBase) {\r
        SwapImageRecordCodeSection (ImageRecordCodeSection, NextImageRecordCodeSection);\r
      }\r
      NextImageRecordCodeSectionLink = NextImageRecordCodeSectionLink->ForwardLink;\r
    }\r
\r
    ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;\r
    NextImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;\r
  }\r
}\r
\r
/**\r
  Check if code section in image record is valid.\r
\r
  @param[in]  ImageRecord    image record to be checked\r
\r
  @retval TRUE  image record is valid\r
  @retval FALSE image record is invalid\r
**/\r
STATIC\r
BOOLEAN\r
IsImageRecordCodeSectionValid (\r
  IN IMAGE_PROPERTIES_RECORD              *ImageRecord\r
  )\r
{\r
  IMAGE_PROPERTIES_RECORD_CODE_SECTION      *ImageRecordCodeSection;\r
  IMAGE_PROPERTIES_RECORD_CODE_SECTION      *LastImageRecordCodeSection;\r
  LIST_ENTRY                                *ImageRecordCodeSectionLink;\r
  LIST_ENTRY                                *ImageRecordCodeSectionEndLink;\r
  LIST_ENTRY                                *ImageRecordCodeSectionList;\r
\r
  DEBUG ((DEBUG_VERBOSE, "SMM ImageCode SegmentCount - 0x%x\n", ImageRecord->CodeSegmentCount));\r
\r
  ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;\r
\r
  ImageRecordCodeSectionLink = ImageRecordCodeSectionList->ForwardLink;\r
  ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;\r
  LastImageRecordCodeSection = NULL;\r
  while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {\r
    ImageRecordCodeSection = CR (\r
                               ImageRecordCodeSectionLink,\r
                               IMAGE_PROPERTIES_RECORD_CODE_SECTION,\r
                               Link,\r
                               IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE\r
                               );\r
    if (ImageRecordCodeSection->CodeSegmentSize == 0) {\r
      return FALSE;\r
    }\r
    if (ImageRecordCodeSection->CodeSegmentBase < ImageRecord->ImageBase) {\r
      return FALSE;\r
    }\r
    if (ImageRecordCodeSection->CodeSegmentBase >= MAX_ADDRESS - ImageRecordCodeSection->CodeSegmentSize) {\r
      return FALSE;\r
    }\r
    if ((ImageRecordCodeSection->CodeSegmentBase + ImageRecordCodeSection->CodeSegmentSize) > (ImageRecord->ImageBase + ImageRecord->ImageSize)) {\r
      return FALSE;\r
    }\r
    if (LastImageRecordCodeSection != NULL) {\r
      if ((LastImageRecordCodeSection->CodeSegmentBase + LastImageRecordCodeSection->CodeSegmentSize) > ImageRecordCodeSection->CodeSegmentBase) {\r
        return FALSE;\r
      }\r
    }\r
\r
    LastImageRecordCodeSection = ImageRecordCodeSection;\r
    ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;\r
  }\r
\r
  return TRUE;\r
}\r
\r
/**\r
  Swap two image records.\r
\r
  @param[in]  FirstImageRecord   first image record.\r
  @param[in]  SecondImageRecord  second image record.\r
**/\r
STATIC\r
VOID\r
SwapImageRecord (\r
  IN IMAGE_PROPERTIES_RECORD      *FirstImageRecord,\r
  IN IMAGE_PROPERTIES_RECORD      *SecondImageRecord\r
  )\r
{\r
  IMAGE_PROPERTIES_RECORD      TempImageRecord;\r
\r
  TempImageRecord.ImageBase = FirstImageRecord->ImageBase;\r
  TempImageRecord.ImageSize = FirstImageRecord->ImageSize;\r
  TempImageRecord.CodeSegmentCount = FirstImageRecord->CodeSegmentCount;\r
\r
  FirstImageRecord->ImageBase = SecondImageRecord->ImageBase;\r
  FirstImageRecord->ImageSize = SecondImageRecord->ImageSize;\r
  FirstImageRecord->CodeSegmentCount = SecondImageRecord->CodeSegmentCount;\r
\r
  SecondImageRecord->ImageBase = TempImageRecord.ImageBase;\r
  SecondImageRecord->ImageSize = TempImageRecord.ImageSize;\r
  SecondImageRecord->CodeSegmentCount = TempImageRecord.CodeSegmentCount;\r
\r
  SwapListEntries (&FirstImageRecord->CodeSegmentList, &SecondImageRecord->CodeSegmentList);\r
}\r
\r
/**\r
  Sort image record based upon the ImageBase from low to high.\r
**/\r
STATIC\r
VOID\r
SortImageRecord (\r
  VOID\r
  )\r
{\r
  IMAGE_PROPERTIES_RECORD      *ImageRecord;\r
  IMAGE_PROPERTIES_RECORD      *NextImageRecord;\r
  LIST_ENTRY                   *ImageRecordLink;\r
  LIST_ENTRY                   *NextImageRecordLink;\r
  LIST_ENTRY                   *ImageRecordEndLink;\r
  LIST_ENTRY                   *ImageRecordList;\r
\r
  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;\r
\r
  ImageRecordLink = ImageRecordList->ForwardLink;\r
  NextImageRecordLink = ImageRecordLink->ForwardLink;\r
  ImageRecordEndLink = ImageRecordList;\r
  while (ImageRecordLink != ImageRecordEndLink) {\r
    ImageRecord = CR (\r
                    ImageRecordLink,\r
                    IMAGE_PROPERTIES_RECORD,\r
                    Link,\r
                    IMAGE_PROPERTIES_RECORD_SIGNATURE\r
                    );\r
    while (NextImageRecordLink != ImageRecordEndLink) {\r
      NextImageRecord = CR (\r
                          NextImageRecordLink,\r
                          IMAGE_PROPERTIES_RECORD,\r
                          Link,\r
                          IMAGE_PROPERTIES_RECORD_SIGNATURE\r
                          );\r
      if (ImageRecord->ImageBase > NextImageRecord->ImageBase) {\r
        SwapImageRecord (ImageRecord, NextImageRecord);\r
      }\r
      NextImageRecordLink = NextImageRecordLink->ForwardLink;\r
    }\r
\r
    ImageRecordLink = ImageRecordLink->ForwardLink;\r
    NextImageRecordLink = ImageRecordLink->ForwardLink;\r
  }\r
}\r
\r
/**\r
  Dump image record.\r
**/\r
STATIC\r
VOID\r
DumpImageRecord (\r
  VOID\r
  )\r
{\r
  IMAGE_PROPERTIES_RECORD      *ImageRecord;\r
  LIST_ENTRY                   *ImageRecordLink;\r
  LIST_ENTRY                   *ImageRecordList;\r
  UINTN                        Index;\r
\r
  ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;\r
\r
  for (ImageRecordLink = ImageRecordList->ForwardLink, Index= 0;\r
       ImageRecordLink != ImageRecordList;\r
       ImageRecordLink = ImageRecordLink->ForwardLink, Index++) {\r
    ImageRecord = CR (\r
                    ImageRecordLink,\r
                    IMAGE_PROPERTIES_RECORD,\r
                    Link,\r
                    IMAGE_PROPERTIES_RECORD_SIGNATURE\r
                    );\r
    DEBUG ((DEBUG_VERBOSE, "SMM  Image[%d]: 0x%016lx - 0x%016lx\n", Index, ImageRecord->ImageBase, ImageRecord->ImageSize));\r
  }\r
}\r
\r
/**\r
  Insert image record.\r
\r
  @param[in]  DriverEntry    Driver information\r
**/\r
VOID\r
SmmInsertImageRecord (\r
  IN EFI_SMM_DRIVER_ENTRY  *DriverEntry\r
  )\r
{\r
  VOID                                 *ImageAddress;\r
  EFI_IMAGE_DOS_HEADER                 *DosHdr;\r
  UINT32                               PeCoffHeaderOffset;\r
  UINT32                               SectionAlignment;\r
  EFI_IMAGE_SECTION_HEADER             *Section;\r
  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;\r
  UINT8                                *Name;\r
  UINTN                                Index;\r
  IMAGE_PROPERTIES_RECORD              *ImageRecord;\r
  CHAR8                                *PdbPointer;\r
  IMAGE_PROPERTIES_RECORD_CODE_SECTION *ImageRecordCodeSection;\r
\r
  DEBUG ((DEBUG_VERBOSE, "SMM InsertImageRecord - 0x%x\n", DriverEntry));\r
  DEBUG ((DEBUG_VERBOSE, "SMM InsertImageRecord - 0x%016lx - 0x%08x\n", DriverEntry->ImageBuffer, DriverEntry->NumberOfPage));\r
\r
  ImageRecord = AllocatePool (sizeof(*ImageRecord));\r
  if (ImageRecord == NULL) {\r
    return ;\r
  }\r
  ImageRecord->Signature = IMAGE_PROPERTIES_RECORD_SIGNATURE;\r
\r
  DEBUG ((DEBUG_VERBOSE, "SMM ImageRecordCount - 0x%x\n", mImagePropertiesPrivateData.ImageRecordCount));\r
\r
  //\r
  // Step 1: record whole region\r
  //\r
  ImageRecord->ImageBase = DriverEntry->ImageBuffer;\r
  ImageRecord->ImageSize = EfiPagesToSize(DriverEntry->NumberOfPage);\r
\r
  ImageAddress = (VOID *)(UINTN)DriverEntry->ImageBuffer;\r
\r
  PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);\r
  if (PdbPointer != NULL) {\r
    DEBUG ((DEBUG_VERBOSE, "SMM   Image - %a\n", PdbPointer));\r
  }\r
\r
  //\r
  // Check PE/COFF image\r
  //\r
  DosHdr = (EFI_IMAGE_DOS_HEADER *) (UINTN) ImageAddress;\r
  PeCoffHeaderOffset = 0;\r
  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {\r
    PeCoffHeaderOffset = DosHdr->e_lfanew;\r
  }\r
\r
  Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINT8 *) (UINTN) ImageAddress + PeCoffHeaderOffset);\r
  if (Hdr.Pe32->Signature != EFI_IMAGE_NT_SIGNATURE) {\r
    DEBUG ((DEBUG_VERBOSE, "SMM Hdr.Pe32->Signature invalid - 0x%x\n", Hdr.Pe32->Signature));\r
    goto Finish;\r
  }\r
\r
  //\r
  // Get SectionAlignment\r
  //\r
  if (Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
    SectionAlignment  = Hdr.Pe32->OptionalHeader.SectionAlignment;\r
  } else {\r
    SectionAlignment  = Hdr.Pe32Plus->OptionalHeader.SectionAlignment;\r
  }\r
\r
  SetMemoryAttributesTableSectionAlignment (SectionAlignment);\r
  if ((SectionAlignment & (RUNTIME_PAGE_ALLOCATION_GRANULARITY - 1)) != 0) {\r
    DEBUG ((DEBUG_WARN, "SMM !!!!!!!!  InsertImageRecord - Section Alignment(0x%x) is not %dK  !!!!!!!!\n",\r
      SectionAlignment, RUNTIME_PAGE_ALLOCATION_GRANULARITY >> 10));\r
    PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);\r
    if (PdbPointer != NULL) {\r
      DEBUG ((DEBUG_WARN, "SMM !!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));\r
    }\r
    goto Finish;\r
  }\r
\r
  Section = (EFI_IMAGE_SECTION_HEADER *) (\r
               (UINT8 *) (UINTN) ImageAddress +\r
               PeCoffHeaderOffset +\r
               sizeof(UINT32) +\r
               sizeof(EFI_IMAGE_FILE_HEADER) +\r
               Hdr.Pe32->FileHeader.SizeOfOptionalHeader\r
               );\r
  ImageRecord->CodeSegmentCount = 0;\r
  InitializeListHead (&ImageRecord->CodeSegmentList);\r
  for (Index = 0; Index < Hdr.Pe32->FileHeader.NumberOfSections; Index++) {\r
    Name = Section[Index].Name;\r
    DEBUG ((\r
      DEBUG_VERBOSE,\r
      "SMM   Section - '%c%c%c%c%c%c%c%c'\n",\r
      Name[0],\r
      Name[1],\r
      Name[2],\r
      Name[3],\r
      Name[4],\r
      Name[5],\r
      Name[6],\r
      Name[7]\r
      ));\r
\r
    if ((Section[Index].Characteristics & EFI_IMAGE_SCN_CNT_CODE) != 0) {\r
      DEBUG ((DEBUG_VERBOSE, "SMM   VirtualSize          - 0x%08x\n", Section[Index].Misc.VirtualSize));\r
      DEBUG ((DEBUG_VERBOSE, "SMM   VirtualAddress       - 0x%08x\n", Section[Index].VirtualAddress));\r
      DEBUG ((DEBUG_VERBOSE, "SMM   SizeOfRawData        - 0x%08x\n", Section[Index].SizeOfRawData));\r
      DEBUG ((DEBUG_VERBOSE, "SMM   PointerToRawData     - 0x%08x\n", Section[Index].PointerToRawData));\r
      DEBUG ((DEBUG_VERBOSE, "SMM   PointerToRelocations - 0x%08x\n", Section[Index].PointerToRelocations));\r
      DEBUG ((DEBUG_VERBOSE, "SMM   PointerToLinenumbers - 0x%08x\n", Section[Index].PointerToLinenumbers));\r
      DEBUG ((DEBUG_VERBOSE, "SMM   NumberOfRelocations  - 0x%08x\n", Section[Index].NumberOfRelocations));\r
      DEBUG ((DEBUG_VERBOSE, "SMM   NumberOfLinenumbers  - 0x%08x\n", Section[Index].NumberOfLinenumbers));\r
      DEBUG ((DEBUG_VERBOSE, "SMM   Characteristics      - 0x%08x\n", Section[Index].Characteristics));\r
\r
      //\r
      // Step 2: record code section\r
      //\r
      ImageRecordCodeSection = AllocatePool (sizeof(*ImageRecordCodeSection));\r
      if (ImageRecordCodeSection == NULL) {\r
        return ;\r
      }\r
      ImageRecordCodeSection->Signature = IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE;\r
\r
      ImageRecordCodeSection->CodeSegmentBase = (UINTN)ImageAddress + Section[Index].VirtualAddress;\r
      ImageRecordCodeSection->CodeSegmentSize = Section[Index].SizeOfRawData;\r
\r
      DEBUG ((DEBUG_VERBOSE, "SMM ImageCode: 0x%016lx - 0x%016lx\n", ImageRecordCodeSection->CodeSegmentBase, ImageRecordCodeSection->CodeSegmentSize));\r
\r
      InsertTailList (&ImageRecord->CodeSegmentList, &ImageRecordCodeSection->Link);\r
      ImageRecord->CodeSegmentCount++;\r
    }\r
  }\r
\r
  if (ImageRecord->CodeSegmentCount == 0) {\r
    SetMemoryAttributesTableSectionAlignment (1);\r
    DEBUG ((DEBUG_ERROR, "SMM !!!!!!!!  InsertImageRecord - CodeSegmentCount is 0  !!!!!!!!\n"));\r
    PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);\r
    if (PdbPointer != NULL) {\r
      DEBUG ((DEBUG_ERROR, "SMM !!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));\r
    }\r
    goto Finish;\r
  }\r
\r
  //\r
  // Final\r
  //\r
  SortImageRecordCodeSection (ImageRecord);\r
  //\r
  // Check overlap all section in ImageBase/Size\r
  //\r
  if (!IsImageRecordCodeSectionValid (ImageRecord)) {\r
    DEBUG ((DEBUG_ERROR, "SMM IsImageRecordCodeSectionValid - FAIL\n"));\r
    goto Finish;\r
  }\r
\r
  InsertTailList (&mImagePropertiesPrivateData.ImageRecordList, &ImageRecord->Link);\r
  mImagePropertiesPrivateData.ImageRecordCount++;\r
\r
  if (mImagePropertiesPrivateData.CodeSegmentCountMax < ImageRecord->CodeSegmentCount) {\r
    mImagePropertiesPrivateData.CodeSegmentCountMax = ImageRecord->CodeSegmentCount;\r
  }\r
\r
  SortImageRecord ();\r
\r
Finish:\r
  return ;\r
}\r
\r
\r
/**\r
  Publish MemoryAttributesTable to SMM configuration table.\r
**/\r
VOID\r
PublishMemoryAttributesTable (\r
  VOID\r
  )\r
{\r
  UINTN                                MemoryMapSize;\r
  EFI_MEMORY_DESCRIPTOR                *MemoryMap;\r
  UINTN                                MapKey;\r
  UINTN                                DescriptorSize;\r
  UINT32                               DescriptorVersion;\r
  UINTN                                Index;\r
  EFI_STATUS                           Status;\r
  UINTN                                RuntimeEntryCount;\r
  EDKII_PI_SMM_MEMORY_ATTRIBUTES_TABLE *MemoryAttributesTable;\r
  EFI_MEMORY_DESCRIPTOR                *MemoryAttributesEntry;\r
  UINTN                                MemoryAttributesTableSize;\r
\r
  MemoryMapSize = 0;\r
  MemoryMap = NULL;\r
  Status = SmmCoreGetMemoryMapMemoryAttributesTable (\r
             &MemoryMapSize,\r
             MemoryMap,\r
             &MapKey,\r
             &DescriptorSize,\r
             &DescriptorVersion\r
             );\r
  ASSERT (Status == EFI_BUFFER_TOO_SMALL);\r
\r
  do {\r
    DEBUG ((DEBUG_INFO, "MemoryMapSize - 0x%x\n", MemoryMapSize));\r
    MemoryMap = AllocatePool (MemoryMapSize);\r
    ASSERT (MemoryMap != NULL);\r
    DEBUG ((DEBUG_INFO, "MemoryMap - 0x%x\n", MemoryMap));\r
\r
    Status = SmmCoreGetMemoryMapMemoryAttributesTable (\r
               &MemoryMapSize,\r
               MemoryMap,\r
               &MapKey,\r
               &DescriptorSize,\r
               &DescriptorVersion\r
               );\r
    if (EFI_ERROR (Status)) {\r
      FreePool (MemoryMap);\r
    }\r
  } while (Status == EFI_BUFFER_TOO_SMALL);\r
\r
  //\r
  // Allocate MemoryAttributesTable\r
  //\r
  RuntimeEntryCount = MemoryMapSize/DescriptorSize;\r
  MemoryAttributesTableSize = sizeof(EDKII_PI_SMM_MEMORY_ATTRIBUTES_TABLE) + DescriptorSize * RuntimeEntryCount;\r
  MemoryAttributesTable = AllocatePool (sizeof(EDKII_PI_SMM_MEMORY_ATTRIBUTES_TABLE) + DescriptorSize * RuntimeEntryCount);\r
  ASSERT (MemoryAttributesTable != NULL);\r
  MemoryAttributesTable->Version         = EDKII_PI_SMM_MEMORY_ATTRIBUTES_TABLE_VERSION;\r
  MemoryAttributesTable->NumberOfEntries = (UINT32)RuntimeEntryCount;\r
  MemoryAttributesTable->DescriptorSize  = (UINT32)DescriptorSize;\r
  MemoryAttributesTable->Reserved        = 0;\r
  DEBUG ((DEBUG_INFO, "MemoryAttributesTable:\n"));\r
  DEBUG ((DEBUG_INFO, "  Version              - 0x%08x\n", MemoryAttributesTable->Version));\r
  DEBUG ((DEBUG_INFO, "  NumberOfEntries      - 0x%08x\n", MemoryAttributesTable->NumberOfEntries));\r
  DEBUG ((DEBUG_INFO, "  DescriptorSize       - 0x%08x\n", MemoryAttributesTable->DescriptorSize));\r
  MemoryAttributesEntry = (EFI_MEMORY_DESCRIPTOR *)(MemoryAttributesTable + 1);\r
  for (Index = 0; Index < MemoryMapSize/DescriptorSize; Index++) {\r
    CopyMem (MemoryAttributesEntry, MemoryMap, DescriptorSize);\r
    DEBUG ((DEBUG_INFO, "Entry (0x%x)\n", MemoryAttributesEntry));\r
    DEBUG ((DEBUG_INFO, "  Type              - 0x%x\n", MemoryAttributesEntry->Type));\r
    DEBUG ((DEBUG_INFO, "  PhysicalStart     - 0x%016lx\n", MemoryAttributesEntry->PhysicalStart));\r
    DEBUG ((DEBUG_INFO, "  VirtualStart      - 0x%016lx\n", MemoryAttributesEntry->VirtualStart));\r
    DEBUG ((DEBUG_INFO, "  NumberOfPages     - 0x%016lx\n", MemoryAttributesEntry->NumberOfPages));\r
    DEBUG ((DEBUG_INFO, "  Attribute         - 0x%016lx\n", MemoryAttributesEntry->Attribute));\r
    MemoryAttributesEntry = NEXT_MEMORY_DESCRIPTOR(MemoryAttributesEntry, DescriptorSize);\r
\r
    MemoryMap = NEXT_MEMORY_DESCRIPTOR(MemoryMap, DescriptorSize);\r
  }\r
\r
  Status = gSmst->SmmInstallConfigurationTable (gSmst, &gEdkiiPiSmmMemoryAttributesTableGuid, MemoryAttributesTable, MemoryAttributesTableSize);\r
  ASSERT_EFI_ERROR (Status);\r
}\r
\r
\r
/**\r
  This function installs all SMM image record information.\r
**/\r
VOID\r
SmmInstallImageRecord (\r
  VOID\r
  )\r
{\r
  EFI_STATUS                  Status;\r
  UINTN                       NoHandles;\r
  EFI_HANDLE                  *HandleBuffer;\r
  EFI_LOADED_IMAGE_PROTOCOL   *LoadedImage;\r
  UINTN                       Index;\r
  EFI_SMM_DRIVER_ENTRY        DriverEntry;\r
\r
  Status = SmmLocateHandleBuffer (\r
             ByProtocol,\r
             &gEfiLoadedImageProtocolGuid,\r
             NULL,\r
             &NoHandles,\r
             &HandleBuffer\r
             );\r
  if (EFI_ERROR (Status)) {\r
    return ;\r
  }\r
\r
  for (Index = 0; Index < NoHandles; Index++) {\r
    Status = gSmst->SmmHandleProtocol (\r
                      HandleBuffer[Index],\r
                      &gEfiLoadedImageProtocolGuid,\r
                      (VOID **)&LoadedImage\r
                      );\r
    if (EFI_ERROR (Status)) {\r
      continue;\r
    }\r
    DEBUG ((DEBUG_VERBOSE, "LoadedImage - 0x%x 0x%x ", LoadedImage->ImageBase, LoadedImage->ImageSize));\r
    {\r
      VOID *PdbPointer;\r
      PdbPointer = PeCoffLoaderGetPdbPointer (LoadedImage->ImageBase);\r
      if (PdbPointer != NULL) {\r
        DEBUG ((DEBUG_VERBOSE, "(%a) ", PdbPointer));\r
      }\r
    }\r
    DEBUG ((DEBUG_VERBOSE, "\n"));\r
    ZeroMem (&DriverEntry, sizeof(DriverEntry));\r
    DriverEntry.ImageBuffer  = (UINTN)LoadedImage->ImageBase;\r
    DriverEntry.NumberOfPage = EFI_SIZE_TO_PAGES((UINTN)LoadedImage->ImageSize);\r
    SmmInsertImageRecord (&DriverEntry);\r
  }\r
\r
  FreePool (HandleBuffer);\r
}\r
\r
/**\r
  Install MemoryAttributesTable.\r
\r
  @param[in] Protocol   Points to the protocol's unique identifier.\r
  @param[in] Interface  Points to the interface instance.\r
  @param[in] Handle     The handle on which the interface was installed.\r
\r
  @retval EFI_SUCCESS   Notification runs successfully.\r
**/\r
EFI_STATUS\r
EFIAPI\r
SmmInstallMemoryAttributesTable (\r
  IN CONST EFI_GUID  *Protocol,\r
  IN VOID            *Interface,\r
  IN EFI_HANDLE      Handle\r
  )\r
{\r
  SmmInstallImageRecord ();\r
\r
  DEBUG ((DEBUG_INFO, "SMM MemoryProtectionAttribute - 0x%016lx\n", mMemoryProtectionAttribute));\r
  if ((mMemoryProtectionAttribute & EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA) == 0) {\r
    return EFI_SUCCESS;\r
  }\r
\r
  DEBUG ((DEBUG_VERBOSE, "SMM Total Image Count - 0x%x\n", mImagePropertiesPrivateData.ImageRecordCount));\r
  DEBUG ((DEBUG_VERBOSE, "SMM Dump ImageRecord:\n"));\r
  DumpImageRecord ();\r
\r
  PublishMemoryAttributesTable ();\r
\r
  return EFI_SUCCESS;\r
}\r
\r
/**\r
  Initialize MemoryAttributesTable support.\r
**/\r
VOID\r
EFIAPI\r
SmmCoreInitializeMemoryAttributesTable (\r
  VOID\r
  )\r
{\r
  EFI_STATUS                        Status;\r
  VOID                              *Registration;\r
\r
  Status = gSmst->SmmRegisterProtocolNotify (\r
                    &gEfiSmmEndOfDxeProtocolGuid,\r
                    SmmInstallMemoryAttributesTable,\r
                    &Registration\r
                    );\r
  ASSERT_EFI_ERROR (Status);\r
\r
  return ;\r
}\r