MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c

   1 /** @file
   2   PI SMM MemoryAttributes support
   3
   4 Copyright (c) 2016, Intel Corporation. All rights reserved.<BR>
   5 This program and the accompanying materials
   6 are licensed and made available under the terms and conditions of the BSD License
   7 which accompanies this distribution.  The full text of the license may be found at
   8 http://opensource.org/licenses/bsd-license.php
   9
  10 THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  11 WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
  12
  13 **/
  14
  15 #include <PiDxe.h>
  16 #include <Library/BaseLib.h>
  17 #include <Library/BaseMemoryLib.h>
  18 #include <Library/MemoryAllocationLib.h>
  19 #include <Library/UefiBootServicesTableLib.h>
  20 #include <Library/SmmServicesTableLib.h>
  21 #include <Library/DebugLib.h>
  22 #include <Library/PcdLib.h>
  23
  24 #include <Library/PeCoffLib.h>
  25 #include <Library/PeCoffGetEntryPointLib.h>
  26
  27 #include <Guid/PiSmmMemoryAttributesTable.h>
  28
  29 #include "PiSmmCore.h"
  30
  31 #define PREVIOUS_MEMORY_DESCRIPTOR(MemoryDescriptor, Size) \
  32   ((EFI_MEMORY_DESCRIPTOR *)((UINT8 *)(MemoryDescriptor) - (Size)))
  33
  34 #define IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE SIGNATURE_32 ('I','P','R','C')
  35
  36 typedef struct {
  37   UINT32                 Signature;
  38   LIST_ENTRY             Link;
  39   EFI_PHYSICAL_ADDRESS   CodeSegmentBase;
  40   UINT64                 CodeSegmentSize;
  41 } IMAGE_PROPERTIES_RECORD_CODE_SECTION;
  42
  43 #define IMAGE_PROPERTIES_RECORD_SIGNATURE SIGNATURE_32 ('I','P','R','D')
  44
  45 typedef struct {
  46   UINT32                 Signature;
  47   LIST_ENTRY             Link;
  48   EFI_PHYSICAL_ADDRESS   ImageBase;
  49   UINT64                 ImageSize;
  50   UINTN                  CodeSegmentCount;
  51   LIST_ENTRY             CodeSegmentList;
  52 } IMAGE_PROPERTIES_RECORD;
  53
  54 #define IMAGE_PROPERTIES_PRIVATE_DATA_SIGNATURE SIGNATURE_32 ('I','P','P','D')
  55
  56 typedef struct {
  57   UINT32                 Signature;
  58   UINTN                  ImageRecordCount;
  59   UINTN                  CodeSegmentCountMax;
  60   LIST_ENTRY             ImageRecordList;
  61 } IMAGE_PROPERTIES_PRIVATE_DATA;
  62
  63 IMAGE_PROPERTIES_PRIVATE_DATA  mImagePropertiesPrivateData = {
  64   IMAGE_PROPERTIES_PRIVATE_DATA_SIGNATURE,
  65   0,
  66   0,
  67   INITIALIZE_LIST_HEAD_VARIABLE (mImagePropertiesPrivateData.ImageRecordList)
  68 };
  69
  70 #define EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA  BIT0
  71
  72 UINT64 mMemoryProtectionAttribute = EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA;
  73
  74 //
  75 // Below functions are for MemoryMap
  76 //
  77
  78 /**
  79   Converts a number of EFI_PAGEs to a size in bytes.
  80
  81   NOTE: Do not use EFI_PAGES_TO_SIZE because it handles UINTN only.
  82
  83   @param[in]  Pages     The number of EFI_PAGES.
  84
  85   @return  The number of bytes associated with the number of EFI_PAGEs specified
  86            by Pages.
  87 **/
  88 STATIC
  89 UINT64
  90 EfiPagesToSize (
  91   IN UINT64 Pages
  92   )
  93 {
  94   return LShiftU64 (Pages, EFI_PAGE_SHIFT);
  95 }
  96
  97 /**
  98   Converts a size, in bytes, to a number of EFI_PAGESs.
  99
 100   NOTE: Do not use EFI_SIZE_TO_PAGES because it handles UINTN only.
 101
 102   @param[in]  Size      A size in bytes.
 103
 104   @return  The number of EFI_PAGESs associated with the number of bytes specified
 105            by Size.
 106
 107 **/
 108 STATIC
 109 UINT64
 110 EfiSizeToPages (
 111   IN UINT64 Size
 112   )
 113 {
 114   return RShiftU64 (Size, EFI_PAGE_SHIFT) + ((((UINTN)Size) & EFI_PAGE_MASK) ? 1 : 0);
 115 }
 116
 117 /**
 118   Check the consistency of Smm memory attributes table.
 119
 120   @param[in] MemoryAttributesTable  PI SMM memory attributes table
 121 **/
 122 VOID
 123 SmmMemoryAttributesTableConsistencyCheck (
 124   IN EDKII_PI_SMM_MEMORY_ATTRIBUTES_TABLE *MemoryAttributesTable
 125   )
 126 {
 127   EFI_MEMORY_DESCRIPTOR                     *MemoryMap;
 128   UINTN                                     MemoryMapEntryCount;
 129   UINTN                                     DescriptorSize;
 130   UINTN                                     Index;
 131   UINT64                                    Address;
 132
 133   Address = 0;
 134   MemoryMapEntryCount = MemoryAttributesTable->NumberOfEntries;
 135   DescriptorSize = MemoryAttributesTable->DescriptorSize;
 136   MemoryMap = (EFI_MEMORY_DESCRIPTOR *)(MemoryAttributesTable + 1);
 137   for (Index = 0; Index < MemoryMapEntryCount; Index++) {
 138     if (Address != 0) {
 139       ASSERT (Address == MemoryMap->PhysicalStart);
 140     }
 141     Address = MemoryMap->PhysicalStart + EFI_PAGES_TO_SIZE(MemoryMap->NumberOfPages);
 142     MemoryMap = NEXT_MEMORY_DESCRIPTOR(MemoryMap, DescriptorSize);
 143   }
 144 }
 145
 146 /**
 147   Sort memory map entries based upon PhysicalStart, from low to high.
 148
 149   @param[in]  MemoryMap              A pointer to the buffer in which firmware places
 150                                  the current memory map.
 151   @param[in]  MemoryMapSize          Size, in bytes, of the MemoryMap buffer.
 152   @param[in]  DescriptorSize         Size, in bytes, of an individual EFI_MEMORY_DESCRIPTOR.
 153 **/
 154 STATIC
 155 VOID
 156 SortMemoryMap (
 157   IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,
 158   IN UINTN                      MemoryMapSize,
 159   IN UINTN                      DescriptorSize
 160   )
 161 {
 162   EFI_MEMORY_DESCRIPTOR       *MemoryMapEntry;
 163   EFI_MEMORY_DESCRIPTOR       *NextMemoryMapEntry;
 164   EFI_MEMORY_DESCRIPTOR       *MemoryMapEnd;
 165   EFI_MEMORY_DESCRIPTOR       TempMemoryMap;
 166
 167   MemoryMapEntry = MemoryMap;
 168   NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);
 169   MemoryMapEnd = (EFI_MEMORY_DESCRIPTOR *) ((UINT8 *) MemoryMap + MemoryMapSize);
 170   while (MemoryMapEntry < MemoryMapEnd) {
 171     while (NextMemoryMapEntry < MemoryMapEnd) {
 172       if (MemoryMapEntry->PhysicalStart > NextMemoryMapEntry->PhysicalStart) {
 173         CopyMem (&TempMemoryMap, MemoryMapEntry, sizeof(EFI_MEMORY_DESCRIPTOR));
 174         CopyMem (MemoryMapEntry, NextMemoryMapEntry, sizeof(EFI_MEMORY_DESCRIPTOR));
 175         CopyMem (NextMemoryMapEntry, &TempMemoryMap, sizeof(EFI_MEMORY_DESCRIPTOR));
 176       }
 177
 178       NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (NextMemoryMapEntry, DescriptorSize);
 179     }
 180
 181     MemoryMapEntry      = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);
 182     NextMemoryMapEntry  = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);
 183   }
 184
 185   return ;
 186 }
 187
 188 /**
 189   Merge continous memory map entries whose have same attributes.
 190
 191   @param[in, out]  MemoryMap              A pointer to the buffer in which firmware places
 192                                           the current memory map.
 193   @param[in, out]  MemoryMapSize          A pointer to the size, in bytes, of the
 194                                           MemoryMap buffer. On input, this is the size of
 195                                           the current memory map.  On output,
 196                                           it is the size of new memory map after merge.
 197   @param[in]       DescriptorSize         Size, in bytes, of an individual EFI_MEMORY_DESCRIPTOR.
 198 **/
 199 STATIC
 200 VOID
 201 MergeMemoryMap (
 202   IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,
 203   IN OUT UINTN                  *MemoryMapSize,
 204   IN UINTN                      DescriptorSize
 205   )
 206 {
 207   EFI_MEMORY_DESCRIPTOR       *MemoryMapEntry;
 208   EFI_MEMORY_DESCRIPTOR       *MemoryMapEnd;
 209   UINT64                      MemoryBlockLength;
 210   EFI_MEMORY_DESCRIPTOR       *NewMemoryMapEntry;
 211   EFI_MEMORY_DESCRIPTOR       *NextMemoryMapEntry;
 212
 213   MemoryMapEntry = MemoryMap;
 214   NewMemoryMapEntry = MemoryMap;
 215   MemoryMapEnd = (EFI_MEMORY_DESCRIPTOR *) ((UINT8 *) MemoryMap + *MemoryMapSize);
 216   while ((UINTN)MemoryMapEntry < (UINTN)MemoryMapEnd) {
 217     CopyMem (NewMemoryMapEntry, MemoryMapEntry, sizeof(EFI_MEMORY_DESCRIPTOR));
 218     NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);
 219
 220     do {
 221       MemoryBlockLength = (UINT64) (EfiPagesToSize (MemoryMapEntry->NumberOfPages));
 222       if (((UINTN)NextMemoryMapEntry < (UINTN)MemoryMapEnd) &&
 223           (MemoryMapEntry->Type == NextMemoryMapEntry->Type) &&
 224           (MemoryMapEntry->Attribute == NextMemoryMapEntry->Attribute) &&
 225           ((MemoryMapEntry->PhysicalStart + MemoryBlockLength) == NextMemoryMapEntry->PhysicalStart)) {
 226         MemoryMapEntry->NumberOfPages += NextMemoryMapEntry->NumberOfPages;
 227         if (NewMemoryMapEntry != MemoryMapEntry) {
 228           NewMemoryMapEntry->NumberOfPages += NextMemoryMapEntry->NumberOfPages;
 229         }
 230
 231         NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (NextMemoryMapEntry, DescriptorSize);
 232         continue;
 233       } else {
 234         MemoryMapEntry = PREVIOUS_MEMORY_DESCRIPTOR (NextMemoryMapEntry, DescriptorSize);
 235         break;
 236       }
 237     } while (TRUE);
 238
 239     MemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);
 240     NewMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (NewMemoryMapEntry, DescriptorSize);
 241   }
 242
 243   *MemoryMapSize = (UINTN)NewMemoryMapEntry - (UINTN)MemoryMap;
 244
 245   return ;
 246 }
 247
 248 /**
 249   Enforce memory map attributes.
 250   This function will set EfiRuntimeServicesData/EfiMemoryMappedIO/EfiMemoryMappedIOPortSpace to be EFI_MEMORY_XP.
 251
 252   @param[in, out]  MemoryMap              A pointer to the buffer in which firmware places
 253                                           the current memory map.
 254   @param[in]       MemoryMapSize          Size, in bytes, of the MemoryMap buffer.
 255   @param[in]       DescriptorSize         Size, in bytes, of an individual EFI_MEMORY_DESCRIPTOR.
 256 **/
 257 STATIC
 258 VOID
 259 EnforceMemoryMapAttribute (
 260   IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,
 261   IN UINTN                      MemoryMapSize,
 262   IN UINTN                      DescriptorSize
 263   )
 264 {
 265   EFI_MEMORY_DESCRIPTOR       *MemoryMapEntry;
 266   EFI_MEMORY_DESCRIPTOR       *MemoryMapEnd;
 267
 268   MemoryMapEntry = MemoryMap;
 269   MemoryMapEnd   = (EFI_MEMORY_DESCRIPTOR *) ((UINT8 *) MemoryMap + MemoryMapSize);
 270   while ((UINTN)MemoryMapEntry < (UINTN)MemoryMapEnd) {
 271     switch (MemoryMapEntry->Type) {
 272     case EfiRuntimeServicesCode:
 273       MemoryMapEntry->Attribute |= EFI_MEMORY_RO;
 274       break;
 275     case EfiRuntimeServicesData:
 276       MemoryMapEntry->Attribute |= EFI_MEMORY_XP;
 277       break;
 278     }
 279
 280     MemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);
 281   }
 282
 283   return ;
 284 }
 285
 286 /**
 287   Return the first image record, whose [ImageBase, ImageSize] covered by [Buffer, Length].
 288
 289   @param[in] Buffer  Start Address
 290   @param[in] Length  Address length
 291
 292   @return first image record covered by [buffer, length]
 293 **/
 294 STATIC
 295 IMAGE_PROPERTIES_RECORD *
 296 GetImageRecordByAddress (
 297   IN EFI_PHYSICAL_ADDRESS  Buffer,
 298   IN UINT64                Length
 299   )
 300 {
 301   IMAGE_PROPERTIES_RECORD    *ImageRecord;
 302   LIST_ENTRY                 *ImageRecordLink;
 303   LIST_ENTRY                 *ImageRecordList;
 304
 305   ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
 306
 307   for (ImageRecordLink = ImageRecordList->ForwardLink;
 308        ImageRecordLink != ImageRecordList;
 309        ImageRecordLink = ImageRecordLink->ForwardLink) {
 310     ImageRecord = CR (
 311                     ImageRecordLink,
 312                     IMAGE_PROPERTIES_RECORD,
 313                     Link,
 314                     IMAGE_PROPERTIES_RECORD_SIGNATURE
 315                     );
 316
 317     if ((Buffer <= ImageRecord->ImageBase) &&
 318         (Buffer + Length >= ImageRecord->ImageBase + ImageRecord->ImageSize)) {
 319       return ImageRecord;
 320     }
 321   }
 322
 323   return NULL;
 324 }
 325
 326 /**
 327   Set the memory map to new entries, according to one old entry,
 328   based upon PE code section and data section in image record
 329
 330   @param[in]       ImageRecord            An image record whose [ImageBase, ImageSize] covered
 331                                           by old memory map entry.
 332   @param[in, out]  NewRecord              A pointer to several new memory map entries.
 333                                           The caller gurantee the buffer size be 1 +
 334                                           (SplitRecordCount * DescriptorSize) calculated
 335                                           below.
 336   @param[in]       OldRecord              A pointer to one old memory map entry.
 337   @param[in]       DescriptorSize         Size, in bytes, of an individual EFI_MEMORY_DESCRIPTOR.
 338 **/
 339 STATIC
 340 UINTN
 341 SetNewRecord (
 342   IN IMAGE_PROPERTIES_RECORD       *ImageRecord,
 343   IN OUT EFI_MEMORY_DESCRIPTOR     *NewRecord,
 344   IN EFI_MEMORY_DESCRIPTOR         *OldRecord,
 345   IN UINTN                         DescriptorSize
 346   )
 347 {
 348   EFI_MEMORY_DESCRIPTOR                     TempRecord;
 349   IMAGE_PROPERTIES_RECORD_CODE_SECTION      *ImageRecordCodeSection;
 350   LIST_ENTRY                                *ImageRecordCodeSectionLink;
 351   LIST_ENTRY                                *ImageRecordCodeSectionEndLink;
 352   LIST_ENTRY                                *ImageRecordCodeSectionList;
 353   UINTN                                     NewRecordCount;
 354   UINT64                                    PhysicalEnd;
 355   UINT64                                    ImageEnd;
 356
 357   CopyMem (&TempRecord, OldRecord, sizeof(EFI_MEMORY_DESCRIPTOR));
 358   PhysicalEnd = TempRecord.PhysicalStart + EfiPagesToSize(TempRecord.NumberOfPages);
 359   NewRecordCount = 0;
 360
 361   ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;
 362
 363   ImageRecordCodeSectionLink = ImageRecordCodeSectionList->ForwardLink;
 364   ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;
 365   while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
 366     ImageRecordCodeSection = CR (
 367                                ImageRecordCodeSectionLink,
 368                                IMAGE_PROPERTIES_RECORD_CODE_SECTION,
 369                                Link,
 370                                IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
 371                                );
 372     ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
 373
 374     if (TempRecord.PhysicalStart <= ImageRecordCodeSection->CodeSegmentBase) {
 375       //
 376       // DATA
 377       //
 378       NewRecord->Type = EfiRuntimeServicesData;
 379       NewRecord->PhysicalStart = TempRecord.PhysicalStart;
 380       NewRecord->VirtualStart  = 0;
 381       NewRecord->NumberOfPages = EfiSizeToPages(ImageRecordCodeSection->CodeSegmentBase - NewRecord->PhysicalStart);
 382       NewRecord->Attribute     = TempRecord.Attribute | EFI_MEMORY_XP;
 383       if (NewRecord->NumberOfPages != 0) {
 384         NewRecord = NEXT_MEMORY_DESCRIPTOR (NewRecord, DescriptorSize);
 385         NewRecordCount ++;
 386       }
 387
 388       //
 389       // CODE
 390       //
 391       NewRecord->Type = EfiRuntimeServicesCode;
 392       NewRecord->PhysicalStart = ImageRecordCodeSection->CodeSegmentBase;
 393       NewRecord->VirtualStart  = 0;
 394       NewRecord->NumberOfPages = EfiSizeToPages(ImageRecordCodeSection->CodeSegmentSize);
 395       NewRecord->Attribute     = (TempRecord.Attribute & (~EFI_MEMORY_XP)) | EFI_MEMORY_RO;
 396       if (NewRecord->NumberOfPages != 0) {
 397         NewRecord = NEXT_MEMORY_DESCRIPTOR (NewRecord, DescriptorSize);
 398         NewRecordCount ++;
 399       }
 400
 401       TempRecord.PhysicalStart = ImageRecordCodeSection->CodeSegmentBase + EfiPagesToSize (EfiSizeToPages(ImageRecordCodeSection->CodeSegmentSize));
 402       TempRecord.NumberOfPages = EfiSizeToPages(PhysicalEnd - TempRecord.PhysicalStart);
 403       if (TempRecord.NumberOfPages == 0) {
 404         break;
 405       }
 406     }
 407   }
 408
 409   ImageEnd = ImageRecord->ImageBase + ImageRecord->ImageSize;
 410
 411   //
 412   // Final DATA
 413   //
 414   if (TempRecord.PhysicalStart < ImageEnd) {
 415     NewRecord->Type = EfiRuntimeServicesData;
 416     NewRecord->PhysicalStart = TempRecord.PhysicalStart;
 417     NewRecord->VirtualStart  = 0;
 418     NewRecord->NumberOfPages = EfiSizeToPages (ImageEnd - TempRecord.PhysicalStart);
 419     NewRecord->Attribute     = TempRecord.Attribute | EFI_MEMORY_XP;
 420     NewRecordCount ++;
 421   }
 422
 423   return NewRecordCount;
 424 }
 425
 426 /**
 427   Return the max number of new splitted entries, according to one old entry,
 428   based upon PE code section and data section.
 429
 430   @param[in]  OldRecord              A pointer to one old memory map entry.
 431
 432   @retval  0 no entry need to be splitted.
 433   @return  the max number of new splitted entries
 434 **/
 435 STATIC
 436 UINTN
 437 GetMaxSplitRecordCount (
 438   IN EFI_MEMORY_DESCRIPTOR *OldRecord
 439   )
 440 {
 441   IMAGE_PROPERTIES_RECORD *ImageRecord;
 442   UINTN                   SplitRecordCount;
 443   UINT64                  PhysicalStart;
 444   UINT64                  PhysicalEnd;
 445
 446   SplitRecordCount = 0;
 447   PhysicalStart = OldRecord->PhysicalStart;
 448   PhysicalEnd = OldRecord->PhysicalStart + EfiPagesToSize(OldRecord->NumberOfPages);
 449
 450   do {
 451     ImageRecord = GetImageRecordByAddress (PhysicalStart, PhysicalEnd - PhysicalStart);
 452     if (ImageRecord == NULL) {
 453       break;
 454     }
 455     SplitRecordCount += (2 * ImageRecord->CodeSegmentCount + 1);
 456     PhysicalStart = ImageRecord->ImageBase + ImageRecord->ImageSize;
 457   } while ((ImageRecord != NULL) && (PhysicalStart < PhysicalEnd));
 458
 459   if (SplitRecordCount != 0) {
 460     SplitRecordCount--;
 461   }
 462
 463   return SplitRecordCount;
 464 }
 465
 466 /**
 467   Split the memory map to new entries, according to one old entry,
 468   based upon PE code section and data section.
 469
 470   @param[in]       OldRecord              A pointer to one old memory map entry.
 471   @param[in, out]  NewRecord              A pointer to several new memory map entries.
 472                                           The caller gurantee the buffer size be 1 +
 473                                           (SplitRecordCount * DescriptorSize) calculated
 474                                           below.
 475   @param[in]       MaxSplitRecordCount    The max number of splitted entries
 476   @param[in]       DescriptorSize         Size, in bytes, of an individual EFI_MEMORY_DESCRIPTOR.
 477
 478   @retval  0 no entry is splitted.
 479   @return  the real number of splitted record.
 480 **/
 481 STATIC
 482 UINTN
 483 SplitRecord (
 484   IN EFI_MEMORY_DESCRIPTOR     *OldRecord,
 485   IN OUT EFI_MEMORY_DESCRIPTOR *NewRecord,
 486   IN UINTN                     MaxSplitRecordCount,
 487   IN UINTN                     DescriptorSize
 488   )
 489 {
 490   EFI_MEMORY_DESCRIPTOR   TempRecord;
 491   IMAGE_PROPERTIES_RECORD *ImageRecord;
 492   IMAGE_PROPERTIES_RECORD *NewImageRecord;
 493   UINT64                  PhysicalStart;
 494   UINT64                  PhysicalEnd;
 495   UINTN                   NewRecordCount;
 496   UINTN                   TotalNewRecordCount;
 497
 498   if (MaxSplitRecordCount == 0) {
 499     CopyMem (NewRecord, OldRecord, DescriptorSize);
 500     return 0;
 501   }
 502
 503   TotalNewRecordCount = 0;
 504
 505   //
 506   // Override previous record
 507   //
 508   CopyMem (&TempRecord, OldRecord, sizeof(EFI_MEMORY_DESCRIPTOR));
 509   PhysicalStart = TempRecord.PhysicalStart;
 510   PhysicalEnd = TempRecord.PhysicalStart + EfiPagesToSize(TempRecord.NumberOfPages);
 511
 512   ImageRecord = NULL;
 513   do {
 514     NewImageRecord = GetImageRecordByAddress (PhysicalStart, PhysicalEnd - PhysicalStart);
 515     if (NewImageRecord == NULL) {
 516       //
 517       // No more image covered by this range, stop
 518       //
 519       if ((PhysicalEnd > PhysicalStart) && (ImageRecord != NULL)) {
 520         //
 521         // If this is still address in this record, need record.
 522         //
 523         NewRecord = PREVIOUS_MEMORY_DESCRIPTOR (NewRecord, DescriptorSize);
 524         if (NewRecord->Type == EfiRuntimeServicesData) {
 525           //
 526           // Last record is DATA, just merge it.
 527           //
 528           NewRecord->NumberOfPages = EfiSizeToPages(PhysicalEnd - NewRecord->PhysicalStart);
 529         } else {
 530           //
 531           // Last record is CODE, create a new DATA entry.
 532           //
 533           NewRecord = NEXT_MEMORY_DESCRIPTOR (NewRecord, DescriptorSize);
 534           NewRecord->Type = EfiRuntimeServicesData;
 535           NewRecord->PhysicalStart = TempRecord.PhysicalStart;
 536           NewRecord->VirtualStart  = 0;
 537           NewRecord->NumberOfPages = TempRecord.NumberOfPages;
 538           NewRecord->Attribute     = TempRecord.Attribute | EFI_MEMORY_XP;
 539           TotalNewRecordCount ++;
 540         }
 541       }
 542       break;
 543     }
 544     ImageRecord = NewImageRecord;
 545
 546     //
 547     // Set new record
 548     //
 549     NewRecordCount = SetNewRecord (ImageRecord, NewRecord, &TempRecord, DescriptorSize);
 550     TotalNewRecordCount += NewRecordCount;
 551     NewRecord = (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)NewRecord + NewRecordCount * DescriptorSize);
 552
 553     //
 554     // Update PhysicalStart, in order to exclude the image buffer already splitted.
 555     //
 556     PhysicalStart = ImageRecord->ImageBase + ImageRecord->ImageSize;
 557     TempRecord.PhysicalStart = PhysicalStart;
 558     TempRecord.NumberOfPages = EfiSizeToPages (PhysicalEnd - PhysicalStart);
 559   } while ((ImageRecord != NULL) && (PhysicalStart < PhysicalEnd));
 560
 561   return TotalNewRecordCount - 1;
 562 }
 563
 564 /**
 565   Split the original memory map, and add more entries to describe PE code section and data section.
 566   This function will set EfiRuntimeServicesData to be EFI_MEMORY_XP.
 567   This function will merge entries with same attributes finally.
 568
 569   NOTE: It assumes PE code/data section are page aligned.
 570   NOTE: It assumes enough entry is prepared for new memory map.
 571
 572   Split table:
 573    +---------------+
 574    | Record X      |
 575    +---------------+
 576    | Record RtCode |
 577    +---------------+
 578    | Record Y      |
 579    +---------------+
 580    ==>
 581    +---------------+
 582    | Record X      |
 583    +---------------+ ----
 584    | Record RtData |     |
 585    +---------------+     |
 586    | Record RtCode |     |-> PE/COFF1
 587    +---------------+     |
 588    | Record RtData |     |
 589    +---------------+ ----
 590    | Record RtData |     |
 591    +---------------+     |
 592    | Record RtCode |     |-> PE/COFF2
 593    +---------------+     |
 594    | Record RtData |     |
 595    +---------------+ ----
 596    | Record Y      |
 597    +---------------+
 598
 599   @param[in, out]  MemoryMapSize          A pointer to the size, in bytes, of the
 600                                           MemoryMap buffer. On input, this is the size of
 601                                           old MemoryMap before split. The actual buffer
 602                                           size of MemoryMap is MemoryMapSize +
 603                                           (AdditionalRecordCount * DescriptorSize) calculated
 604                                           below. On output, it is the size of new MemoryMap
 605                                           after split.
 606   @param[in, out]  MemoryMap              A pointer to the buffer in which firmware places
 607                                           the current memory map.
 608   @param[in]       DescriptorSize         Size, in bytes, of an individual EFI_MEMORY_DESCRIPTOR.
 609 **/
 610 STATIC
 611 VOID
 612 SplitTable (
 613   IN OUT UINTN                  *MemoryMapSize,
 614   IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,
 615   IN UINTN                      DescriptorSize
 616   )
 617 {
 618   INTN        IndexOld;
 619   INTN        IndexNew;
 620   UINTN       MaxSplitRecordCount;
 621   UINTN       RealSplitRecordCount;
 622   UINTN       TotalSplitRecordCount;
 623   UINTN       AdditionalRecordCount;
 624
 625   AdditionalRecordCount = (2 * mImagePropertiesPrivateData.CodeSegmentCountMax + 1) * mImagePropertiesPrivateData.ImageRecordCount;
 626
 627   TotalSplitRecordCount = 0;
 628   //
 629   // Let old record point to end of valid MemoryMap buffer.
 630   //
 631   IndexOld = ((*MemoryMapSize) / DescriptorSize) - 1;
 632   //
 633   // Let new record point to end of full MemoryMap buffer.
 634   //
 635   IndexNew = ((*MemoryMapSize) / DescriptorSize) - 1 + AdditionalRecordCount;
 636   for (; IndexOld >= 0; IndexOld--) {
 637     MaxSplitRecordCount = GetMaxSplitRecordCount ((EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + IndexOld * DescriptorSize));
 638     //
 639     // Split this MemoryMap record
 640     //
 641     IndexNew -= MaxSplitRecordCount;
 642     RealSplitRecordCount = SplitRecord (
 643                              (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + IndexOld * DescriptorSize),
 644                              (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + IndexNew * DescriptorSize),
 645                              MaxSplitRecordCount,
 646                              DescriptorSize
 647                              );
 648     //
 649     // Adjust IndexNew according to real split.
 650     //
 651     CopyMem (
 652       ((UINT8 *)MemoryMap + (IndexNew + MaxSplitRecordCount - RealSplitRecordCount) * DescriptorSize),
 653       ((UINT8 *)MemoryMap + IndexNew * DescriptorSize),
 654       RealSplitRecordCount * DescriptorSize
 655       );
 656     IndexNew = IndexNew + MaxSplitRecordCount - RealSplitRecordCount;
 657     TotalSplitRecordCount += RealSplitRecordCount;
 658     IndexNew --;
 659   }
 660   //
 661   // Move all records to the beginning.
 662   //
 663   CopyMem (
 664     MemoryMap,
 665     (UINT8 *)MemoryMap + (AdditionalRecordCount - TotalSplitRecordCount) * DescriptorSize,
 666     (*MemoryMapSize) + TotalSplitRecordCount * DescriptorSize
 667     );
 668
 669   *MemoryMapSize = (*MemoryMapSize) + DescriptorSize * TotalSplitRecordCount;
 670
 671   //
 672   // Sort from low to high (Just in case)
 673   //
 674   SortMemoryMap (MemoryMap, *MemoryMapSize, DescriptorSize);
 675
 676   //
 677   // Set RuntimeData to XP
 678   //
 679   EnforceMemoryMapAttribute (MemoryMap, *MemoryMapSize, DescriptorSize);
 680
 681   //
 682   // Merge same type to save entry size
 683   //
 684   MergeMemoryMap (MemoryMap, MemoryMapSize, DescriptorSize);
 685
 686   return ;
 687 }
 688
 689 /**
 690   This function for GetMemoryMap() with memory attributes table.
 691
 692   It calls original GetMemoryMap() to get the original memory map information. Then
 693   plus the additional memory map entries for PE Code/Data seperation.
 694
 695   @param[in, out]  MemoryMapSize          A pointer to the size, in bytes, of the
 696                                           MemoryMap buffer. On input, this is the size of
 697                                           the buffer allocated by the caller.  On output,
 698                                           it is the size of the buffer returned by the
 699                                           firmware  if the buffer was large enough, or the
 700                                           size of the buffer needed  to contain the map if
 701                                           the buffer was too small.
 702   @param[in, out]  MemoryMap              A pointer to the buffer in which firmware places
 703                                           the current memory map.
 704   @param[out]      MapKey                 A pointer to the location in which firmware
 705                                           returns the key for the current memory map.
 706   @param[out]      DescriptorSize         A pointer to the location in which firmware
 707                                           returns the size, in bytes, of an individual
 708                                           EFI_MEMORY_DESCRIPTOR.
 709   @param[out]      DescriptorVersion      A pointer to the location in which firmware
 710                                           returns the version number associated with the
 711                                           EFI_MEMORY_DESCRIPTOR.
 712
 713   @retval EFI_SUCCESS            The memory map was returned in the MemoryMap
 714                                  buffer.
 715   @retval EFI_BUFFER_TOO_SMALL   The MemoryMap buffer was too small. The current
 716                                  buffer size needed to hold the memory map is
 717                                  returned in MemoryMapSize.
 718   @retval EFI_INVALID_PARAMETER  One of the parameters has an invalid value.
 719
 720 **/
 721 STATIC
 722 EFI_STATUS
 723 EFIAPI
 724 SmmCoreGetMemoryMapMemoryAttributesTable (
 725   IN OUT UINTN                  *MemoryMapSize,
 726   IN OUT EFI_MEMORY_DESCRIPTOR  *MemoryMap,
 727   OUT UINTN                     *MapKey,
 728   OUT UINTN                     *DescriptorSize,
 729   OUT UINT32                    *DescriptorVersion
 730   )
 731 {
 732   EFI_STATUS  Status;
 733   UINTN       OldMemoryMapSize;
 734   UINTN       AdditionalRecordCount;
 735
 736   //
 737   // If PE code/data is not aligned, just return.
 738   //
 739   if ((mMemoryProtectionAttribute & EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA) == 0) {
 740     return SmmCoreGetMemoryMap (MemoryMapSize, MemoryMap, MapKey, DescriptorSize, DescriptorVersion);
 741   }
 742
 743   if (MemoryMapSize == NULL) {
 744     return EFI_INVALID_PARAMETER;
 745   }
 746
 747   AdditionalRecordCount = (2 * mImagePropertiesPrivateData.CodeSegmentCountMax + 1) * mImagePropertiesPrivateData.ImageRecordCount;
 748
 749   OldMemoryMapSize = *MemoryMapSize;
 750   Status = SmmCoreGetMemoryMap (MemoryMapSize, MemoryMap, MapKey, DescriptorSize, DescriptorVersion);
 751   if (Status == EFI_BUFFER_TOO_SMALL) {
 752     *MemoryMapSize = *MemoryMapSize + (*DescriptorSize) * AdditionalRecordCount;
 753   } else if (Status == EFI_SUCCESS) {
 754     if (OldMemoryMapSize - *MemoryMapSize < (*DescriptorSize) * AdditionalRecordCount) {
 755       *MemoryMapSize = *MemoryMapSize + (*DescriptorSize) * AdditionalRecordCount;
 756       //
 757       // Need update status to buffer too small
 758       //
 759       Status = EFI_BUFFER_TOO_SMALL;
 760     } else {
 761       //
 762       // Split PE code/data
 763       //
 764       ASSERT(MemoryMap != NULL);
 765       SplitTable (MemoryMapSize, MemoryMap, *DescriptorSize);
 766     }
 767   }
 768
 769   return Status;
 770 }
 771
 772 //
 773 // Below functions are for ImageRecord
 774 //
 775
 776 /**
 777   Set MemoryProtectionAttribute accroding to PE/COFF image section alignment.
 778
 779   @param[in]  SectionAlignment    PE/COFF section alignment
 780 **/
 781 STATIC
 782 VOID
 783 SetMemoryAttributesTableSectionAlignment (
 784   IN UINT32  SectionAlignment
 785   )
 786 {
 787   if (((SectionAlignment & (EFI_ACPI_RUNTIME_PAGE_ALLOCATION_ALIGNMENT - 1)) != 0) &&
 788       ((mMemoryProtectionAttribute & EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA) != 0)) {
 789     DEBUG ((DEBUG_VERBOSE, "SMM SetMemoryAttributesTableSectionAlignment - Clear\n"));
 790     mMemoryProtectionAttribute &= ~((UINT64)EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA);
 791   }
 792 }
 793
 794 /**
 795   Swap two code sections in image record.
 796
 797   @param[in]  FirstImageRecordCodeSection    first code section in image record
 798   @param[in]  SecondImageRecordCodeSection   second code section in image record
 799 **/
 800 STATIC
 801 VOID
 802 SwapImageRecordCodeSection (
 803   IN IMAGE_PROPERTIES_RECORD_CODE_SECTION      *FirstImageRecordCodeSection,
 804   IN IMAGE_PROPERTIES_RECORD_CODE_SECTION      *SecondImageRecordCodeSection
 805   )
 806 {
 807   IMAGE_PROPERTIES_RECORD_CODE_SECTION      TempImageRecordCodeSection;
 808
 809   TempImageRecordCodeSection.CodeSegmentBase = FirstImageRecordCodeSection->CodeSegmentBase;
 810   TempImageRecordCodeSection.CodeSegmentSize = FirstImageRecordCodeSection->CodeSegmentSize;
 811
 812   FirstImageRecordCodeSection->CodeSegmentBase = SecondImageRecordCodeSection->CodeSegmentBase;
 813   FirstImageRecordCodeSection->CodeSegmentSize = SecondImageRecordCodeSection->CodeSegmentSize;
 814
 815   SecondImageRecordCodeSection->CodeSegmentBase = TempImageRecordCodeSection.CodeSegmentBase;
 816   SecondImageRecordCodeSection->CodeSegmentSize = TempImageRecordCodeSection.CodeSegmentSize;
 817 }
 818
 819 /**
 820   Sort code section in image record, based upon CodeSegmentBase from low to high.
 821
 822   @param[in]  ImageRecord    image record to be sorted
 823 **/
 824 STATIC
 825 VOID
 826 SortImageRecordCodeSection (
 827   IN IMAGE_PROPERTIES_RECORD              *ImageRecord
 828   )
 829 {
 830   IMAGE_PROPERTIES_RECORD_CODE_SECTION      *ImageRecordCodeSection;
 831   IMAGE_PROPERTIES_RECORD_CODE_SECTION      *NextImageRecordCodeSection;
 832   LIST_ENTRY                                *ImageRecordCodeSectionLink;
 833   LIST_ENTRY                                *NextImageRecordCodeSectionLink;
 834   LIST_ENTRY                                *ImageRecordCodeSectionEndLink;
 835   LIST_ENTRY                                *ImageRecordCodeSectionList;
 836
 837   ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;
 838
 839   ImageRecordCodeSectionLink = ImageRecordCodeSectionList->ForwardLink;
 840   NextImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
 841   ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;
 842   while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
 843     ImageRecordCodeSection = CR (
 844                                ImageRecordCodeSectionLink,
 845                                IMAGE_PROPERTIES_RECORD_CODE_SECTION,
 846                                Link,
 847                                IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
 848                                );
 849     while (NextImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
 850       NextImageRecordCodeSection = CR (
 851                                      NextImageRecordCodeSectionLink,
 852                                      IMAGE_PROPERTIES_RECORD_CODE_SECTION,
 853                                      Link,
 854                                      IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
 855                                      );
 856       if (ImageRecordCodeSection->CodeSegmentBase > NextImageRecordCodeSection->CodeSegmentBase) {
 857         SwapImageRecordCodeSection (ImageRecordCodeSection, NextImageRecordCodeSection);
 858       }
 859       NextImageRecordCodeSectionLink = NextImageRecordCodeSectionLink->ForwardLink;
 860     }
 861
 862     ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
 863     NextImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
 864   }
 865 }
 866
 867 /**
 868   Check if code section in image record is valid.
 869
 870   @param[in]  ImageRecord    image record to be checked
 871
 872   @retval TRUE  image record is valid
 873   @retval FALSE image record is invalid
 874 **/
 875 STATIC
 876 BOOLEAN
 877 IsImageRecordCodeSectionValid (
 878   IN IMAGE_PROPERTIES_RECORD              *ImageRecord
 879   )
 880 {
 881   IMAGE_PROPERTIES_RECORD_CODE_SECTION      *ImageRecordCodeSection;
 882   IMAGE_PROPERTIES_RECORD_CODE_SECTION      *LastImageRecordCodeSection;
 883   LIST_ENTRY                                *ImageRecordCodeSectionLink;
 884   LIST_ENTRY                                *ImageRecordCodeSectionEndLink;
 885   LIST_ENTRY                                *ImageRecordCodeSectionList;
 886
 887   DEBUG ((DEBUG_VERBOSE, "SMM ImageCode SegmentCount - 0x%x\n", ImageRecord->CodeSegmentCount));
 888
 889   ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;
 890
 891   ImageRecordCodeSectionLink = ImageRecordCodeSectionList->ForwardLink;
 892   ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;
 893   LastImageRecordCodeSection = NULL;
 894   while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
 895     ImageRecordCodeSection = CR (
 896                                ImageRecordCodeSectionLink,
 897                                IMAGE_PROPERTIES_RECORD_CODE_SECTION,
 898                                Link,
 899                                IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
 900                                );
 901     if (ImageRecordCodeSection->CodeSegmentSize == 0) {
 902       return FALSE;
 903     }
 904     if (ImageRecordCodeSection->CodeSegmentBase < ImageRecord->ImageBase) {
 905       return FALSE;
 906     }
 907     if (ImageRecordCodeSection->CodeSegmentBase >= MAX_ADDRESS - ImageRecordCodeSection->CodeSegmentSize) {
 908       return FALSE;
 909     }
 910     if ((ImageRecordCodeSection->CodeSegmentBase + ImageRecordCodeSection->CodeSegmentSize) > (ImageRecord->ImageBase + ImageRecord->ImageSize)) {
 911       return FALSE;
 912     }
 913     if (LastImageRecordCodeSection != NULL) {
 914       if ((LastImageRecordCodeSection->CodeSegmentBase + LastImageRecordCodeSection->CodeSegmentSize) > ImageRecordCodeSection->CodeSegmentBase) {
 915         return FALSE;
 916       }
 917     }
 918
 919     LastImageRecordCodeSection = ImageRecordCodeSection;
 920     ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
 921   }
 922
 923   return TRUE;
 924 }
 925
 926 /**
 927   Swap two image records.
 928
 929   @param[in]  FirstImageRecord   first image record.
 930   @param[in]  SecondImageRecord  second image record.
 931 **/
 932 STATIC
 933 VOID
 934 SwapImageRecord (
 935   IN IMAGE_PROPERTIES_RECORD      *FirstImageRecord,
 936   IN IMAGE_PROPERTIES_RECORD      *SecondImageRecord
 937   )
 938 {
 939   IMAGE_PROPERTIES_RECORD      TempImageRecord;
 940
 941   TempImageRecord.ImageBase = FirstImageRecord->ImageBase;
 942   TempImageRecord.ImageSize = FirstImageRecord->ImageSize;
 943   TempImageRecord.CodeSegmentCount = FirstImageRecord->CodeSegmentCount;
 944
 945   FirstImageRecord->ImageBase = SecondImageRecord->ImageBase;
 946   FirstImageRecord->ImageSize = SecondImageRecord->ImageSize;
 947   FirstImageRecord->CodeSegmentCount = SecondImageRecord->CodeSegmentCount;
 948
 949   SecondImageRecord->ImageBase = TempImageRecord.ImageBase;
 950   SecondImageRecord->ImageSize = TempImageRecord.ImageSize;
 951   SecondImageRecord->CodeSegmentCount = TempImageRecord.CodeSegmentCount;
 952
 953   SwapListEntries (&FirstImageRecord->CodeSegmentList, &SecondImageRecord->CodeSegmentList);
 954 }
 955
 956 /**
 957   Sort image record based upon the ImageBase from low to high.
 958 **/
 959 STATIC
 960 VOID
 961 SortImageRecord (
 962   VOID
 963   )
 964 {
 965   IMAGE_PROPERTIES_RECORD      *ImageRecord;
 966   IMAGE_PROPERTIES_RECORD      *NextImageRecord;
 967   LIST_ENTRY                   *ImageRecordLink;
 968   LIST_ENTRY                   *NextImageRecordLink;
 969   LIST_ENTRY                   *ImageRecordEndLink;
 970   LIST_ENTRY                   *ImageRecordList;
 971
 972   ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
 973
 974   ImageRecordLink = ImageRecordList->ForwardLink;
 975   NextImageRecordLink = ImageRecordLink->ForwardLink;
 976   ImageRecordEndLink = ImageRecordList;
 977   while (ImageRecordLink != ImageRecordEndLink) {
 978     ImageRecord = CR (
 979                     ImageRecordLink,
 980                     IMAGE_PROPERTIES_RECORD,
 981                     Link,
 982                     IMAGE_PROPERTIES_RECORD_SIGNATURE
 983                     );
 984     while (NextImageRecordLink != ImageRecordEndLink) {
 985       NextImageRecord = CR (
 986                           NextImageRecordLink,
 987                           IMAGE_PROPERTIES_RECORD,
 988                           Link,
 989                           IMAGE_PROPERTIES_RECORD_SIGNATURE
 990                           );
 991       if (ImageRecord->ImageBase > NextImageRecord->ImageBase) {
 992         SwapImageRecord (ImageRecord, NextImageRecord);
 993       }
 994       NextImageRecordLink = NextImageRecordLink->ForwardLink;
 995     }
 996
 997     ImageRecordLink = ImageRecordLink->ForwardLink;
 998     NextImageRecordLink = ImageRecordLink->ForwardLink;
 999   }
1000 }
1001
1002 /**
1003   Dump image record.
1004 **/
1005 STATIC
1006 VOID
1007 DumpImageRecord (
1008   VOID
1009   )
1010 {
1011   IMAGE_PROPERTIES_RECORD      *ImageRecord;
1012   LIST_ENTRY                   *ImageRecordLink;
1013   LIST_ENTRY                   *ImageRecordList;
1014   UINTN                        Index;
1015
1016   ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
1017
1018   for (ImageRecordLink = ImageRecordList->ForwardLink, Index= 0;
1019        ImageRecordLink != ImageRecordList;
1020        ImageRecordLink = ImageRecordLink->ForwardLink, Index++) {
1021     ImageRecord = CR (
1022                     ImageRecordLink,
1023                     IMAGE_PROPERTIES_RECORD,
1024                     Link,
1025                     IMAGE_PROPERTIES_RECORD_SIGNATURE
1026                     );
1027     DEBUG ((DEBUG_VERBOSE, "SMM  Image[%d]: 0x%016lx - 0x%016lx\n", Index, ImageRecord->ImageBase, ImageRecord->ImageSize));
1028   }
1029 }
1030
1031 /**
1032   Insert image record.
1033
1034   @param[in]  DriverEntry    Driver information
1035 **/
1036 VOID
1037 SmmInsertImageRecord (
1038   IN EFI_SMM_DRIVER_ENTRY  *DriverEntry
1039   )
1040 {
1041   VOID                                 *ImageAddress;
1042   EFI_IMAGE_DOS_HEADER                 *DosHdr;
1043   UINT32                               PeCoffHeaderOffset;
1044   UINT32                               SectionAlignment;
1045   EFI_IMAGE_SECTION_HEADER             *Section;
1046   EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
1047   UINT8                                *Name;
1048   UINTN                                Index;
1049   IMAGE_PROPERTIES_RECORD              *ImageRecord;
1050   CHAR8                                *PdbPointer;
1051   IMAGE_PROPERTIES_RECORD_CODE_SECTION *ImageRecordCodeSection;
1052   UINT16                               Magic;
1053
1054   DEBUG ((DEBUG_VERBOSE, "SMM InsertImageRecord - 0x%x\n", DriverEntry));
1055   DEBUG ((DEBUG_VERBOSE, "SMM InsertImageRecord - 0x%016lx - 0x%08x\n", DriverEntry->ImageBuffer, DriverEntry->NumberOfPage));
1056
1057   ImageRecord = AllocatePool (sizeof(*ImageRecord));
1058   if (ImageRecord == NULL) {
1059     return ;
1060   }
1061   ImageRecord->Signature = IMAGE_PROPERTIES_RECORD_SIGNATURE;
1062
1063   DEBUG ((DEBUG_VERBOSE, "SMM ImageRecordCount - 0x%x\n", mImagePropertiesPrivateData.ImageRecordCount));
1064
1065   //
1066   // Step 1: record whole region
1067   //
1068   ImageRecord->ImageBase = DriverEntry->ImageBuffer;
1069   ImageRecord->ImageSize = EFI_PAGES_TO_SIZE(DriverEntry->NumberOfPage);
1070
1071   ImageAddress = (VOID *)(UINTN)DriverEntry->ImageBuffer;
1072
1073   PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);
1074   if (PdbPointer != NULL) {
1075     DEBUG ((DEBUG_VERBOSE, "SMM   Image - %a\n", PdbPointer));
1076   }
1077
1078   //
1079   // Check PE/COFF image
1080   //
1081   DosHdr = (EFI_IMAGE_DOS_HEADER *) (UINTN) ImageAddress;
1082   PeCoffHeaderOffset = 0;
1083   if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
1084     PeCoffHeaderOffset = DosHdr->e_lfanew;
1085   }
1086
1087   Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINT8 *) (UINTN) ImageAddress + PeCoffHeaderOffset);
1088   if (Hdr.Pe32->Signature != EFI_IMAGE_NT_SIGNATURE) {
1089     DEBUG ((DEBUG_VERBOSE, "SMM Hdr.Pe32->Signature invalid - 0x%x\n", Hdr.Pe32->Signature));
1090     goto Finish;
1091   }
1092
1093   //
1094   // Get SectionAlignment
1095   //
1096   if (Hdr.Pe32->FileHeader.Machine == IMAGE_FILE_MACHINE_IA64 && Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
1097     //
1098     // NOTE: Some versions of Linux ELILO for Itanium have an incorrect magic value
1099     //       in the PE/COFF Header. If the MachineType is Itanium(IA64) and the
1100     //       Magic value in the OptionalHeader is EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC
1101     //       then override the magic value to EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC
1102     //
1103     Magic = EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;
1104   } else {
1105     //
1106     // Get the magic value from the PE/COFF Optional Header
1107     //
1108     Magic = Hdr.Pe32->OptionalHeader.Magic;
1109   }
1110   if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
1111     SectionAlignment  = Hdr.Pe32->OptionalHeader.SectionAlignment;
1112   } else {
1113     SectionAlignment  = Hdr.Pe32Plus->OptionalHeader.SectionAlignment;
1114   }
1115
1116   SetMemoryAttributesTableSectionAlignment (SectionAlignment);
1117   if ((SectionAlignment & (EFI_ACPI_RUNTIME_PAGE_ALLOCATION_ALIGNMENT - 1)) != 0) {
1118     DEBUG ((DEBUG_ERROR, "SMM !!!!!!!!  InsertImageRecord - Section Alignment(0x%x) is not %dK  !!!!!!!!\n",
1119       SectionAlignment, EFI_ACPI_RUNTIME_PAGE_ALLOCATION_ALIGNMENT >> 10));
1120     PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);
1121     if (PdbPointer != NULL) {
1122       DEBUG ((DEBUG_ERROR, "SMM !!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));
1123     }
1124     goto Finish;
1125   }
1126
1127   Section = (EFI_IMAGE_SECTION_HEADER *) (
1128                (UINT8 *) (UINTN) ImageAddress +
1129                PeCoffHeaderOffset +
1130                sizeof(UINT32) +
1131                sizeof(EFI_IMAGE_FILE_HEADER) +
1132                Hdr.Pe32->FileHeader.SizeOfOptionalHeader
1133                );
1134   ImageRecord->CodeSegmentCount = 0;
1135   InitializeListHead (&ImageRecord->CodeSegmentList);
1136   for (Index = 0; Index < Hdr.Pe32->FileHeader.NumberOfSections; Index++) {
1137     Name = Section[Index].Name;
1138     DEBUG ((
1139       DEBUG_VERBOSE,
1140       "SMM   Section - '%c%c%c%c%c%c%c%c'\n",
1141       Name[0],
1142       Name[1],
1143       Name[2],
1144       Name[3],
1145       Name[4],
1146       Name[5],
1147       Name[6],
1148       Name[7]
1149       ));
1150
1151     if ((Section[Index].Characteristics & EFI_IMAGE_SCN_CNT_CODE) != 0) {
1152       DEBUG ((DEBUG_VERBOSE, "SMM   VirtualSize          - 0x%08x\n", Section[Index].Misc.VirtualSize));
1153       DEBUG ((DEBUG_VERBOSE, "SMM   VirtualAddress       - 0x%08x\n", Section[Index].VirtualAddress));
1154       DEBUG ((DEBUG_VERBOSE, "SMM   SizeOfRawData        - 0x%08x\n", Section[Index].SizeOfRawData));
1155       DEBUG ((DEBUG_VERBOSE, "SMM   PointerToRawData     - 0x%08x\n", Section[Index].PointerToRawData));
1156       DEBUG ((DEBUG_VERBOSE, "SMM   PointerToRelocations - 0x%08x\n", Section[Index].PointerToRelocations));
1157       DEBUG ((DEBUG_VERBOSE, "SMM   PointerToLinenumbers - 0x%08x\n", Section[Index].PointerToLinenumbers));
1158       DEBUG ((DEBUG_VERBOSE, "SMM   NumberOfRelocations  - 0x%08x\n", Section[Index].NumberOfRelocations));
1159       DEBUG ((DEBUG_VERBOSE, "SMM   NumberOfLinenumbers  - 0x%08x\n", Section[Index].NumberOfLinenumbers));
1160       DEBUG ((DEBUG_VERBOSE, "SMM   Characteristics      - 0x%08x\n", Section[Index].Characteristics));
1161
1162       //
1163       // Step 2: record code section
1164       //
1165       ImageRecordCodeSection = AllocatePool (sizeof(*ImageRecordCodeSection));
1166       if (ImageRecordCodeSection == NULL) {
1167         return ;
1168       }
1169       ImageRecordCodeSection->Signature = IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE;
1170
1171       ImageRecordCodeSection->CodeSegmentBase = (UINTN)ImageAddress + Section[Index].VirtualAddress;
1172       ImageRecordCodeSection->CodeSegmentSize = Section[Index].SizeOfRawData;
1173
1174       DEBUG ((DEBUG_VERBOSE, "SMM ImageCode: 0x%016lx - 0x%016lx\n", ImageRecordCodeSection->CodeSegmentBase, ImageRecordCodeSection->CodeSegmentSize));
1175
1176       InsertTailList (&ImageRecord->CodeSegmentList, &ImageRecordCodeSection->Link);
1177       ImageRecord->CodeSegmentCount++;
1178     }
1179   }
1180
1181   if (ImageRecord->CodeSegmentCount == 0) {
1182     SetMemoryAttributesTableSectionAlignment (1);
1183     DEBUG ((DEBUG_ERROR, "SMM !!!!!!!!  InsertImageRecord - CodeSegmentCount is 0  !!!!!!!!\n"));
1184     PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);
1185     if (PdbPointer != NULL) {
1186       DEBUG ((DEBUG_ERROR, "SMM !!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));
1187     }
1188     goto Finish;
1189   }
1190
1191   //
1192   // Final
1193   //
1194   SortImageRecordCodeSection (ImageRecord);
1195   //
1196   // Check overlap all section in ImageBase/Size
1197   //
1198   if (!IsImageRecordCodeSectionValid (ImageRecord)) {
1199     DEBUG ((DEBUG_ERROR, "SMM IsImageRecordCodeSectionValid - FAIL\n"));
1200     goto Finish;
1201   }
1202
1203   InsertTailList (&mImagePropertiesPrivateData.ImageRecordList, &ImageRecord->Link);
1204   mImagePropertiesPrivateData.ImageRecordCount++;
1205
1206   SortImageRecord ();
1207
1208   if (mImagePropertiesPrivateData.CodeSegmentCountMax < ImageRecord->CodeSegmentCount) {
1209     mImagePropertiesPrivateData.CodeSegmentCountMax = ImageRecord->CodeSegmentCount;
1210   }
1211
1212 Finish:
1213   return ;
1214 }
1215
1216 /**
1217   Find image record accroding to image base and size.
1218
1219   @param[in]  ImageBase    Base of PE image
1220   @param[in]  ImageSize    Size of PE image
1221
1222   @return image record
1223 **/
1224 STATIC
1225 IMAGE_PROPERTIES_RECORD *
1226 FindImageRecord (
1227   IN EFI_PHYSICAL_ADDRESS  ImageBase,
1228   IN UINT64                ImageSize
1229   )
1230 {
1231   IMAGE_PROPERTIES_RECORD    *ImageRecord;
1232   LIST_ENTRY                 *ImageRecordLink;
1233   LIST_ENTRY                 *ImageRecordList;
1234
1235   ImageRecordList = &mImagePropertiesPrivateData.ImageRecordList;
1236
1237   for (ImageRecordLink = ImageRecordList->ForwardLink;
1238        ImageRecordLink != ImageRecordList;
1239        ImageRecordLink = ImageRecordLink->ForwardLink) {
1240     ImageRecord = CR (
1241                     ImageRecordLink,
1242                     IMAGE_PROPERTIES_RECORD,
1243                     Link,
1244                     IMAGE_PROPERTIES_RECORD_SIGNATURE
1245                     );
1246
1247     if ((ImageBase == ImageRecord->ImageBase) &&
1248         (ImageSize == ImageRecord->ImageSize)) {
1249       return ImageRecord;
1250     }
1251   }
1252
1253   return NULL;
1254 }
1255
1256 /**
1257   Remove Image record.
1258
1259   @param[in]  DriverEntry    Driver information
1260 **/
1261 VOID
1262 SmmRemoveImageRecord (
1263   IN EFI_SMM_DRIVER_ENTRY  *DriverEntry
1264   )
1265 {
1266   IMAGE_PROPERTIES_RECORD              *ImageRecord;
1267   LIST_ENTRY                           *CodeSegmentListHead;
1268   IMAGE_PROPERTIES_RECORD_CODE_SECTION *ImageRecordCodeSection;
1269
1270   DEBUG ((DEBUG_VERBOSE, "SMM RemoveImageRecord - 0x%x\n", DriverEntry));
1271   DEBUG ((DEBUG_VERBOSE, "SMM RemoveImageRecord - 0x%016lx - 0x%016lx\n", DriverEntry->ImageBuffer, DriverEntry->NumberOfPage));
1272
1273   ImageRecord = FindImageRecord (DriverEntry->ImageBuffer, EFI_PAGES_TO_SIZE(DriverEntry->NumberOfPage));
1274   if (ImageRecord == NULL) {
1275     DEBUG ((DEBUG_ERROR, "SMM !!!!!!!! ImageRecord not found !!!!!!!!\n"));
1276     return ;
1277   }
1278
1279   CodeSegmentListHead = &ImageRecord->CodeSegmentList;
1280   while (!IsListEmpty (CodeSegmentListHead)) {
1281     ImageRecordCodeSection = CR (
1282                                CodeSegmentListHead->ForwardLink,
1283                                IMAGE_PROPERTIES_RECORD_CODE_SECTION,
1284                                Link,
1285                                IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
1286                                );
1287     RemoveEntryList (&ImageRecordCodeSection->Link);
1288     FreePool (ImageRecordCodeSection);
1289   }
1290
1291   RemoveEntryList (&ImageRecord->Link);
1292   FreePool (ImageRecord);
1293   mImagePropertiesPrivateData.ImageRecordCount--;
1294 }
1295
1296 /**
1297   Publish MemoryAttributesTable to SMM configuration table.
1298 **/
1299 VOID
1300 PublishMemoryAttributesTable (
1301   VOID
1302   )
1303 {
1304   UINTN                                MemoryMapSize;
1305   EFI_MEMORY_DESCRIPTOR                *MemoryMap;
1306   UINTN                                MapKey;
1307   UINTN                                DescriptorSize;
1308   UINT32                               DescriptorVersion;
1309   UINTN                                Index;
1310   EFI_STATUS                           Status;
1311   UINTN                                RuntimeEntryCount;
1312   EDKII_PI_SMM_MEMORY_ATTRIBUTES_TABLE *MemoryAttributesTable;
1313   EFI_MEMORY_DESCRIPTOR                *MemoryAttributesEntry;
1314   UINTN                                MemoryAttributesTableSize;
1315
1316   MemoryMapSize = 0;
1317   MemoryMap = NULL;
1318   Status = SmmCoreGetMemoryMapMemoryAttributesTable (
1319              &MemoryMapSize,
1320              MemoryMap,
1321              &MapKey,
1322              &DescriptorSize,
1323              &DescriptorVersion
1324              );
1325   ASSERT (Status == EFI_BUFFER_TOO_SMALL);
1326
1327   do {
1328     DEBUG ((DEBUG_INFO, "MemoryMapSize - 0x%x\n", MemoryMapSize));
1329     MemoryMap = AllocatePool (MemoryMapSize);
1330     ASSERT (MemoryMap != NULL);
1331     DEBUG ((DEBUG_INFO, "MemoryMap - 0x%x\n", MemoryMap));
1332
1333     Status = SmmCoreGetMemoryMapMemoryAttributesTable (
1334                &MemoryMapSize,
1335                MemoryMap,
1336                &MapKey,
1337                &DescriptorSize,
1338                &DescriptorVersion
1339                );
1340     if (EFI_ERROR (Status)) {
1341       FreePool (MemoryMap);
1342     }
1343   } while (Status == EFI_BUFFER_TOO_SMALL);
1344
1345   //
1346   // Allocate MemoryAttributesTable
1347   //
1348   RuntimeEntryCount = MemoryMapSize/DescriptorSize;
1349   MemoryAttributesTableSize = sizeof(EDKII_PI_SMM_MEMORY_ATTRIBUTES_TABLE) + DescriptorSize * RuntimeEntryCount;
1350   MemoryAttributesTable = AllocatePool (sizeof(EDKII_PI_SMM_MEMORY_ATTRIBUTES_TABLE) + DescriptorSize * RuntimeEntryCount);
1351   ASSERT (MemoryAttributesTable != NULL);
1352   MemoryAttributesTable->Version         = EDKII_PI_SMM_MEMORY_ATTRIBUTES_TABLE_VERSION;
1353   MemoryAttributesTable->NumberOfEntries = (UINT32)RuntimeEntryCount;
1354   MemoryAttributesTable->DescriptorSize  = (UINT32)DescriptorSize;
1355   MemoryAttributesTable->Reserved        = 0;
1356   DEBUG ((DEBUG_INFO, "MemoryAttributesTable:\n"));
1357   DEBUG ((DEBUG_INFO, "  Version              - 0x%08x\n", MemoryAttributesTable->Version));
1358   DEBUG ((DEBUG_INFO, "  NumberOfEntries      - 0x%08x\n", MemoryAttributesTable->NumberOfEntries));
1359   DEBUG ((DEBUG_INFO, "  DescriptorSize       - 0x%08x\n", MemoryAttributesTable->DescriptorSize));
1360   MemoryAttributesEntry = (EFI_MEMORY_DESCRIPTOR *)(MemoryAttributesTable + 1);
1361   for (Index = 0; Index < MemoryMapSize/DescriptorSize; Index++) {
1362     CopyMem (MemoryAttributesEntry, MemoryMap, DescriptorSize);
1363     DEBUG ((DEBUG_INFO, "Entry (0x%x)\n", MemoryAttributesEntry));
1364     DEBUG ((DEBUG_INFO, "  Type              - 0x%x\n", MemoryAttributesEntry->Type));
1365     DEBUG ((DEBUG_INFO, "  PhysicalStart     - 0x%016lx\n", MemoryAttributesEntry->PhysicalStart));
1366     DEBUG ((DEBUG_INFO, "  VirtualStart      - 0x%016lx\n", MemoryAttributesEntry->VirtualStart));
1367     DEBUG ((DEBUG_INFO, "  NumberOfPages     - 0x%016lx\n", MemoryAttributesEntry->NumberOfPages));
1368     DEBUG ((DEBUG_INFO, "  Attribute         - 0x%016lx\n", MemoryAttributesEntry->Attribute));
1369     MemoryAttributesEntry = NEXT_MEMORY_DESCRIPTOR(MemoryAttributesEntry, DescriptorSize);
1370
1371     MemoryMap = NEXT_MEMORY_DESCRIPTOR(MemoryMap, DescriptorSize);
1372   }
1373
1374   Status = gSmst->SmmInstallConfigurationTable (gSmst, &gEdkiiPiSmmMemoryAttributesTableGuid, MemoryAttributesTable, MemoryAttributesTableSize);
1375   ASSERT_EFI_ERROR (Status);
1376 }
1377
1378 /**
1379   This function returns if image is inside SMRAM.
1380
1381   @param[in] LoadedImage LoadedImage protocol instance for an image.
1382
1383   @retval TRUE  the image is inside SMRAM.
1384   @retval FALSE the image is outside SMRAM.
1385 **/
1386 BOOLEAN
1387 IsImageInsideSmram (
1388   IN EFI_LOADED_IMAGE_PROTOCOL   *LoadedImage
1389   )
1390 {
1391   UINTN  Index;
1392
1393   for (Index = 0; Index < mFullSmramRangeCount; Index++) {
1394     if ((mFullSmramRanges[Index].PhysicalStart <= (UINTN)LoadedImage->ImageBase)&&
1395         (mFullSmramRanges[Index].PhysicalStart + mFullSmramRanges[Index].PhysicalSize >= (UINTN)LoadedImage->ImageBase + LoadedImage->ImageSize)) {
1396       return TRUE;
1397     }
1398   }
1399
1400   return FALSE;
1401 }
1402
1403 /**
1404   This function installs all SMM image record information.
1405 **/
1406 VOID
1407 SmmInstallImageRecord (
1408   VOID
1409   )
1410 {
1411   EFI_STATUS                  Status;
1412   UINTN                       NoHandles;
1413   EFI_HANDLE                  *HandleBuffer;
1414   EFI_LOADED_IMAGE_PROTOCOL   *LoadedImage;
1415   UINTN                       Index;
1416   EFI_SMM_DRIVER_ENTRY        DriverEntry;
1417
1418   Status = SmmLocateHandleBuffer (
1419              ByProtocol,
1420              &gEfiLoadedImageProtocolGuid,
1421              NULL,
1422              &NoHandles,
1423              &HandleBuffer
1424              );
1425   if (EFI_ERROR (Status)) {
1426     return ;
1427   }
1428
1429   for (Index = 0; Index < NoHandles; Index++) {
1430     Status = gSmst->SmmHandleProtocol (
1431                       HandleBuffer[Index],
1432                       &gEfiLoadedImageProtocolGuid,
1433                       (VOID **)&LoadedImage
1434                       );
1435     if (EFI_ERROR (Status)) {
1436       continue;
1437     }
1438     DEBUG ((DEBUG_VERBOSE, "LoadedImage - 0x%x 0x%x ", LoadedImage->ImageBase, LoadedImage->ImageSize));
1439     {
1440       VOID *PdbPointer;
1441       PdbPointer = PeCoffLoaderGetPdbPointer (LoadedImage->ImageBase);
1442       if (PdbPointer != NULL) {
1443         DEBUG ((DEBUG_VERBOSE, "(%a) ", PdbPointer));
1444       }
1445     }
1446     DEBUG ((DEBUG_VERBOSE, "\n"));
1447     ZeroMem (&DriverEntry, sizeof(DriverEntry));
1448     DriverEntry.ImageBuffer  = (UINTN)LoadedImage->ImageBase;
1449     DriverEntry.NumberOfPage = EFI_SIZE_TO_PAGES((UINTN)LoadedImage->ImageSize);
1450     SmmInsertImageRecord (&DriverEntry);
1451   }
1452
1453   FreePool (HandleBuffer);
1454 }
1455
1456 /**
1457   Install MemoryAttributesTable.
1458
1459   @param[in] Protocol   Points to the protocol's unique identifier.
1460   @param[in] Interface  Points to the interface instance.
1461   @param[in] Handle     The handle on which the interface was installed.
1462
1463   @retval EFI_SUCCESS   Notification runs successfully.
1464 **/
1465 EFI_STATUS
1466 EFIAPI
1467 SmmInstallMemoryAttributesTable (
1468   IN CONST EFI_GUID  *Protocol,
1469   IN VOID            *Interface,
1470   IN EFI_HANDLE      Handle
1471   )
1472 {
1473   SmmInstallImageRecord ();
1474
1475   DEBUG ((DEBUG_INFO, "SMM MemoryProtectionAttribute - 0x%016lx\n", mMemoryProtectionAttribute));
1476   if ((mMemoryProtectionAttribute & EFI_MEMORY_ATTRIBUTES_RUNTIME_MEMORY_PROTECTION_NON_EXECUTABLE_PE_DATA) == 0) {
1477     return EFI_SUCCESS;
1478   }
1479
1480   DEBUG ((DEBUG_VERBOSE, "SMM Total Image Count - 0x%x\n", mImagePropertiesPrivateData.ImageRecordCount));
1481   DEBUG ((DEBUG_VERBOSE, "SMM Dump ImageRecord:\n"));
1482   DumpImageRecord ();
1483
1484   PublishMemoryAttributesTable ();
1485
1486   return EFI_SUCCESS;
1487 }
1488
1489 /**
1490   Initialize MemoryAttributesTable support.
1491 **/
1492 VOID
1493 EFIAPI
1494 SmmCoreInitializeMemoryAttributesTable (
1495   VOID
1496   )
1497 {
1498   EFI_STATUS                        Status;
1499   VOID                              *Registration;
1500
1501   Status = gSmst->SmmRegisterProtocolNotify (
1502                     &gEfiSmmEndOfDxeProtocolGuid,
1503                     SmmInstallMemoryAttributesTable,
1504                     &Registration
1505                     );
1506   ASSERT_EFI_ERROR (Status);
1507
1508   return ;
1509 }