MdeModulePkg DxeCore/PiSmmCore/MemoryProfileInfo: Fix EBC and VS2013 build failure.

[mirror_edk2.git] / MdeModulePkg / Application / MemoryProfileInfo / MemoryProfileInfo.c

diff --git a/MdeModulePkg/Application/MemoryProfileInfo/MemoryProfileInfo.c b/MdeModulePkg/Application/MemoryProfileInfo/MemoryProfileInfo.c
index 4896ebd1ff0bedd1fda747a046a2a48e9cc8b8f5..084b069fb8a70c4c57900213d2ed98d09e3e0dab 100644 (file)
--- a/MdeModulePkg/Application/MemoryProfileInfo/MemoryProfileInfo.c
+++ b/MdeModulePkg/Application/MemoryProfileInfo/MemoryProfileInfo.c
@@ -557,6 +557,7 @@ GetUefiMemoryProfileData (
   Size = Size + sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
   Data = AllocateZeroPool ((UINTN) Size);\r
   if (Data == NULL) {\r
+    Status = EFI_OUT_OF_RESOURCES;\r
     Print (L"UefiMemoryProfile: AllocateZeroPool (0x%x) - %r\n", Size, Status);\r
     return Status;\r
   }\r
@@ -597,7 +598,7 @@ GetSmramProfileData (
 {\r
   EFI_STATUS                                    Status;\r
   UINTN                                         CommSize;\r
-  UINT8                                         CommBuffer[sizeof (EFI_GUID) + sizeof (UINTN) + sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA)];\r
+  UINT8                                         *CommBuffer;\r
   EFI_SMM_COMMUNICATE_HEADER                    *CommHeader;\r
   SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO      *CommGetProfileInfo;\r
   SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA      *CommGetProfileData;\r
@@ -611,6 +612,14 @@ GetSmramProfileData (
     return Status;\r
   }\r
 \r
+  CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA);\r
+  CommBuffer = AllocateZeroPool (CommSize);\r
+  if (CommBuffer == NULL) {\r
+    Status = EFI_OUT_OF_RESOURCES;\r
+    Print (L"SmramProfile: AllocateZeroPool (0x%x) for comm buffer - %r\n", CommSize, Status);\r
+    return Status;\r
+  }\r
+\r
   //\r
   // Get Size\r
   //\r
@@ -627,6 +636,7 @@ GetSmramProfileData (
   CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + CommHeader->MessageLength;\r
   Status = SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);\r
   if (EFI_ERROR (Status)) {\r
+    FreePool (CommBuffer);\r
     DEBUG ((EFI_D_ERROR, "SmramProfile: SmmCommunication - %r\n", Status));\r
     return Status;\r
   }\r
@@ -643,15 +653,17 @@ GetSmramProfileData (
   //\r
   ProfileBuffer = (PHYSICAL_ADDRESS) (UINTN) AllocateZeroPool ((UINTN) ProfileSize);\r
   if (ProfileBuffer == 0) {\r
-    Print (L"UefiMemoryProfile: AllocateZeroPool (0x%x) - %r\n", (UINTN) ProfileSize, Status);\r
-    return EFI_SUCCESS;\r
+    FreePool (CommBuffer);\r
+    Status = EFI_OUT_OF_RESOURCES;\r
+    Print (L"SmramProfile: AllocateZeroPool (0x%x) for profile buffer - %r\n", (UINTN) ProfileSize, Status);\r
+    return Status;\r
   }\r
 \r
   CommHeader = (EFI_SMM_COMMUNICATE_HEADER *) &CommBuffer[0];\r
   CopyMem (&CommHeader->HeaderGuid, &gEdkiiMemoryProfileGuid, sizeof(gEdkiiMemoryProfileGuid));\r
   CommHeader->MessageLength = sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA);\r
 \r
-  CommGetProfileData = (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA *)&CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];\r
+  CommGetProfileData = (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA *) &CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];\r
   CommGetProfileData->Header.Command      = SMRAM_PROFILE_COMMAND_GET_PROFILE_DATA;\r
   CommGetProfileData->Header.DataLength   = sizeof (*CommGetProfileData);\r
   CommGetProfileData->Header.ReturnStatus = (UINT64)-1;\r
@@ -663,6 +675,8 @@ GetSmramProfileData (
   ASSERT_EFI_ERROR (Status);\r
 \r
   if (CommGetProfileData->Header.ReturnStatus != 0) {\r
+    FreePool ((VOID *) (UINTN) CommGetProfileData->ProfileBuffer);\r
+    FreePool (CommBuffer);\r
     Print (L"GetProfileData - 0x%x\n", CommGetProfileData->Header.ReturnStatus);\r
     return EFI_SUCCESS;\r
   }\r
@@ -674,6 +688,7 @@ GetSmramProfileData (
   Print (L"======= SmramProfile end =======\n\n\n");\r
 \r
   FreePool ((VOID *) (UINTN) CommGetProfileData->ProfileBuffer);\r
+  FreePool (CommBuffer);\r
 \r
   return EFI_SUCCESS;\r
 }\r