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