MdeModulePkg/Core/PiSmmCore/MemoryAttributesTable.c

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