From d758f80971a7447328d67dd5d4dd2a94dabe1656 Mon Sep 17 00:00:00 2001 From: Ruiyu Ni Date: Tue, 12 Jul 2016 13:42:59 +0800 Subject: [PATCH] ShellPkg/Dp: Handle memory allocation failure Contributed-under: TianoCore Contribution Agreement 1.0 Signed-off-by: Ruiyu Ni Reviewed-by: Jaben Carsey --- ShellPkg/Library/UefiDpLib/Dp.c | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) 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); } -- 2.39.2