MdeModulePkg/DxeCore: Re-use ImageIsFromFv for memory protection

MdeModulePkg/Core/Dxe/DxeMain.h

@@ -2727,11 +2727,14 @@ RemoveImageRecord (
   Protect UEFI image.
 
   @param[in]  LoadedImage              The loaded image protocol
+  @param[in]  ImageType                Whether File comes from FV. Must be FALSE
+                                       or TRUE.
   @param[in]  LoadedImageDevicePath    The loaded image device path protocol
 **/
 VOID
 ProtectUefiImage (
-  IN LOADED_IMAGE_PRIVATE_DATA  *Image,
+  IN EFI_LOADED_IMAGE_PROTOCOL     *LoadedImage,
+  IN BOOLEAN                       ImageIsFromFv,
   UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   );
 

MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c

@@ -331,7 +331,7 @@ DxeMain (
 
   CoreInitializeMemoryProtection ();
 
-  ProtectUefiImage (mCurrentImage, &ImageContext);
+  ProtectUefiImage (&mCurrentImage->Info, TRUE, &ImageContext);
 
   //
   // Call constructor for all libraries

MdeModulePkg/Core/Dxe/Image/Image.c

@@ -1383,7 +1383,7 @@ CoreLoadImageCommon (
   }
 
   Status = EFI_SUCCESS;
-  ProtectUefiImage (Image, &ImageContext);
+  ProtectUefiImage (&Image->Info, ImageIsFromFv, &ImageContext);
 
   RegisterMemoryProfileImage (
     Image->LoadedImageDevicePath,

MdeModulePkg/Core/Dxe/Misc/MemoryProtection.c

@@ -47,12 +47,6 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include "ProcessorBind.h"
 #include "Uefi/UefiMultiPhase.h"
 
-//
-// Image type definitions
-//
-#define IMAGE_UNKNOWN  0x00000001
-#define IMAGE_FROM_FV  0x00000002
-
 //
 // Protection policy bit definition
 //
@@ -71,67 +65,21 @@ extern LIST_ENTRY  mGcdMemorySpaceMap;
 
 STATIC LIST_ENTRY  mProtectedImageRecordList;
 
-/**
-  Get the image type.
-
-  @param[in]    File       This is a pointer to the device path of the file that is
-                           being dispatched.
-
-  @return UINT32           Image Type
-**/
-UINT32
-GetImageType (
-  IN  CONST EFI_DEVICE_PATH_PROTOCOL  *File
-  )
-{
-  EFI_STATUS                Status;
-  EFI_HANDLE                DeviceHandle;
-  EFI_DEVICE_PATH_PROTOCOL  *TempDevicePath;
-
-  if (File == NULL) {
-    return IMAGE_UNKNOWN;
-  }
-
-  //
-  // First check to see if File is from a Firmware Volume
-  //
-  DeviceHandle   = NULL;
-  TempDevicePath = (EFI_DEVICE_PATH_PROTOCOL *)File;
-  Status         = gBS->LocateDevicePath (
-                          &gEfiFirmwareVolume2ProtocolGuid,
-                          &TempDevicePath,
-                          &DeviceHandle
-                          );
-  if (!EFI_ERROR (Status)) {
-    Status = gBS->OpenProtocol (
-                    DeviceHandle,
-                    &gEfiFirmwareVolume2ProtocolGuid,
-                    NULL,
-                    NULL,
-                    NULL,
-                    EFI_OPEN_PROTOCOL_TEST_PROTOCOL
-                    );
-    if (!EFI_ERROR (Status)) {
-      return IMAGE_FROM_FV;
-    }
-  }
-
-  return IMAGE_UNKNOWN;
-}
-
 /**
   Get UEFI image protection policy based upon image type.
 
-  @param[in]  ImageType    The UEFI image type
+  @param[in]  ImageIsFromFv  Whether File comes from FV. Must be FALSE or TRUE.
 
   @return UEFI image protection policy
 **/
 UINT32
 GetProtectionPolicyFromImageType (
-  IN UINT32  ImageType
+  IN BOOLEAN  ImageIsFromFv
   )
 {
-  if ((ImageType & mImageProtectionPolicy) == 0) {
+  ASSERT (ImageIsFromFv == FALSE || ImageIsFromFv == TRUE);
+
+  if (((ImageIsFromFv + 1) & mImageProtectionPolicy) == 0) {
     return DO_NOT_PROTECT;
   } else {
     return PROTECT_IF_ALIGNED_ELSE_ALLOW;
@@ -141,19 +89,16 @@ GetProtectionPolicyFromImageType (
 /**
   Get UEFI image protection policy based upon loaded image device path.
 
-  @param[in]  LoadedImage              The loaded image protocol
+  @param[in]  ImageIsFromFv  Whether File comes from FV. Must be FALSE or TRUE.
-  @param[in]  LoadedImageDevicePath    The loaded image device path protocol
 
   @return UEFI image protection policy
 **/
 UINT32
 GetUefiImageProtectionPolicy (
-  IN EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage,
+  IN BOOLEAN  ImageIsFromFv
-  IN EFI_DEVICE_PATH_PROTOCOL   *LoadedImageDevicePath
   )
 {
   BOOLEAN  InSmm;
-  UINT32   ImageType;
   UINT32   ProtectionPolicy;
 
   //
@@ -168,16 +113,7 @@ GetUefiImageProtectionPolicy (
     return FALSE;
   }
 
-  //
+  ProtectionPolicy = GetProtectionPolicyFromImageType (ImageIsFromFv);
-  // Check DevicePath
-  //
-  if (LoadedImage == gDxeCoreLoadedImage) {
-    ImageType = IMAGE_FROM_FV;
-  } else {
-    ImageType = GetImageType (LoadedImageDevicePath);
-  }
-
-  ProtectionPolicy = GetProtectionPolicyFromImageType (ImageType);
   return ProtectionPolicy;
 }
 
@@ -288,17 +224,18 @@ IsMemoryProtectionSectionAligned (
   Protect UEFI PE/COFF image.
 
   @param[in]  LoadedImage              The loaded image protocol
+  @param[in]  ImageIsFromFv            Whether File comes from FV. Must be FALSE
+                                       or TRUE.
   @param[in]  LoadedImageDevicePath    The loaded image device path protocol
 **/
 VOID
 ProtectUefiImage (
-  IN LOADED_IMAGE_PRIVATE_DATA     *Image,
+  IN EFI_LOADED_IMAGE_PROTOCOL     *LoadedImage,
+  IN BOOLEAN                       ImageIsFromFv,
   UEFI_IMAGE_LOADER_IMAGE_CONTEXT  *ImageContext
   )
 {
   RETURN_STATUS              PdbStatus;
-  EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage;
-  EFI_DEVICE_PATH_PROTOCOL   *LoadedImageDevicePath;
   UINT32                     SectionAlignment;
   UEFI_IMAGE_RECORD          *ImageRecord;
   CONST CHAR8                *PdbPointer;
@@ -306,13 +243,10 @@ ProtectUefiImage (
   BOOLEAN                    IsAligned;
   UINT32                     ProtectionPolicy;
 
-  LoadedImage = &Image->Info;
-  LoadedImageDevicePath = Image->LoadedImageDevicePath;
-
   DEBUG ((DEBUG_INFO, "ProtectUefiImageCommon - 0x%x\n", LoadedImage));
   DEBUG ((DEBUG_INFO, "  - 0x%016lx - 0x%016lx\n", (EFI_PHYSICAL_ADDRESS)(UINTN)LoadedImage->ImageBase, LoadedImage->ImageSize));
 
-  ProtectionPolicy = GetUefiImageProtectionPolicy (LoadedImage, LoadedImageDevicePath);
+  ProtectionPolicy = GetUefiImageProtectionPolicy (ImageIsFromFv);
   switch (ProtectionPolicy) {
     case DO_NOT_PROTECT:
       return;