MdeModulePkg/MemoryProtection: split protect and unprotect paths

Currently, the PE/COFF image memory protection code uses the same code
paths for protecting and unprotecting an image. This is strange, since
unprotecting an image involves a single call into the CPU arch protocol
to clear the permission attributes of the entire range, and there is no
need to parse the PE/COFF headers again.

So let's store the ImageRecord entries in a linked list, so we can find
it again at unprotect time, and simply clear the permissions.

Note that this fixes a DEBUG hang on an ASSERT() that occurs when the
PE/COFF image fails to load, which causes UnprotectUefiImage() to be
invoked before the image is fully loaded.

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Ard Biesheuvel <ard.biesheuvel@linaro.org>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>

diff --git a/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c b/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
index 451cc35b92908f9ba1c83724cdab06ae6f593f8f..93f96f0c9502e28b58e3809ab6de4650d746be09 100644 (file)
--- a/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
+++ b/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
@@ -74,6 +74,8 @@ UINT32   mImageProtectionPolicy;
 \r
 extern LIST_ENTRY         mGcdMemorySpaceMap;\r
 \r
+STATIC LIST_ENTRY         mProtectedImageRecordList;\r
+\r
 /**\r
   Sort code section in image record, based upon CodeSegmentBase from low to high.\r
 \r
@@ -238,13 +240,10 @@ SetUefiImageMemoryAttributes (
   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
+  IN IMAGE_PROPERTIES_RECORD     *ImageRecord\r
   )\r
 {\r
   IMAGE_PROPERTIES_RECORD_CODE_SECTION      *ImageRecordCodeSection;\r
@@ -253,7 +252,6 @@ SetUefiImageProtectionAttributes (
   LIST_ENTRY                                *ImageRecordCodeSectionList;\r
   UINT64                                    CurrentBase;\r
   UINT64                                    ImageEnd;\r
-  UINT64                                    Attribute;\r
 \r
   ImageRecordCodeSectionList = &ImageRecord->CodeSegmentList;\r
 \r
@@ -276,29 +274,19 @@ SetUefiImageProtectionAttributes (
       //\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
+        EFI_MEMORY_XP\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
+      EFI_MEMORY_RO\r
       );\r
     CurrentBase = ImageRecordCodeSection->CodeSegmentBase + ImageRecordCodeSection->CodeSegmentSize;\r
   }\r
@@ -310,15 +298,10 @@ SetUefiImageProtectionAttributes (
     //\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
+      EFI_MEMORY_XP\r
       );\r
   }\r
   return ;\r
@@ -401,18 +384,15 @@ FreeImageRecord (
 }\r
 \r
 /**\r
-  Protect or unprotect UEFI image common function.\r
+  Protect UEFI PE/COFF image\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
+ProtectUefiImage (\r
   IN EFI_LOADED_IMAGE_PROTOCOL   *LoadedImage,\r
-  IN EFI_DEVICE_PATH_PROTOCOL    *LoadedImageDevicePath,\r
-  IN BOOLEAN                     Protect\r
+  IN EFI_DEVICE_PATH_PROTOCOL    *LoadedImageDevicePath\r
   )\r
 {\r
   VOID                                 *ImageAddress;\r
@@ -617,34 +597,17 @@ ProtectUefiImageCommon (
   //\r
   // CPU ARCH present. Update memory attribute directly.\r
   //\r
-  SetUefiImageProtectionAttributes (ImageRecord, Protect);\r
+  SetUefiImageProtectionAttributes (ImageRecord);\r
 \r
   //\r
-  // Clean up\r
+  // Record the image record in the list so we can undo the protections later\r
   //\r
-  FreeImageRecord (ImageRecord);\r
+  InsertTailList (&mProtectedImageRecordList, &ImageRecord->Link);\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
@@ -657,8 +620,28 @@ UnprotectUefiImage (
   IN EFI_DEVICE_PATH_PROTOCOL    *LoadedImageDevicePath\r
   )\r
 {\r
+  IMAGE_PROPERTIES_RECORD    *ImageRecord;\r
+  LIST_ENTRY                 *ImageRecordLink;\r
+\r
   if (PcdGet32(PcdImageProtectionPolicy) != 0) {\r
-    ProtectUefiImageCommon (LoadedImage, LoadedImageDevicePath, FALSE);\r
+    for (ImageRecordLink = mProtectedImageRecordList.ForwardLink;\r
+         ImageRecordLink != &mProtectedImageRecordList;\r
+         ImageRecordLink = ImageRecordLink->ForwardLink) {\r
+      ImageRecord = CR (\r
+                      ImageRecordLink,\r
+                      IMAGE_PROPERTIES_RECORD,\r
+                      Link,\r
+                      IMAGE_PROPERTIES_RECORD_SIGNATURE\r
+                      );\r
+\r
+      if (ImageRecord->ImageBase == (EFI_PHYSICAL_ADDRESS)(UINTN)LoadedImage->ImageBase) {\r
+        SetUefiImageMemoryAttributes (ImageRecord->ImageBase,\r
+                                      ImageRecord->ImageSize,\r
+                                      0);\r
+        FreeImageRecord (ImageRecord);\r
+        return;\r
+      }\r
+    }\r
   }\r
 }\r
 \r
@@ -1027,6 +1010,8 @@ CoreInitializeMemoryProtection (
 \r
   mImageProtectionPolicy = PcdGet32(PcdImageProtectionPolicy);\r
 \r
+  InitializeListHead (&mProtectedImageRecordList);\r
+\r
   //\r
   // Sanity check the PcdDxeNxMemoryProtectionPolicy setting:\r
   // - code regions should have no EFI_MEMORY_XP attribute\r