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