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