ShellPkg/Dp: Handle memory allocation failure

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Ruiyu Ni <ruiyu.ni@intel.com>
Reviewed-by: Jaben Carsey <jaben.carsey@intel.com>

1 files changed, 7 insertions, 1 deletions

diff --git a/ShellPkg/Library/UefiDpLib/Dp.c b/ShellPkg/Library/UefiDpLib/Dp.c
index 4bad3c2f72..75c7d11dc3 100644
--- a/ShellPkg/Library/UefiDpLib/Dp.c
+++ b/ShellPkg/Library/UefiDpLib/Dp.c
@@ -259,13 +259,19 @@ ShellCommandRunDp (
   CustomCumulativeToken = ShellCommandLineGetValue (ParamPackage, L"-c");

   if (CustomCumulativeToken != NULL) {

     CustomCumulativeData = AllocateZeroPool (sizeof (PERF_CUM_DATA));

-    ASSERT (CustomCumulativeData != NULL);

+    if (CustomCumulativeData == NULL) {

+      return SHELL_OUT_OF_RESOURCES;

+    }

     CustomCumulativeData->MinDur = 0;

     CustomCumulativeData->MaxDur = 0;

     CustomCumulativeData->Count  = 0;

     CustomCumulativeData->Duration = 0;

     NameSize = StrLen (CustomCumulativeToken) + 1;

     CustomCumulativeData->Name   = AllocateZeroPool (NameSize);

+    if (CustomCumulativeData->Name == NULL) {

+      FreePool (CustomCumulativeData);

+      return SHELL_OUT_OF_RESOURCES;

+    }

     UnicodeStrToAsciiStrS (CustomCumulativeToken, CustomCumulativeData->Name, NameSize);

   }