Update DxeCore to set efi memory type for the different image to be loaded.

[mirror_edk2.git] / MdeModulePkg / Core / Dxe / Image / Image.c

diff --git a/MdeModulePkg/Core/Dxe/Image/Image.c b/MdeModulePkg/Core/Dxe/Image/Image.c
index 9232cabcef78cf5d21ae1d7ead4f43ad072643a9..31b1445bf3698d207be0f890b580da76fadceb1c 100644 (file)
--- a/MdeModulePkg/Core/Dxe/Image/Image.c
+++ b/MdeModulePkg/Core/Dxe/Image/Image.c
@@ -225,7 +225,28 @@ Returns:
     //\r
     return EFI_UNSUPPORTED;\r
   }\r
-\r
+  \r
+  //\r
+  // Set EFI memory type based on ImageType\r
+  //\r
+  switch (Image->ImageContext.ImageType) {\r
+  case EFI_IMAGE_SUBSYSTEM_EFI_APPLICATION:\r
+    Image->ImageContext.ImageCodeMemoryType = EfiLoaderCode;\r
+    Image->ImageContext.ImageDataMemoryType = EfiLoaderData;\r
+    break;\r
+  case EFI_IMAGE_SUBSYSTEM_EFI_BOOT_SERVICE_DRIVER:\r
+    Image->ImageContext.ImageCodeMemoryType = EfiBootServicesCode;\r
+    Image->ImageContext.ImageDataMemoryType = EfiBootServicesData;\r
+    break;\r
+  case EFI_IMAGE_SUBSYSTEM_EFI_RUNTIME_DRIVER:\r
+  case EFI_IMAGE_SUBSYSTEM_SAL_RUNTIME_DRIVER:\r
+    Image->ImageContext.ImageCodeMemoryType = EfiRuntimeServicesCode;\r
+    Image->ImageContext.ImageDataMemoryType = EfiRuntimeServicesData;\r
+    break;\r
+  default:\r
+    Image->ImageContext.ImageError = IMAGE_ERROR_INVALID_SUBSYSTEM;\r
+    return EFI_UNSUPPORTED;\r
+  }\r
 \r
   //\r
   // Allocate memory of the correct memory type aligned on the required image boundry\r