MdeModulePkg/Core/Dxe/Misc/MemoryAttributesTable.c

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