ShellPkg/Edit: 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>

diff --git a/ShellPkg/Library/UefiShellDebug1CommandsLib/Edit/FileBuffer.c b/ShellPkg/Library/UefiShellDebug1CommandsLib/Edit/FileBuffer.c
index acd8512ff3e876e4318c4f50c5ff79092932e577..efe0df0253c0c8e4d8cf05666f472a2ce5121291 100644 (file)
--- a/ShellPkg/Library/UefiShellDebug1CommandsLib/Edit/FileBuffer.c
+++ b/ShellPkg/Library/UefiShellDebug1CommandsLib/Edit/FileBuffer.c
@@ -1,7 +1,7 @@
 /** @file\r
   Implements filebuffer interface functions.\r
 \r
-  Copyright (c) 2005 - 2015, Intel Corporation. All rights reserved. <BR>\r
+  Copyright (c) 2005 - 2016, Intel Corporation. All rights reserved. <BR>\r
   This program and the accompanying materials\r
   are licensed and made available under the terms and conditions of the BSD License\r
   which accompanies this distribution.  The full text of the license may be found at\r
@@ -505,29 +505,28 @@ FileBufferPrintLine (
 \r
   BufLen = (MainEditor.ScreenSize.Column + 1) * sizeof (CHAR16);\r
   PrintLine = AllocatePool (BufLen);\r
-  ASSERT (PrintLine != NULL);\r
-\r
-  StrnCpyS (PrintLine, BufLen/sizeof(CHAR16), Buffer, MIN(Limit, MainEditor.ScreenSize.Column));\r
-  for (; Limit < MainEditor.ScreenSize.Column; Limit++) {\r
-    PrintLine[Limit] = L' ';\r
-  }\r
-\r
-  PrintLine[MainEditor.ScreenSize.Column] = CHAR_NULL;\r
-\r
-  PrintLine2 = AllocatePool (BufLen * 2);\r
-  ASSERT (PrintLine2 != NULL);\r
+  if (PrintLine != NULL) {\r
+    StrnCpyS (PrintLine, BufLen/sizeof(CHAR16), Buffer, MIN(Limit, MainEditor.ScreenSize.Column));\r
+    for (; Limit < MainEditor.ScreenSize.Column; Limit++) {\r
+      PrintLine[Limit] = L' ';\r
+    }\r
 \r
-  ShellCopySearchAndReplace(PrintLine, PrintLine2, BufLen * 2, L"%", L"^%", FALSE, FALSE);\r
+    PrintLine[MainEditor.ScreenSize.Column] = CHAR_NULL;\r
 \r
-  ShellPrintEx (\r
-    0,\r
-    (INT32)Row - 1,\r
-    L"%s",\r
-    PrintLine2\r
-    );\r
+    PrintLine2 = AllocatePool (BufLen * 2);\r
+    if (PrintLine2 != NULL) {\r
+      ShellCopySearchAndReplace(PrintLine, PrintLine2, BufLen * 2, L"%", L"^%", FALSE, FALSE);\r
 \r
-  FreePool (PrintLine);\r
-  FreePool (PrintLine2);\r
+      ShellPrintEx (\r
+        0,\r
+        (INT32)Row - 1,\r
+        L"%s",\r
+        PrintLine2\r
+        );\r
+      FreePool (PrintLine2);\r
+    }\r
+    FreePool (PrintLine);\r
+  }\r
 \r
   return EFI_SUCCESS;\r
 }\r