--- /dev/null
+/** @file\r
+ UEFI Memory Protection support.\r
+\r
+ If the UEFI image is page aligned, the image code section is set to read only\r
+ and the image data section is set to non-executable.\r
+\r
+ 1) This policy is applied for all UEFI image including boot service driver,\r
+ runtime driver or application.\r
+ 2) This policy is applied only if the UEFI image meets the page alignment\r
+ requirement.\r
+ 3) This policy is applied only if the Source UEFI image matches the\r
+ PcdImageProtectionPolicy definition.\r
+ 4) This policy is not applied to the non-PE image region.\r
+\r
+ The DxeCore calls CpuArchProtocol->SetMemoryAttributes() to protect\r
+ the image. If the CpuArch protocol is not installed yet, the DxeCore\r
+ enqueues the protection request. Once the CpuArch is installed, the\r
+ DxeCore dequeues the protection request and applies policy.\r
+\r
+ Once the image is unloaded, the protection is removed automatically.\r
+\r
+Copyright (c) 2017, Intel Corporation. All rights reserved.<BR>\r
+This program and the accompanying materials\r
+are licensed and made available under the terms and conditions of the BSD License\r
+which accompanies this distribution. The full text of the license may be found at\r
+http://opensource.org/licenses/bsd-license.php\r
+\r
+THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,\r
+WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
+\r
+**/\r
+\r
+#include <PiDxe.h>\r
+#include <Library/BaseLib.h>\r
+#include <Library/BaseMemoryLib.h>\r
+#include <Library/MemoryAllocationLib.h>\r
+#include <Library/UefiBootServicesTableLib.h>\r
+#include <Library/DxeServicesTableLib.h>\r
+#include <Library/DebugLib.h>\r
+#include <Library/UefiLib.h>\r
+\r
+#include <Guid/EventGroup.h>\r
+#include <Guid/MemoryAttributesTable.h>\r
+#include <Guid/PropertiesTable.h>\r
+\r
+#include <Protocol/FirmwareVolume2.h>\r
+#include <Protocol/BlockIo.h>\r
+#include <Protocol/SimpleFileSystem.h>\r
+\r
+#include "DxeMain.h"\r
+\r
+#define CACHE_ATTRIBUTE_MASK (EFI_MEMORY_UC | EFI_MEMORY_WC | EFI_MEMORY_WT | EFI_MEMORY_WB | EFI_MEMORY_UCE | EFI_MEMORY_WP)\r
+#define MEMORY_ATTRIBUTE_MASK (EFI_MEMORY_RP | EFI_MEMORY_XP | EFI_MEMORY_RO)\r
+\r
+//\r
+// Image type definitions\r
+//\r
+#define IMAGE_UNKNOWN 0x00000001\r
+#define IMAGE_FROM_FV 0x00000002\r
+\r
+//\r
+// Protection policy bit definition\r
+//\r
+#define DO_NOT_PROTECT 0x00000000\r
+#define PROTECT_IF_ALIGNED_ELSE_ALLOW 0x00000001\r
+\r
+UINT32 mImageProtectionPolicy;\r
+\r
+/**\r
+ Sort code section in image record, based upon CodeSegmentBase from low to high.\r
+\r
+ @param ImageRecord image record to be sorted\r
+**/\r
+VOID\r
+SortImageRecordCodeSection (\r
+ IN IMAGE_PROPERTIES_RECORD *ImageRecord\r
+ );\r
+\r
+/**\r
+ Check if code section in image record is valid.\r
+\r
+ @param ImageRecord image record to be checked\r
+\r
+ @retval TRUE image record is valid\r
+ @retval FALSE image record is invalid\r
+**/\r
+BOOLEAN\r
+IsImageRecordCodeSectionValid (\r
+ IN IMAGE_PROPERTIES_RECORD *ImageRecord\r
+ );\r
+\r
+/**\r
+ Get the image type.\r
+\r
+ @param[in] File This is a pointer to the device path of the file that is\r
+ being dispatched.\r
+\r
+ @return UINT32 Image Type\r
+**/\r
+UINT32\r
+GetImageType (\r
+ IN CONST EFI_DEVICE_PATH_PROTOCOL *File\r
+ )\r
+{\r
+ EFI_STATUS Status;\r
+ EFI_HANDLE DeviceHandle;\r
+ EFI_DEVICE_PATH_PROTOCOL *TempDevicePath;\r
+\r
+ if (File == NULL) {\r
+ return IMAGE_UNKNOWN;\r
+ }\r
+\r
+ //\r
+ // First check to see if File is from a Firmware Volume\r
+ //\r
+ DeviceHandle = NULL;\r
+ TempDevicePath = (EFI_DEVICE_PATH_PROTOCOL *) File;\r
+ Status = gBS->LocateDevicePath (\r
+ &gEfiFirmwareVolume2ProtocolGuid,\r
+ &TempDevicePath,\r
+ &DeviceHandle\r
+ );\r
+ if (!EFI_ERROR (Status)) {\r
+ Status = gBS->OpenProtocol (\r
+ DeviceHandle,\r
+ &gEfiFirmwareVolume2ProtocolGuid,\r
+ NULL,\r
+ NULL,\r
+ NULL,\r
+ EFI_OPEN_PROTOCOL_TEST_PROTOCOL\r
+ );\r
+ if (!EFI_ERROR (Status)) {\r
+ return IMAGE_FROM_FV;\r
+ }\r
+ }\r
+ return IMAGE_UNKNOWN;\r
+}\r
+\r
+/**\r
+ Get UEFI image protection policy based upon image type.\r
+\r
+ @param[in] ImageType The UEFI image type\r
+\r
+ @return UEFI image protection policy\r
+**/\r
+UINT32\r
+GetProtectionPolicyFromImageType (\r
+ IN UINT32 ImageType\r
+ )\r
+{\r
+ if ((ImageType & mImageProtectionPolicy) == 0) {\r
+ return DO_NOT_PROTECT;\r
+ } else {\r
+ return PROTECT_IF_ALIGNED_ELSE_ALLOW;\r
+ }\r
+}\r
+\r
+/**\r
+ Get UEFI image protection policy based upon loaded image device path.\r
+\r
+ @param[in] LoadedImage The loaded image protocol\r
+ @param[in] LoadedImageDevicePath The loaded image device path protocol\r
+\r
+ @return UEFI image protection policy\r
+**/\r
+UINT32\r
+GetUefiImageProtectionPolicy (\r
+ IN EFI_LOADED_IMAGE_PROTOCOL *LoadedImage,\r
+ IN EFI_DEVICE_PATH_PROTOCOL *LoadedImageDevicePath\r
+ )\r
+{\r
+ BOOLEAN InSmm;\r
+ UINT32 ImageType;\r
+ UINT32 ProtectionPolicy;\r
+\r
+ //\r
+ // Check SMM\r
+ //\r
+ InSmm = FALSE;\r
+ if (gSmmBase2 != NULL) {\r
+ gSmmBase2->InSmm (gSmmBase2, &InSmm);\r
+ }\r
+ if (InSmm) {\r
+ return FALSE;\r
+ }\r
+\r
+ //\r
+ // Check DevicePath\r
+ //\r
+ if (LoadedImage == gDxeCoreLoadedImage) {\r
+ ImageType = IMAGE_FROM_FV;\r
+ } else {\r
+ ImageType = GetImageType (LoadedImageDevicePath);\r
+ }\r
+ ProtectionPolicy = GetProtectionPolicyFromImageType (ImageType);\r
+ return ProtectionPolicy;\r
+}\r
+\r
+\r
+/**\r
+ Set UEFI image memory attributes.\r
+\r
+ @param[in] BaseAddress Specified start address\r
+ @param[in] Length Specified length\r
+ @param[in] Attributes Specified attributes\r
+**/\r
+VOID\r
+SetUefiImageMemoryAttributes (\r
+ IN UINT64 BaseAddress,\r
+ IN UINT64 Length,\r
+ IN UINT64 Attributes\r
+ )\r
+{\r
+ EFI_STATUS Status;\r
+ EFI_GCD_MEMORY_SPACE_DESCRIPTOR Descriptor;\r
+ UINT64 FinalAttributes;\r
+\r
+ Status = CoreGetMemorySpaceDescriptor(BaseAddress, &Descriptor);\r
+ ASSERT_EFI_ERROR(Status);\r
+\r
+ FinalAttributes = (Descriptor.Attributes & CACHE_ATTRIBUTE_MASK) | (Attributes & MEMORY_ATTRIBUTE_MASK);\r
+\r
+ DEBUG ((DEBUG_INFO, "SetUefiImageMemoryAttributes - 0x%016lx - 0x%016lx (0x%016lx)\n", BaseAddress, Length, FinalAttributes));\r
+\r
+ ASSERT(gCpu != NULL);\r
+ gCpu->SetMemoryAttributes (gCpu, BaseAddress, Length, FinalAttributes);\r
+}\r
+\r
+/**\r
+ Set UEFI image protection attributes.\r
+\r
+ @param[in] ImageRecord A UEFI image record\r
+ @param[in] Protect TRUE: Protect the UEFI image.\r
+ FALSE: Unprotect the UEFI image.\r
+**/\r
+VOID\r
+SetUefiImageProtectionAttributes (\r
+ IN IMAGE_PROPERTIES_RECORD *ImageRecord,\r
+ IN BOOLEAN Protect\r
+ )\r
+{\r
+ IMAGE_PROPERTIES_RECORD_CODE_SECTION *ImageRecordCodeSection;\r
+ LIST_ENTRY *ImageRecordCodeSectionLink;\r
+ LIST_ENTRY *ImageRecordCodeSectionEndLink;\r
+ LIST_ENTRY *ImageRecordCodeSectionList;\r
+ UINT64 CurrentBase;\r
+ UINT64 ImageEnd;\r
+ UINT64 Attribute;\r
+\r
+ ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;\r
+\r
+ CurrentBase = ImageRecord->ImageBase;\r
+ ImageEnd = ImageRecord->ImageBase + ImageRecord->ImageSize;\r
+\r
+ ImageRecordCodeSectionLink = ImageRecordCodeSectionList->ForwardLink;\r
+ ImageRecordCodeSectionEndLink = ImageRecordCodeSectionList;\r
+ while (ImageRecordCodeSectionLink != ImageRecordCodeSectionEndLink) {\r
+ ImageRecordCodeSection = CR (\r
+ ImageRecordCodeSectionLink,\r
+ IMAGE_PROPERTIES_RECORD_CODE_SECTION,\r
+ Link,\r
+ IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE\r
+ );\r
+ ImageRecordCodeSectionLink = ImageRecordCodeSectionLink->ForwardLink;\r
+\r
+ ASSERT (CurrentBase <= ImageRecordCodeSection->CodeSegmentBase);\r
+ if (CurrentBase < ImageRecordCodeSection->CodeSegmentBase) {\r
+ //\r
+ // DATA\r
+ //\r
+ if (Protect) {\r
+ Attribute = EFI_MEMORY_XP;\r
+ } else {\r
+ Attribute = 0;\r
+ }\r
+ SetUefiImageMemoryAttributes (\r
+ CurrentBase,\r
+ ImageRecordCodeSection->CodeSegmentBase - CurrentBase,\r
+ Attribute\r
+ );\r
+ }\r
+ //\r
+ // CODE\r
+ //\r
+ if (Protect) {\r
+ Attribute = EFI_MEMORY_RO;\r
+ } else {\r
+ Attribute = 0;\r
+ }\r
+ SetUefiImageMemoryAttributes (\r
+ ImageRecordCodeSection->CodeSegmentBase,\r
+ ImageRecordCodeSection->CodeSegmentSize,\r
+ Attribute\r
+ );\r
+ CurrentBase = ImageRecordCodeSection->CodeSegmentBase + ImageRecordCodeSection->CodeSegmentSize;\r
+ }\r
+ //\r
+ // Last DATA\r
+ //\r
+ ASSERT (CurrentBase <= ImageEnd);\r
+ if (CurrentBase < ImageEnd) {\r
+ //\r
+ // DATA\r
+ //\r
+ if (Protect) {\r
+ Attribute = EFI_MEMORY_XP;\r
+ } else {\r
+ Attribute = 0;\r
+ }\r
+ SetUefiImageMemoryAttributes (\r
+ CurrentBase,\r
+ ImageEnd - CurrentBase,\r
+ Attribute\r
+ );\r
+ }\r
+ return ;\r
+}\r
+\r
+/**\r
+ Return if the PE image section is aligned.\r
+\r
+ @param[in] SectionAlignment PE/COFF section alignment\r
+ @param[in] MemoryType PE/COFF image memory type\r
+\r
+ @retval TRUE The PE image section is aligned.\r
+ @retval FALSE The PE image section is not aligned.\r
+**/\r
+BOOLEAN\r
+IsMemoryProtectionSectionAligned (\r
+ IN UINT32 SectionAlignment,\r
+ IN EFI_MEMORY_TYPE MemoryType\r
+ )\r
+{\r
+ UINT32 PageAlignment;\r
+\r
+ switch (MemoryType) {\r
+ case EfiRuntimeServicesCode:\r
+ case EfiACPIMemoryNVS:\r
+ PageAlignment = EFI_ACPI_RUNTIME_PAGE_ALLOCATION_ALIGNMENT;\r
+ break;\r
+ case EfiRuntimeServicesData:\r
+ case EfiACPIReclaimMemory:\r
+ ASSERT (FALSE);\r
+ PageAlignment = EFI_ACPI_RUNTIME_PAGE_ALLOCATION_ALIGNMENT;\r
+ break;\r
+ case EfiBootServicesCode:\r
+ case EfiLoaderCode:\r
+ case EfiReservedMemoryType:\r
+ PageAlignment = EFI_PAGE_SIZE;\r
+ break;\r
+ default:\r
+ ASSERT (FALSE);\r
+ PageAlignment = EFI_PAGE_SIZE;\r
+ break;\r
+ }\r
+\r
+ if ((SectionAlignment & (PageAlignment - 1)) != 0) {\r
+ return FALSE;\r
+ } else {\r
+ return TRUE;\r
+ }\r
+}\r
+\r
+/**\r
+ Free Image record.\r
+\r
+ @param[in] ImageRecord A UEFI image record\r
+**/\r
+VOID\r
+FreeImageRecord (\r
+ IN IMAGE_PROPERTIES_RECORD *ImageRecord\r
+ )\r
+{\r
+ LIST_ENTRY *CodeSegmentListHead;\r
+ IMAGE_PROPERTIES_RECORD_CODE_SECTION *ImageRecordCodeSection;\r
+\r
+ CodeSegmentListHead = &ImageRecord->CodeSegmentList;\r
+ while (!IsListEmpty (CodeSegmentListHead)) {\r
+ ImageRecordCodeSection = CR (\r
+ CodeSegmentListHead->ForwardLink,\r
+ IMAGE_PROPERTIES_RECORD_CODE_SECTION,\r
+ Link,\r
+ IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE\r
+ );\r
+ RemoveEntryList (&ImageRecordCodeSection->Link);\r
+ FreePool (ImageRecordCodeSection);\r
+ }\r
+\r
+ if (ImageRecord->Link.ForwardLink != NULL) {\r
+ RemoveEntryList (&ImageRecord->Link);\r
+ }\r
+ FreePool (ImageRecord);\r
+}\r
+\r
+/**\r
+ Protect or unprotect UEFI image common function.\r
+\r
+ @param[in] LoadedImage The loaded image protocol\r
+ @param[in] LoadedImageDevicePath The loaded image device path protocol\r
+ @param[in] Protect TRUE: Protect the UEFI image.\r
+ FALSE: Unprotect the UEFI image.\r
+**/\r
+VOID\r
+ProtectUefiImageCommon (\r
+ IN EFI_LOADED_IMAGE_PROTOCOL *LoadedImage,\r
+ IN EFI_DEVICE_PATH_PROTOCOL *LoadedImageDevicePath,\r
+ IN BOOLEAN Protect\r
+ )\r
+{\r
+ VOID *ImageAddress;\r
+ EFI_IMAGE_DOS_HEADER *DosHdr;\r
+ UINT32 PeCoffHeaderOffset;\r
+ UINT32 SectionAlignment;\r
+ EFI_IMAGE_SECTION_HEADER *Section;\r
+ EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION Hdr;\r
+ UINT8 *Name;\r
+ UINTN Index;\r
+ IMAGE_PROPERTIES_RECORD *ImageRecord;\r
+ CHAR8 *PdbPointer;\r
+ IMAGE_PROPERTIES_RECORD_CODE_SECTION *ImageRecordCodeSection;\r
+ UINT16 Magic;\r
+ BOOLEAN IsAligned;\r
+ UINT32 ProtectionPolicy;\r
+\r
+ DEBUG ((DEBUG_INFO, "ProtectUefiImageCommon - 0x%x\n", LoadedImage));\r
+ DEBUG ((DEBUG_INFO, " - 0x%016lx - 0x%016lx\n", (EFI_PHYSICAL_ADDRESS)(UINTN)LoadedImage->ImageBase, LoadedImage->ImageSize));\r
+\r
+ if (gCpu == NULL) {\r
+ return ;\r
+ }\r
+\r
+ ProtectionPolicy = GetUefiImageProtectionPolicy (LoadedImage, LoadedImageDevicePath);\r
+ switch (ProtectionPolicy) {\r
+ case DO_NOT_PROTECT:\r
+ return ;\r
+ case PROTECT_IF_ALIGNED_ELSE_ALLOW:\r
+ break;\r
+ default:\r
+ ASSERT(FALSE);\r
+ return ;\r
+ }\r
+\r
+ ImageRecord = AllocateZeroPool (sizeof(*ImageRecord));\r
+ if (ImageRecord == NULL) {\r
+ return ;\r
+ }\r
+ ImageRecord->Signature = IMAGE_PROPERTIES_RECORD_SIGNATURE;\r
+\r
+ //\r
+ // Step 1: record whole region\r
+ //\r
+ ImageRecord->ImageBase = (EFI_PHYSICAL_ADDRESS)(UINTN)LoadedImage->ImageBase;\r
+ ImageRecord->ImageSize = LoadedImage->ImageSize;\r
+\r
+ ImageAddress = LoadedImage->ImageBase;\r
+\r
+ PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);\r
+ if (PdbPointer != NULL) {\r
+ DEBUG ((DEBUG_VERBOSE, " Image - %a\n", PdbPointer));\r
+ }\r
+\r
+ //\r
+ // Check PE/COFF image\r
+ //\r
+ DosHdr = (EFI_IMAGE_DOS_HEADER *) (UINTN) ImageAddress;\r
+ PeCoffHeaderOffset = 0;\r
+ if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {\r
+ PeCoffHeaderOffset = DosHdr->e_lfanew;\r
+ }\r
+\r
+ Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINT8 *) (UINTN) ImageAddress + PeCoffHeaderOffset);\r
+ if (Hdr.Pe32->Signature != EFI_IMAGE_NT_SIGNATURE) {\r
+ DEBUG ((DEBUG_VERBOSE, "Hdr.Pe32->Signature invalid - 0x%x\n", Hdr.Pe32->Signature));\r
+ // It might be image in SMM.\r
+ goto Finish;\r
+ }\r
+\r
+ //\r
+ // Get SectionAlignment\r
+ //\r
+ if (Hdr.Pe32->FileHeader.Machine == IMAGE_FILE_MACHINE_IA64 && Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+ //\r
+ // NOTE: Some versions of Linux ELILO for Itanium have an incorrect magic value\r
+ // in the PE/COFF Header. If the MachineType is Itanium(IA64) and the\r
+ // Magic value in the OptionalHeader is EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC\r
+ // then override the magic value to EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC\r
+ //\r
+ Magic = EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;\r
+ } else {\r
+ //\r
+ // Get the magic value from the PE/COFF Optional Header\r
+ //\r
+ Magic = Hdr.Pe32->OptionalHeader.Magic;\r
+ }\r
+ if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
+ SectionAlignment = Hdr.Pe32->OptionalHeader.SectionAlignment;\r
+ } else {\r
+ SectionAlignment = Hdr.Pe32Plus->OptionalHeader.SectionAlignment;\r
+ }\r
+\r
+ IsAligned = IsMemoryProtectionSectionAligned (SectionAlignment, LoadedImage->ImageCodeType);\r
+ if (!IsAligned) {\r
+ DEBUG ((DEBUG_VERBOSE, "!!!!!!!! ProtectUefiImageCommon - Section Alignment(0x%x) is incorrect !!!!!!!!\n",\r
+ SectionAlignment));\r
+ PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);\r
+ if (PdbPointer != NULL) {\r
+ DEBUG ((DEBUG_VERBOSE, "!!!!!!!! Image - %a !!!!!!!!\n", PdbPointer));\r
+ }\r
+ goto Finish;\r
+ }\r
+\r
+ Section = (EFI_IMAGE_SECTION_HEADER *) (\r
+ (UINT8 *) (UINTN) ImageAddress +\r
+ PeCoffHeaderOffset +\r
+ sizeof(UINT32) +\r
+ sizeof(EFI_IMAGE_FILE_HEADER) +\r
+ Hdr.Pe32->FileHeader.SizeOfOptionalHeader\r
+ );\r
+ ImageRecord->CodeSegmentCount = 0;\r
+ InitializeListHead (&ImageRecord->CodeSegmentList);\r
+ for (Index = 0; Index < Hdr.Pe32->FileHeader.NumberOfSections; Index++) {\r
+ Name = Section[Index].Name;\r
+ DEBUG ((\r
+ DEBUG_VERBOSE,\r
+ " Section - '%c%c%c%c%c%c%c%c'\n",\r
+ Name[0],\r
+ Name[1],\r
+ Name[2],\r
+ Name[3],\r
+ Name[4],\r
+ Name[5],\r
+ Name[6],\r
+ Name[7]\r
+ ));\r
+\r
+ if ((Section[Index].Characteristics & EFI_IMAGE_SCN_CNT_CODE) != 0) {\r
+ DEBUG ((DEBUG_VERBOSE, " VirtualSize - 0x%08x\n", Section[Index].Misc.VirtualSize));\r
+ DEBUG ((DEBUG_VERBOSE, " VirtualAddress - 0x%08x\n", Section[Index].VirtualAddress));\r
+ DEBUG ((DEBUG_VERBOSE, " SizeOfRawData - 0x%08x\n", Section[Index].SizeOfRawData));\r
+ DEBUG ((DEBUG_VERBOSE, " PointerToRawData - 0x%08x\n", Section[Index].PointerToRawData));\r
+ DEBUG ((DEBUG_VERBOSE, " PointerToRelocations - 0x%08x\n", Section[Index].PointerToRelocations));\r
+ DEBUG ((DEBUG_VERBOSE, " PointerToLinenumbers - 0x%08x\n", Section[Index].PointerToLinenumbers));\r
+ DEBUG ((DEBUG_VERBOSE, " NumberOfRelocations - 0x%08x\n", Section[Index].NumberOfRelocations));\r
+ DEBUG ((DEBUG_VERBOSE, " NumberOfLinenumbers - 0x%08x\n", Section[Index].NumberOfLinenumbers));\r
+ DEBUG ((DEBUG_VERBOSE, " Characteristics - 0x%08x\n", Section[Index].Characteristics));\r
+\r
+ //\r
+ // Step 2: record code section\r
+ //\r
+ ImageRecordCodeSection = AllocatePool (sizeof(*ImageRecordCodeSection));\r
+ if (ImageRecordCodeSection == NULL) {\r
+ return ;\r
+ }\r
+ ImageRecordCodeSection->Signature = IMAGE_PROPERTIES_RECORD_CODE_SECTION_SIGNATURE;\r
+\r
+ ImageRecordCodeSection->CodeSegmentBase = (UINTN)ImageAddress + Section[Index].VirtualAddress;\r
+ ImageRecordCodeSection->CodeSegmentSize = ALIGN_VALUE(Section[Index].SizeOfRawData, SectionAlignment);\r
+\r
+ DEBUG ((DEBUG_VERBOSE, "ImageCode: 0x%016lx - 0x%016lx\n", ImageRecordCodeSection->CodeSegmentBase, ImageRecordCodeSection->CodeSegmentSize));\r
+\r
+ InsertTailList (&ImageRecord->CodeSegmentList, &ImageRecordCodeSection->Link);\r
+ ImageRecord->CodeSegmentCount++;\r
+ }\r
+ }\r
+\r
+ if (ImageRecord->CodeSegmentCount == 0) {\r
+ DEBUG ((DEBUG_ERROR, "!!!!!!!! ProtectUefiImageCommon - CodeSegmentCount is 0 !!!!!!!!\n"));\r
+ PdbPointer = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageAddress);\r
+ if (PdbPointer != NULL) {\r
+ DEBUG ((DEBUG_ERROR, "!!!!!!!! Image - %a !!!!!!!!\n", PdbPointer));\r
+ }\r
+ goto Finish;\r
+ }\r
+\r
+ //\r
+ // Final\r
+ //\r
+ SortImageRecordCodeSection (ImageRecord);\r
+ //\r
+ // Check overlap all section in ImageBase/Size\r
+ //\r
+ if (!IsImageRecordCodeSectionValid (ImageRecord)) {\r
+ DEBUG ((DEBUG_ERROR, "IsImageRecordCodeSectionValid - FAIL\n"));\r
+ goto Finish;\r
+ }\r
+\r
+ //\r
+ // Round up the ImageSize, some CPU arch may return EFI_UNSUPPORTED if ImageSize is not aligned.\r
+ // Given that the loader always allocates full pages, we know the space after the image is not used.\r
+ //\r
+ ImageRecord->ImageSize = ALIGN_VALUE(LoadedImage->ImageSize, EFI_PAGE_SIZE);\r
+\r
+ //\r
+ // CPU ARCH present. Update memory attribute directly.\r
+ //\r
+ SetUefiImageProtectionAttributes (ImageRecord, Protect);\r
+\r
+ //\r
+ // Clean up\r
+ //\r
+ FreeImageRecord (ImageRecord);\r
+\r
+Finish:\r
+ return ;\r
+}\r
+\r
+/**\r
+ Protect UEFI image.\r
+\r
+ @param[in] LoadedImage The loaded image protocol\r
+ @param[in] LoadedImageDevicePath The loaded image device path protocol\r
+**/\r
+VOID\r
+ProtectUefiImage (\r
+ IN EFI_LOADED_IMAGE_PROTOCOL *LoadedImage,\r
+ IN EFI_DEVICE_PATH_PROTOCOL *LoadedImageDevicePath\r
+ )\r
+{\r
+ if (PcdGet32(PcdImageProtectionPolicy) != 0) {\r
+ ProtectUefiImageCommon (LoadedImage, LoadedImageDevicePath, TRUE);\r
+ }\r
+}\r
+\r
+/**\r
+ Unprotect UEFI image.\r
+\r
+ @param[in] LoadedImage The loaded image protocol\r
+ @param[in] LoadedImageDevicePath The loaded image device path protocol\r
+**/\r
+VOID\r
+UnprotectUefiImage (\r
+ IN EFI_LOADED_IMAGE_PROTOCOL *LoadedImage,\r
+ IN EFI_DEVICE_PATH_PROTOCOL *LoadedImageDevicePath\r
+ )\r
+{\r
+ if (PcdGet32(PcdImageProtectionPolicy) != 0) {\r
+ ProtectUefiImageCommon (LoadedImage, LoadedImageDevicePath, FALSE);\r
+ }\r
+}\r
+\r
+/**\r
+ A notification for CPU_ARCH protocol.\r
+\r
+ @param[in] Event Event whose notification function is being invoked.\r
+ @param[in] Context Pointer to the notification function's context,\r
+ which is implementation-dependent.\r
+\r
+**/\r
+VOID\r
+EFIAPI\r
+MemoryProtectionCpuArchProtocolNotify (\r
+ IN EFI_EVENT Event,\r
+ IN VOID *Context\r
+ )\r
+{\r
+ EFI_STATUS Status;\r
+ EFI_LOADED_IMAGE_PROTOCOL *LoadedImage;\r
+ EFI_DEVICE_PATH_PROTOCOL *LoadedImageDevicePath;\r
+ UINTN NoHandles;\r
+ EFI_HANDLE *HandleBuffer;\r
+ UINTN Index;\r
+\r
+ DEBUG ((DEBUG_INFO, "MemoryProtectionCpuArchProtocolNotify:\n"));\r
+ Status = CoreLocateProtocol (&gEfiCpuArchProtocolGuid, NULL, (VOID **)&gCpu);\r
+ if (EFI_ERROR (Status)) {\r
+ return;\r
+ }\r
+\r
+ Status = gBS->LocateHandleBuffer (\r
+ ByProtocol,\r
+ &gEfiLoadedImageProtocolGuid,\r
+ NULL,\r
+ &NoHandles,\r
+ &HandleBuffer\r
+ );\r
+ if (EFI_ERROR (Status) && (NoHandles == 0)) {\r
+ return ;\r
+ }\r
+\r
+ for (Index = 0; Index < NoHandles; Index++) {\r
+ Status = gBS->HandleProtocol (\r
+ HandleBuffer[Index],\r
+ &gEfiLoadedImageProtocolGuid,\r
+ (VOID **)&LoadedImage\r
+ );\r
+ if (EFI_ERROR(Status)) {\r
+ continue;\r
+ }\r
+ Status = gBS->HandleProtocol (\r
+ HandleBuffer[Index],\r
+ &gEfiLoadedImageDevicePathProtocolGuid,\r
+ (VOID **)&LoadedImageDevicePath\r
+ );\r
+ if (EFI_ERROR(Status)) {\r
+ LoadedImageDevicePath = NULL;\r
+ }\r
+\r
+ ProtectUefiImage (LoadedImage, LoadedImageDevicePath);\r
+ }\r
+\r
+ CoreCloseEvent (Event);\r
+ return;\r
+}\r
+\r
+/**\r
+ ExitBootServices Callback function for memory protection.\r
+**/\r
+VOID\r
+MemoryProtectionExitBootServicesCallback (\r
+ VOID\r
+ )\r
+{\r
+ EFI_RUNTIME_IMAGE_ENTRY *RuntimeImage;\r
+ LIST_ENTRY *Link;\r
+\r
+ //\r
+ // We need remove the RT protection, because RT relocation need write code segment\r
+ // at SetVirtualAddressMap(). We cannot assume OS/Loader has taken over page table at that time.\r
+ //\r
+ // Firmware does not own page tables after ExitBootServices(), so the OS would\r
+ // have to relax protection of RT code pages across SetVirtualAddressMap(), or\r
+ // delay setting protections on RT code pages until after SetVirtualAddressMap().\r
+ // OS may set protection on RT based upon EFI_MEMORY_ATTRIBUTES_TABLE later.\r
+ //\r
+ if (mImageProtectionPolicy != 0) {\r
+ for (Link = gRuntime->ImageHead.ForwardLink; Link != &gRuntime->ImageHead; Link = Link->ForwardLink) {\r
+ RuntimeImage = BASE_CR (Link, EFI_RUNTIME_IMAGE_ENTRY, Link);\r
+ SetUefiImageMemoryAttributes ((UINT64)(UINTN)RuntimeImage->ImageBase, ALIGN_VALUE(RuntimeImage->ImageSize, EFI_PAGE_SIZE), 0);\r
+ }\r
+ }\r
+}\r
+\r
+/**\r
+ Initialize Memory Protection support.\r
+**/\r
+VOID\r
+EFIAPI\r
+CoreInitializeMemoryProtection (\r
+ VOID\r
+ )\r
+{\r
+ EFI_STATUS Status;\r
+ EFI_EVENT Event;\r
+ VOID *Registration;\r
+\r
+ mImageProtectionPolicy = PcdGet32(PcdImageProtectionPolicy);\r
+\r
+ if (mImageProtectionPolicy != 0) {\r
+ Status = CoreCreateEvent (\r
+ EVT_NOTIFY_SIGNAL,\r
+ TPL_CALLBACK,\r
+ MemoryProtectionCpuArchProtocolNotify,\r
+ NULL,\r
+ &Event\r
+ );\r
+ ASSERT_EFI_ERROR(Status);\r
+\r
+ //\r
+ // Register for protocol notifactions on this event\r
+ //\r
+ Status = CoreRegisterProtocolNotify (\r
+ &gEfiCpuArchProtocolGuid,\r
+ Event,\r
+ &Registration\r
+ );\r
+ ASSERT_EFI_ERROR(Status);\r
+ }\r
+ return ;\r
+}\r