diff mbox

[edk2] MdeModulePkg/MemoryProtection: split protect and unprotect paths

Message ID 1490104148-4193-1-git-send-email-ard.biesheuvel@linaro.org
State Accepted
Commit 5920a9d16b1ab887c2858224316a98e961d71b05
Headers show

Commit Message

Ard Biesheuvel March 21, 2017, 1:49 p.m. UTC 
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>

---

This is an alternative fix for the issue addressed by patch 

   'MdeModulePkg/DxeCore: ignore PdbPointer if ImageAddress == 0'

(https://lists.01.org/pipermail/edk2-devel/2017-March/008766.html)

 MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c | 85 ++++++++------------
 1 file changed, 35 insertions(+), 50 deletions(-)

-- 
2.7.4

_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel

Comments

Yao, Jiewen March 21, 2017, 2:14 p.m. UTC | #1 
Good idea.

Reviewed-by: Jiewen.yao@Intel.com




> -----Original Message-----

> From: Ard Biesheuvel [mailto:ard.biesheuvel@linaro.org]

> Sent: Tuesday, March 21, 2017 9:49 PM

> To: edk2-devel@lists.01.org; Yao, Jiewen <jiewen.yao@intel.com>

> Cc: Zeng, Star <star.zeng@intel.com>; Tian, Feng <feng.tian@intel.com>; Gao,

> Liming <liming.gao@intel.com>; Ard Biesheuvel <ard.biesheuvel@linaro.org>

> Subject: [PATCH] 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>

> ---

> 

> This is an alternative fix for the issue addressed by patch

> 

>    'MdeModulePkg/DxeCore: ignore PdbPointer if ImageAddress == 0'

> 

> (https://lists.01.org/pipermail/edk2-devel/2017-March/008766.html)

> 

>  MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c | 85 ++++++++------------

>  1 file changed, 35 insertions(+), 50 deletions(-)

> 

> diff --git a/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c

> b/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c

> index 451cc35b9290..93f96f0c9502 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

> --

> 2.7.4


_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel
Ard Biesheuvel March 21, 2017, 9:30 p.m. UTC | #2 
On 21 March 2017 at 14:14, Yao, Jiewen <jiewen.yao@intel.com> wrote:
> Good idea.

>

> Reviewed-by: Jiewen.yao@Intel.com

>


Pushed, thanks.


>

>

>> -----Original Message-----

>> From: Ard Biesheuvel [mailto:ard.biesheuvel@linaro.org]

>> Sent: Tuesday, March 21, 2017 9:49 PM

>> To: edk2-devel@lists.01.org; Yao, Jiewen <jiewen.yao@intel.com>

>> Cc: Zeng, Star <star.zeng@intel.com>; Tian, Feng <feng.tian@intel.com>; Gao,

>> Liming <liming.gao@intel.com>; Ard Biesheuvel <ard.biesheuvel@linaro.org>

>> Subject: [PATCH] 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>

>> ---

>>

>> This is an alternative fix for the issue addressed by patch

>>

>>    'MdeModulePkg/DxeCore: ignore PdbPointer if ImageAddress == 0'

>>

>> (https://lists.01.org/pipermail/edk2-devel/2017-March/008766.html)

>>

>>  MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c | 85 ++++++++------------

>>  1 file changed, 35 insertions(+), 50 deletions(-)

>>

>> diff --git a/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c

>> b/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c

>> index 451cc35b9290..93f96f0c9502 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

>> --

>> 2.7.4

>

_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel
diff mbox

Patch

diff --git a/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c b/MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c
index 451cc35b9290..93f96f0c9502 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