3 files changed, 23 insertions, 4 deletions
diff --git a/MdeModulePkg/Application/MemoryProfileInfo/MemoryProfileInfo.c b/MdeModulePkg/Application/MemoryProfileInfo/MemoryProfileInfo.c
index 4896ebd1ff..084b069fb8 100644
--- a/MdeModulePkg/Application/MemoryProfileInfo/MemoryProfileInfo.c
+++ b/MdeModulePkg/Application/MemoryProfileInfo/MemoryProfileInfo.c
@@ -557,6 +557,7 @@ GetUefiMemoryProfileData (
   Size = Size + sizeof (MEMORY_PROFILE_ALLOC_INFO);
   Data = AllocateZeroPool ((UINTN) Size);
   if (Data == NULL) {
+    Status = EFI_OUT_OF_RESOURCES;
     Print (L"UefiMemoryProfile: AllocateZeroPool (0x%x) - %r\n", Size, Status);
     return Status;
   }
@@ -597,7 +598,7 @@ GetSmramProfileData (
 {
   EFI_STATUS                                    Status;
   UINTN                                         CommSize;
-  UINT8                                         CommBuffer[sizeof (EFI_GUID) + sizeof (UINTN) + sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA)];
+  UINT8                                         *CommBuffer;
   EFI_SMM_COMMUNICATE_HEADER                    *CommHeader;
   SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO      *CommGetProfileInfo;
   SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA      *CommGetProfileData;
@@ -611,6 +612,14 @@ GetSmramProfileData (
     return Status;
   }
 
+  CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA);
+  CommBuffer = AllocateZeroPool (CommSize);
+  if (CommBuffer == NULL) {
+    Status = EFI_OUT_OF_RESOURCES;
+    Print (L"SmramProfile: AllocateZeroPool (0x%x) for comm buffer - %r\n", CommSize, Status);
+    return Status;
+  }
+
   //
   // Get Size
   //
@@ -627,6 +636,7 @@ GetSmramProfileData (
   CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + CommHeader->MessageLength;
   Status = SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);
   if (EFI_ERROR (Status)) {
+    FreePool (CommBuffer);
     DEBUG ((EFI_D_ERROR, "SmramProfile: SmmCommunication - %r\n", Status));
     return Status;
   }
@@ -643,15 +653,17 @@ GetSmramProfileData (
   //
   ProfileBuffer = (PHYSICAL_ADDRESS) (UINTN) AllocateZeroPool ((UINTN) ProfileSize);
   if (ProfileBuffer == 0) {
-    Print (L"UefiMemoryProfile: AllocateZeroPool (0x%x) - %r\n", (UINTN) ProfileSize, Status);
-    return EFI_SUCCESS;
+    FreePool (CommBuffer);
+    Status = EFI_OUT_OF_RESOURCES;
+    Print (L"SmramProfile: AllocateZeroPool (0x%x) for profile buffer - %r\n", (UINTN) ProfileSize, Status);
+    return Status;
   }
 
   CommHeader = (EFI_SMM_COMMUNICATE_HEADER *) &CommBuffer[0];
   CopyMem (&CommHeader->HeaderGuid, &gEdkiiMemoryProfileGuid, sizeof(gEdkiiMemoryProfileGuid));
   CommHeader->MessageLength = sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA);
 
-  CommGetProfileData = (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA *)&CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];
+  CommGetProfileData = (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA *) &CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];
   CommGetProfileData->Header.Command      = SMRAM_PROFILE_COMMAND_GET_PROFILE_DATA;
   CommGetProfileData->Header.DataLength   = sizeof (*CommGetProfileData);
   CommGetProfileData->Header.ReturnStatus = (UINT64)-1;
@@ -663,6 +675,8 @@ GetSmramProfileData (
   ASSERT_EFI_ERROR (Status);
 
   if (CommGetProfileData->Header.ReturnStatus != 0) {
+    FreePool ((VOID *) (UINTN) CommGetProfileData->ProfileBuffer);
+    FreePool (CommBuffer);
     Print (L"GetProfileData - 0x%x\n", CommGetProfileData->Header.ReturnStatus);
     return EFI_SUCCESS;
   }
@@ -674,6 +688,7 @@ GetSmramProfileData (
   Print (L"======= SmramProfile end =======\n\n\n");
 
   FreePool ((VOID *) (UINTN) CommGetProfileData->ProfileBuffer);
+  FreePool (CommBuffer);
 
   return EFI_SUCCESS;
 }
diff --git a/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c b/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
index 1602fdb8e1..dca0008f75 100644
--- a/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
+++ b/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
@@ -817,6 +817,8 @@ CoreUpdateProfileAllocate (
   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;
   EFI_MEMORY_TYPE                   ProfileMemoryIndex;
 
+  AllocInfoData = NULL;
+
   ContextData = GetMemoryProfileContext ();
   if (ContextData == NULL) {
     return FALSE;
diff --git a/MdeModulePkg/Core/PiSmmCore/SmramProfileRecord.c b/MdeModulePkg/Core/PiSmmCore/SmramProfileRecord.c
index 24ab1b1de4..e119fadb67 100644
--- a/MdeModulePkg/Core/PiSmmCore/SmramProfileRecord.c
+++ b/MdeModulePkg/Core/PiSmmCore/SmramProfileRecord.c
@@ -824,6 +824,8 @@ SmmCoreUpdateProfileAllocate (
   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;
   EFI_MEMORY_TYPE                   ProfileMemoryIndex;
 
+  AllocInfoData = NULL;
+
   ContextData = GetSmramProfileContext ();
   if (ContextData == NULL) {
     return FALSE;