MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c

   1 /** @file
   2   UEFI Memory Protection support.
   3
   4   If the UEFI image is page aligned, the image code section is set to read only
   5   and the image data section is set to non-executable.
   6
   7   1) This policy is applied for all UEFI image including boot service driver,
   8      runtime driver or application.
   9   2) This policy is applied only if the UEFI image meets the page alignment
  10      requirement.
  11   3) This policy is applied only if the Source UEFI image matches the
  12      PcdImageProtectionPolicy definition.
  13   4) This policy is not applied to the non-PE image region.
  14
  15   The DxeCore calls CpuArchProtocol->SetMemoryAttributes() to protect
  16   the image. If the CpuArch protocol is not installed yet, the DxeCore
  17   enqueues the protection request. Once the CpuArch is installed, the
  18   DxeCore dequeues the protection request and applies policy.
  19
  20   Once the image is unloaded, the protection is removed automatically.
  21
  22 Copyright (c) 2017, Intel Corporation. All rights reserved.<BR>
  23 This program and the accompanying materials
  24 are licensed and made available under the terms and conditions of the BSD License
  25 which accompanies this distribution.  The full text of the license may be found at
  26 http://opensource.org/licenses/bsd-license.php
  27
  28 THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  29 WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
  30
  31 **/
  32
  33 #include <PiDxe.h>
  34 #include <Library/BaseLib.h>
  35 #include <Library/BaseMemoryLib.h>
  36 #include <Library/MemoryAllocationLib.h>
  37 #include <Library/UefiBootServicesTableLib.h>
  38 #include <Library/DxeServicesTableLib.h>
  39 #include <Library/DebugLib.h>
  40 #include <Library/UefiLib.h>
  41
  42 #include <Guid/EventGroup.h>
  43 #include <Guid/MemoryAttributesTable.h>
  44 #include <Guid/PropertiesTable.h>
  45
  46 #include <Protocol/FirmwareVolume2.h>
  47 #include <Protocol/BlockIo.h>
  48 #include <Protocol/SimpleFileSystem.h>
  49
  50 #include "DxeMain.h"
  51
  52 #define CACHE_ATTRIBUTE_MASK   (EFI_MEMORY_UC | EFI_MEMORY_WC | EFI_MEMORY_WT | EFI_MEMORY_WB | EFI_MEMORY_UCE | EFI_MEMORY_WP)
  53 #define MEMORY_ATTRIBUTE_MASK  (EFI_MEMORY_RP | EFI_MEMORY_XP | EFI_MEMORY_RO)
  54
  55 //
  56 // Image type definitions
  57 //
  58 #define IMAGE_UNKNOWN                         0x00000001
  59 #define IMAGE_FROM_FV                         0x00000002
  60
  61 //
  62 // Protection policy bit definition
  63 //
  64 #define DO_NOT_PROTECT                         0x00000000
  65 #define PROTECT_IF_ALIGNED_ELSE_ALLOW          0x00000001
  66
  67 UINT32   mImageProtectionPolicy;
  68
  69 /**
  70   Sort code section in image record, based upon CodeSegmentBase from low to high.
  71
  72   @param  ImageRecord    image record to be sorted
  73 **/
  74 VOID
  75 SortImageRecordCodeSection (
  76   IN IMAGE_PROPERTIES_RECORD              *ImageRecord
  77   );
  78
  79 /**
  80   Check if code section in image record is valid.
  81
  82   @param  ImageRecord    image record to be checked
  83
  84   @retval TRUE  image record is valid
  85   @retval FALSE image record is invalid
  86 **/
  87 BOOLEAN
  88 IsImageRecordCodeSectionValid (
  89   IN IMAGE_PROPERTIES_RECORD              *ImageRecord
  90   );
  91
  92 /**
  93   Get the image type.
  94
  95   @param[in]    File       This is a pointer to the device path of the file that is
  96                            being dispatched.
  97
  98   @return UINT32           Image Type
  99 **/
 100 UINT32
 101 GetImageType (
 102   IN  CONST EFI_DEVICE_PATH_PROTOCOL   *File
 103   )
 104 {
 105   EFI_STATUS                        Status;
 106   EFI_HANDLE                        DeviceHandle;
 107   EFI_DEVICE_PATH_PROTOCOL          *TempDevicePath;
 108
 109   if (File == NULL) {
 110     return IMAGE_UNKNOWN;
 111   }
 112
 113   //
 114   // First check to see if File is from a Firmware Volume
 115   //
 116   DeviceHandle      = NULL;
 117   TempDevicePath = (EFI_DEVICE_PATH_PROTOCOL *) File;
 118   Status = gBS->LocateDevicePath (
 119                   &gEfiFirmwareVolume2ProtocolGuid,
 120                   &TempDevicePath,
 121                   &DeviceHandle
 122                   );
 123   if (!EFI_ERROR (Status)) {
 124     Status = gBS->OpenProtocol (
 125                     DeviceHandle,
 126                     &gEfiFirmwareVolume2ProtocolGuid,
 127                     NULL,
 128                     NULL,
 129                     NULL,
 130                     EFI_OPEN_PROTOCOL_TEST_PROTOCOL
 131                     );
 132     if (!EFI_ERROR (Status)) {
 133       return IMAGE_FROM_FV;
 134     }
 135   }
 136   return IMAGE_UNKNOWN;
 137 }
 138
 139 /**
 140   Get UEFI image protection policy based upon image type.
 141
 142   @param[in]  ImageType    The UEFI image type
 143
 144   @return UEFI image protection policy
 145 **/
 146 UINT32
 147 GetProtectionPolicyFromImageType (
 148   IN UINT32  ImageType
 149   )
 150 {
 151   if ((ImageType & mImageProtectionPolicy) == 0) {
 152     return DO_NOT_PROTECT;
 153   } else {
 154     return PROTECT_IF_ALIGNED_ELSE_ALLOW;
 155   }
 156 }
 157
 158 /**
 159   Get UEFI image protection policy based upon loaded image device path.
 160
 161   @param[in]  LoadedImage              The loaded image protocol
 162   @param[in]  LoadedImageDevicePath    The loaded image device path protocol
 163
 164   @return UEFI image protection policy
 165 **/
 166 UINT32
 167 GetUefiImageProtectionPolicy (
 168   IN EFI_LOADED_IMAGE_PROTOCOL   *LoadedImage,
 169   IN EFI_DEVICE_PATH_PROTOCOL    *LoadedImageDevicePath
 170   )
 171 {
 172   BOOLEAN     InSmm;
 173   UINT32      ImageType;
 174   UINT32      ProtectionPolicy;
 175
 176   //
 177   // Check SMM
 178   //
 179   InSmm = FALSE;
 180   if (gSmmBase2 != NULL) {
 181     gSmmBase2->InSmm (gSmmBase2, &InSmm);
 182   }
 183   if (InSmm) {
 184     return FALSE;
 185   }
 186
 187   //
 188   // Check DevicePath
 189   //
 190   if (LoadedImage == gDxeCoreLoadedImage) {
 191     ImageType = IMAGE_FROM_FV;
 192   } else {
 193     ImageType = GetImageType (LoadedImageDevicePath);
 194   }
 195   ProtectionPolicy = GetProtectionPolicyFromImageType (ImageType);
 196   return ProtectionPolicy;
 197 }
 198
 199
 200 /**
 201   Set UEFI image memory attributes.
 202
 203   @param[in]  BaseAddress            Specified start address
 204   @param[in]  Length                 Specified length
 205   @param[in]  Attributes             Specified attributes
 206 **/
 207 VOID
 208 SetUefiImageMemoryAttributes (
 209   IN UINT64   BaseAddress,
 210   IN UINT64   Length,
 211   IN UINT64   Attributes
 212   )
 213 {
 214   EFI_STATUS                       Status;
 215   EFI_GCD_MEMORY_SPACE_DESCRIPTOR  Descriptor;
 216   UINT64                           FinalAttributes;
 217
 218   Status = CoreGetMemorySpaceDescriptor(BaseAddress, &Descriptor);
 219   ASSERT_EFI_ERROR(Status);
 220
 221   FinalAttributes = (Descriptor.Attributes & CACHE_ATTRIBUTE_MASK) | (Attributes & MEMORY_ATTRIBUTE_MASK);
 222
 223   DEBUG ((DEBUG_INFO, "SetUefiImageMemoryAttributes - 0x%016lx - 0x%016lx (0x%016lx)\n", BaseAddress, Length, FinalAttributes));
 224
 225   ASSERT(gCpu != NULL);
 226   gCpu->SetMemoryAttributes (gCpu, BaseAddress, Length, FinalAttributes);
 227 }
 228
 229 /**
 230   Set UEFI image protection attributes.
 231
 232   @param[in]  ImageRecord    A UEFI image record
 233   @param[in]  Protect        TRUE:  Protect the UEFI image.
 234                              FALSE: Unprotect the UEFI image.
 235 **/
 236 VOID
 237 SetUefiImageProtectionAttributes (
 238   IN IMAGE_PROPERTIES_RECORD     *ImageRecord,
 239   IN BOOLEAN                     Protect
 240   )
 241 {
 242   IMAGE_PROPERTIES_RECORD_CODE_SECTION      *ImageRecordCodeSection;
 243   LIST_ENTRY                                *ImageRecordCodeSectionLink;
 244   LIST_ENTRY                                *ImageRecordCodeSectionEndLink;
 245   LIST_ENTRY                                *ImageRecordCodeSectionList;
 246   UINT64                                    CurrentBase;
 247   UINT64                                    ImageEnd;
 248   UINT64                                    Attribute;
 249
 250   ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;
 251
 252   CurrentBase = ImageRecord->ImageBase;
 253   ImageEnd    = ImageRecord->ImageBase + ImageRecord->ImageSize;
 254
 255   ImageRecordCodeSectionLink = ImageRecordCodeSectionList->ForwardLink;
 256   ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;
 257   while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {
 258     ImageRecordCodeSection = CR (
 259                                ImageRecordCodeSectionLink,
 260                                IMAGE_PROPERTIES_RECORD_CODE_SECTION,
 261                                Link,
 262                                IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
 263                                );
 264     ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;
 265
 266     ASSERT (CurrentBase <= ImageRecordCodeSection->CodeSegmentBase);
 267     if (CurrentBase < ImageRecordCodeSection->CodeSegmentBase) {
 268       //
 269       // DATA
 270       //
 271       if (Protect) {
 272         Attribute = EFI_MEMORY_XP;
 273       } else {
 274         Attribute = 0;
 275       }
 276       SetUefiImageMemoryAttributes (
 277         CurrentBase,
 278         ImageRecordCodeSection->CodeSegmentBase - CurrentBase,
 279         Attribute
 280         );
 281     }
 282     //
 283     // CODE
 284     //
 285     if (Protect) {
 286       Attribute = EFI_MEMORY_RO;
 287     } else {
 288       Attribute = 0;
 289     }
 290     SetUefiImageMemoryAttributes (
 291       ImageRecordCodeSection->CodeSegmentBase,
 292       ImageRecordCodeSection->CodeSegmentSize,
 293       Attribute
 294       );
 295     CurrentBase = ImageRecordCodeSection->CodeSegmentBase + ImageRecordCodeSection->CodeSegmentSize;
 296   }
 297   //
 298   // Last DATA
 299   //
 300   ASSERT (CurrentBase <= ImageEnd);
 301   if (CurrentBase < ImageEnd) {
 302     //
 303     // DATA
 304     //
 305     if (Protect) {
 306       Attribute = EFI_MEMORY_XP;
 307     } else {
 308       Attribute = 0;
 309     }
 310     SetUefiImageMemoryAttributes (
 311       CurrentBase,
 312       ImageEnd - CurrentBase,
 313       Attribute
 314       );
 315   }
 316   return ;
 317 }
 318
 319 /**
 320   Return if the PE image section is aligned.
 321
 322   @param[in]  SectionAlignment    PE/COFF section alignment
 323   @param[in]  MemoryType          PE/COFF image memory type
 324
 325   @retval TRUE  The PE image section is aligned.
 326   @retval FALSE The PE image section is not aligned.
 327 **/
 328 BOOLEAN
 329 IsMemoryProtectionSectionAligned (
 330   IN UINT32           SectionAlignment,
 331   IN EFI_MEMORY_TYPE  MemoryType
 332   )
 333 {
 334   UINT32  PageAlignment;
 335
 336   switch (MemoryType) {
 337   case EfiRuntimeServicesCode:
 338   case EfiACPIMemoryNVS:
 339     PageAlignment = EFI_ACPI_RUNTIME_PAGE_ALLOCATION_ALIGNMENT;
 340     break;
 341   case EfiRuntimeServicesData:
 342   case EfiACPIReclaimMemory:
 343     ASSERT (FALSE);
 344     PageAlignment = EFI_ACPI_RUNTIME_PAGE_ALLOCATION_ALIGNMENT;
 345     break;
 346   case EfiBootServicesCode:
 347   case EfiLoaderCode:
 348   case EfiReservedMemoryType:
 349     PageAlignment = EFI_PAGE_SIZE;
 350     break;
 351   default:
 352     ASSERT (FALSE);
 353     PageAlignment = EFI_PAGE_SIZE;
 354     break;
 355   }
 356
 357   if ((SectionAlignment & (PageAlignment - 1)) != 0) {
 358     return FALSE;
 359   } else {
 360     return TRUE;
 361   }
 362 }
 363
 364 /**
 365   Free Image record.
 366
 367   @param[in]  ImageRecord    A UEFI image record
 368 **/
 369 VOID
 370 FreeImageRecord (
 371   IN IMAGE_PROPERTIES_RECORD              *ImageRecord
 372   )
 373 {
 374   LIST_ENTRY                           *CodeSegmentListHead;
 375   IMAGE_PROPERTIES_RECORD_CODE_SECTION *ImageRecordCodeSection;
 376
 377   CodeSegmentListHead = &ImageRecord->CodeSegmentList;
 378   while (!IsListEmpty (CodeSegmentListHead)) {
 379     ImageRecordCodeSection = CR (
 380                                CodeSegmentListHead->ForwardLink,
 381                                IMAGE_PROPERTIES_RECORD_CODE_SECTION,
 382                                Link,
 383                                IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE
 384                                );
 385     RemoveEntryList (&ImageRecordCodeSection->Link);
 386     FreePool (ImageRecordCodeSection);
 387   }
 388
 389   if (ImageRecord->Link.ForwardLink != NULL) {
 390     RemoveEntryList (&ImageRecord->Link);
 391   }
 392   FreePool (ImageRecord);
 393 }
 394
 395 /**
 396   Protect or unprotect UEFI image common function.
 397
 398   @param[in]  LoadedImage              The loaded image protocol
 399   @param[in]  LoadedImageDevicePath    The loaded image device path protocol
 400   @param[in]  Protect                  TRUE:  Protect the UEFI image.
 401                                        FALSE: Unprotect the UEFI image.
 402 **/
 403 VOID
 404 ProtectUefiImageCommon (
 405   IN EFI_LOADED_IMAGE_PROTOCOL   *LoadedImage,
 406   IN EFI_DEVICE_PATH_PROTOCOL    *LoadedImageDevicePath,
 407   IN BOOLEAN                     Protect
 408   )
 409 {
 410   VOID                                 *ImageAddress;
 411   EFI_IMAGE_DOS_HEADER                 *DosHdr;
 412   UINT32                               PeCoffHeaderOffset;
 413   UINT32                               SectionAlignment;
 414   EFI_IMAGE_SECTION_HEADER             *Section;
 415   EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
 416   UINT8                                *Name;
 417   UINTN                                Index;
 418   IMAGE_PROPERTIES_RECORD              *ImageRecord;
 419   CHAR8                                *PdbPointer;
 420   IMAGE_PROPERTIES_RECORD_CODE_SECTION *ImageRecordCodeSection;
 421   UINT16                               Magic;
 422   BOOLEAN                              IsAligned;
 423   UINT32                               ProtectionPolicy;
 424
 425   DEBUG ((DEBUG_INFO, "ProtectUefiImageCommon - 0x%x\n", LoadedImage));
 426   DEBUG ((DEBUG_INFO, "  - 0x%016lx - 0x%016lx\n", (EFI_PHYSICAL_ADDRESS)(UINTN)LoadedImage->ImageBase, LoadedImage->ImageSize));
 427
 428   if (gCpu == NULL) {
 429     return ;
 430   }
 431
 432   ProtectionPolicy = GetUefiImageProtectionPolicy (LoadedImage, LoadedImageDevicePath);
 433   switch (ProtectionPolicy) {
 434   case DO_NOT_PROTECT:
 435     return ;
 436   case PROTECT_IF_ALIGNED_ELSE_ALLOW:
 437     break;
 438   default:
 439     ASSERT(FALSE);
 440     return ;
 441   }
 442
 443   ImageRecord = AllocateZeroPool (sizeof(*ImageRecord));
 444   if (ImageRecord == NULL) {
 445     return ;
 446   }
 447   ImageRecord->Signature = IMAGE_PROPERTIES_RECORD_SIGNATURE;
 448
 449   //
 450   // Step 1: record whole region
 451   //
 452   ImageRecord->ImageBase = (EFI_PHYSICAL_ADDRESS)(UINTN)LoadedImage->ImageBase;
 453   ImageRecord->ImageSize = LoadedImage->ImageSize;
 454
 455   ImageAddress = LoadedImage->ImageBase;
 456
 457   PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);
 458   if (PdbPointer != NULL) {
 459     DEBUG ((DEBUG_VERBOSE, "  Image - %a\n", PdbPointer));
 460   }
 461
 462   //
 463   // Check PE/COFF image
 464   //
 465   DosHdr = (EFI_IMAGE_DOS_HEADER *) (UINTN) ImageAddress;
 466   PeCoffHeaderOffset = 0;
 467   if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
 468     PeCoffHeaderOffset = DosHdr->e_lfanew;
 469   }
 470
 471   Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINT8 *) (UINTN) ImageAddress + PeCoffHeaderOffset);
 472   if (Hdr.Pe32->Signature != EFI_IMAGE_NT_SIGNATURE) {
 473     DEBUG ((DEBUG_VERBOSE, "Hdr.Pe32->Signature invalid - 0x%x\n", Hdr.Pe32->Signature));
 474     // It might be image in SMM.
 475     goto Finish;
 476   }
 477
 478   //
 479   // Get SectionAlignment
 480   //
 481   if (Hdr.Pe32->FileHeader.Machine == IMAGE_FILE_MACHINE_IA64 && Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
 482     //
 483     // NOTE: Some versions of Linux ELILO for Itanium have an incorrect magic value
 484     //       in the PE/COFF Header. If the MachineType is Itanium(IA64) and the
 485     //       Magic value in the OptionalHeader is EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC
 486     //       then override the magic value to EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC
 487     //
 488     Magic = EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;
 489   } else {
 490     //
 491     // Get the magic value from the PE/COFF Optional Header
 492     //
 493     Magic = Hdr.Pe32->OptionalHeader.Magic;
 494   }
 495   if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
 496     SectionAlignment  = Hdr.Pe32->OptionalHeader.SectionAlignment;
 497   } else {
 498     SectionAlignment  = Hdr.Pe32Plus->OptionalHeader.SectionAlignment;
 499   }
 500
 501   IsAligned = IsMemoryProtectionSectionAligned (SectionAlignment, LoadedImage->ImageCodeType);
 502   if (!IsAligned) {
 503     DEBUG ((DEBUG_VERBOSE, "!!!!!!!!  ProtectUefiImageCommon - Section Alignment(0x%x) is incorrect  !!!!!!!!\n",
 504       SectionAlignment));
 505     PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);
 506     if (PdbPointer != NULL) {
 507       DEBUG ((DEBUG_VERBOSE, "!!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));
 508     }
 509     goto Finish;
 510   }
 511
 512   Section = (EFI_IMAGE_SECTION_HEADER *) (
 513                (UINT8 *) (UINTN) ImageAddress +
 514                PeCoffHeaderOffset +
 515                sizeof(UINT32) +
 516                sizeof(EFI_IMAGE_FILE_HEADER) +
 517                Hdr.Pe32->FileHeader.SizeOfOptionalHeader
 518                );
 519   ImageRecord->CodeSegmentCount = 0;
 520   InitializeListHead (&ImageRecord->CodeSegmentList);
 521   for (Index = 0; Index < Hdr.Pe32->FileHeader.NumberOfSections; Index++) {
 522     Name = Section[Index].Name;
 523     DEBUG ((
 524       DEBUG_VERBOSE,
 525       "  Section - '%c%c%c%c%c%c%c%c'\n",
 526       Name[0],
 527       Name[1],
 528       Name[2],
 529       Name[3],
 530       Name[4],
 531       Name[5],
 532       Name[6],
 533       Name[7]
 534       ));
 535
 536     //
 537     // Instead of assuming that a PE/COFF section of type EFI_IMAGE_SCN_CNT_CODE
 538     // can always be mapped read-only, classify a section as a code section only
 539     // if it has the executable attribute set and the writable attribute cleared.
 540     //
 541     // This adheres more closely to the PE/COFF spec, and avoids issues with
 542     // Linux OS loaders that may consist of a single read/write/execute section.
 543     //
 544     if ((Section[Index].Characteristics & (EFI_IMAGE_SCN_MEM_WRITE | EFI_IMAGE_SCN_MEM_EXECUTE)) == EFI_IMAGE_SCN_MEM_EXECUTE) {
 545       DEBUG ((DEBUG_VERBOSE, "  VirtualSize          - 0x%08x\n", Section[Index].Misc.VirtualSize));
 546       DEBUG ((DEBUG_VERBOSE, "  VirtualAddress       - 0x%08x\n", Section[Index].VirtualAddress));
 547       DEBUG ((DEBUG_VERBOSE, "  SizeOfRawData        - 0x%08x\n", Section[Index].SizeOfRawData));
 548       DEBUG ((DEBUG_VERBOSE, "  PointerToRawData     - 0x%08x\n", Section[Index].PointerToRawData));
 549       DEBUG ((DEBUG_VERBOSE, "  PointerToRelocations - 0x%08x\n", Section[Index].PointerToRelocations));
 550       DEBUG ((DEBUG_VERBOSE, "  PointerToLinenumbers - 0x%08x\n", Section[Index].PointerToLinenumbers));
 551       DEBUG ((DEBUG_VERBOSE, "  NumberOfRelocations  - 0x%08x\n", Section[Index].NumberOfRelocations));
 552       DEBUG ((DEBUG_VERBOSE, "  NumberOfLinenumbers  - 0x%08x\n", Section[Index].NumberOfLinenumbers));
 553       DEBUG ((DEBUG_VERBOSE, "  Characteristics      - 0x%08x\n", Section[Index].Characteristics));
 554
 555       //
 556       // Step 2: record code section
 557       //
 558       ImageRecordCodeSection = AllocatePool (sizeof(*ImageRecordCodeSection));
 559       if (ImageRecordCodeSection == NULL) {
 560         return ;
 561       }
 562       ImageRecordCodeSection->Signature = IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE;
 563
 564       ImageRecordCodeSection->CodeSegmentBase = (UINTN)ImageAddress + Section[Index].VirtualAddress;
 565       ImageRecordCodeSection->CodeSegmentSize = ALIGN_VALUE(Section[Index].SizeOfRawData, SectionAlignment);
 566
 567       DEBUG ((DEBUG_VERBOSE, "ImageCode: 0x%016lx - 0x%016lx\n", ImageRecordCodeSection->CodeSegmentBase, ImageRecordCodeSection->CodeSegmentSize));
 568
 569       InsertTailList (&ImageRecord->CodeSegmentList, &ImageRecordCodeSection->Link);
 570       ImageRecord->CodeSegmentCount++;
 571     }
 572   }
 573
 574   if (ImageRecord->CodeSegmentCount == 0) {
 575     DEBUG ((DEBUG_ERROR, "!!!!!!!!  ProtectUefiImageCommon - CodeSegmentCount is 0  !!!!!!!!\n"));
 576     PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);
 577     if (PdbPointer != NULL) {
 578       DEBUG ((DEBUG_ERROR, "!!!!!!!!  Image - %a  !!!!!!!!\n", PdbPointer));
 579     }
 580     goto Finish;
 581   }
 582
 583   //
 584   // Final
 585   //
 586   SortImageRecordCodeSection (ImageRecord);
 587   //
 588   // Check overlap all section in ImageBase/Size
 589   //
 590   if (!IsImageRecordCodeSectionValid (ImageRecord)) {
 591     DEBUG ((DEBUG_ERROR, "IsImageRecordCodeSectionValid - FAIL\n"));
 592     goto Finish;
 593   }
 594
 595   //
 596   // Round up the ImageSize, some CPU arch may return EFI_UNSUPPORTED if ImageSize is not aligned.
 597   // Given that the loader always allocates full pages, we know the space after the image is not used.
 598   //
 599   ImageRecord->ImageSize = ALIGN_VALUE(LoadedImage->ImageSize, EFI_PAGE_SIZE);
 600
 601   //
 602   // CPU ARCH present. Update memory attribute directly.
 603   //
 604   SetUefiImageProtectionAttributes (ImageRecord, Protect);
 605
 606   //
 607   // Clean up
 608   //
 609   FreeImageRecord (ImageRecord);
 610
 611 Finish:
 612   return ;
 613 }
 614
 615 /**
 616   Protect UEFI image.
 617
 618   @param[in]  LoadedImage              The loaded image protocol
 619   @param[in]  LoadedImageDevicePath    The loaded image device path protocol
 620 **/
 621 VOID
 622 ProtectUefiImage (
 623   IN EFI_LOADED_IMAGE_PROTOCOL   *LoadedImage,
 624   IN EFI_DEVICE_PATH_PROTOCOL    *LoadedImageDevicePath
 625   )
 626 {
 627   if (PcdGet32(PcdImageProtectionPolicy) != 0) {
 628     ProtectUefiImageCommon (LoadedImage, LoadedImageDevicePath, TRUE);
 629   }
 630 }
 631
 632 /**
 633   Unprotect UEFI image.
 634
 635   @param[in]  LoadedImage              The loaded image protocol
 636   @param[in]  LoadedImageDevicePath    The loaded image device path protocol
 637 **/
 638 VOID
 639 UnprotectUefiImage (
 640   IN EFI_LOADED_IMAGE_PROTOCOL   *LoadedImage,
 641   IN EFI_DEVICE_PATH_PROTOCOL    *LoadedImageDevicePath
 642   )
 643 {
 644   if (PcdGet32(PcdImageProtectionPolicy) != 0) {
 645     ProtectUefiImageCommon (LoadedImage, LoadedImageDevicePath, FALSE);
 646   }
 647 }
 648
 649 /**
 650   A notification for CPU_ARCH protocol.
 651
 652   @param[in]  Event                 Event whose notification function is being invoked.
 653   @param[in]  Context               Pointer to the notification function's context,
 654                                     which is implementation-dependent.
 655
 656 **/
 657 VOID
 658 EFIAPI
 659 MemoryProtectionCpuArchProtocolNotify (
 660   IN EFI_EVENT                Event,
 661   IN VOID                     *Context
 662   )
 663 {
 664   EFI_STATUS                  Status;
 665   EFI_LOADED_IMAGE_PROTOCOL   *LoadedImage;
 666   EFI_DEVICE_PATH_PROTOCOL    *LoadedImageDevicePath;
 667   UINTN                       NoHandles;
 668   EFI_HANDLE                  *HandleBuffer;
 669   UINTN                       Index;
 670
 671   DEBUG ((DEBUG_INFO, "MemoryProtectionCpuArchProtocolNotify:\n"));
 672   Status = CoreLocateProtocol (&gEfiCpuArchProtocolGuid, NULL, (VOID **)&gCpu);
 673   if (EFI_ERROR (Status)) {
 674     return;
 675   }
 676
 677   Status = gBS->LocateHandleBuffer (
 678                   ByProtocol,
 679                   &gEfiLoadedImageProtocolGuid,
 680                   NULL,
 681                   &NoHandles,
 682                   &HandleBuffer
 683                   );
 684   if (EFI_ERROR (Status) && (NoHandles == 0)) {
 685     return ;
 686   }
 687
 688   for (Index = 0; Index < NoHandles; Index++) {
 689     Status = gBS->HandleProtocol (
 690                     HandleBuffer[Index],
 691                     &gEfiLoadedImageProtocolGuid,
 692                     (VOID **)&LoadedImage
 693                     );
 694     if (EFI_ERROR(Status)) {
 695       continue;
 696     }
 697     Status = gBS->HandleProtocol (
 698                     HandleBuffer[Index],
 699                     &gEfiLoadedImageDevicePathProtocolGuid,
 700                     (VOID **)&LoadedImageDevicePath
 701                     );
 702     if (EFI_ERROR(Status)) {
 703       LoadedImageDevicePath = NULL;
 704     }
 705
 706     ProtectUefiImage (LoadedImage, LoadedImageDevicePath);
 707   }
 708
 709   CoreCloseEvent (Event);
 710   return;
 711 }
 712
 713 /**
 714   ExitBootServices Callback function for memory protection.
 715 **/
 716 VOID
 717 MemoryProtectionExitBootServicesCallback (
 718   VOID
 719   )
 720 {
 721   EFI_RUNTIME_IMAGE_ENTRY       *RuntimeImage;
 722   LIST_ENTRY                    *Link;
 723
 724   //
 725   // We need remove the RT protection, because RT relocation need write code segment
 726   // at SetVirtualAddressMap(). We cannot assume OS/Loader has taken over page table at that time.
 727   //
 728   // Firmware does not own page tables after ExitBootServices(), so the OS would
 729   // have to relax protection of RT code pages across SetVirtualAddressMap(), or
 730   // delay setting protections on RT code pages until after SetVirtualAddressMap().
 731   // OS may set protection on RT based upon EFI_MEMORY_ATTRIBUTES_TABLE later.
 732   //
 733   if (mImageProtectionPolicy != 0) {
 734     for (Link = gRuntime->ImageHead.ForwardLink; Link != &gRuntime->ImageHead; Link = Link->ForwardLink) {
 735       RuntimeImage = BASE_CR (Link, EFI_RUNTIME_IMAGE_ENTRY, Link);
 736       SetUefiImageMemoryAttributes ((UINT64)(UINTN)RuntimeImage->ImageBase, ALIGN_VALUE(RuntimeImage->ImageSize, EFI_PAGE_SIZE), 0);
 737     }
 738   }
 739 }
 740
 741 /**
 742   Initialize Memory Protection support.
 743 **/
 744 VOID
 745 EFIAPI
 746 CoreInitializeMemoryProtection (
 747   VOID
 748   )
 749 {
 750   EFI_STATUS  Status;
 751   EFI_EVENT   Event;
 752   VOID        *Registration;
 753
 754   mImageProtectionPolicy = PcdGet32(PcdImageProtectionPolicy);
 755
 756   if (mImageProtectionPolicy != 0) {
 757     Status = CoreCreateEvent (
 758                EVT_NOTIFY_SIGNAL,
 759                TPL_CALLBACK,
 760                MemoryProtectionCpuArchProtocolNotify,
 761                NULL,
 762                &Event
 763                );
 764     ASSERT_EFI_ERROR(Status);
 765
 766     //
 767     // Register for protocol notifactions on this event
 768     //
 769     Status = CoreRegisterProtocolNotify (
 770                &gEfiCpuArchProtocolGuid,
 771                Event,
 772                &Registration
 773                );
 774     ASSERT_EFI_ERROR(Status);
 775   }
 776   return ;
 777 }