ShellPkg/UefiShellBcfgCommandLib: 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/UefiShellBcfgCommandLib/UefiShellBcfgCommandLib.c b/ShellPkg/Library/UefiShellBcfgCommandLib/UefiShellBcfgCommandLib.c
index bf35a35b985d00e32af4e4aa5d1cb285c7e9a014..9baeecc339ed3807cb105778d2af3b46ec9303c6 100644 (file)
--- a/ShellPkg/Library/UefiShellBcfgCommandLib/UefiShellBcfgCommandLib.c
+++ b/ShellPkg/Library/UefiShellBcfgCommandLib/UefiShellBcfgCommandLib.c
@@ -2,7 +2,7 @@
   Main file for BCFG command.\r
 \r
   (C) Copyright 2014-2015 Hewlett-Packard Development Company, L.P.<BR>\r
-  Copyright (c) 2010 - 2015, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2010 - 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
@@ -553,33 +553,34 @@ BcfgAdd(
     if (EFI_ERROR(Status)) {\r
       ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_BCFG_SET_VAR_FAIL), gShellBcfgHiiHandle, L"bcfg", OptionStr);  \r
     } else {\r
-      NewOrder = AllocateZeroPool((OrderCount+1)*sizeof(NewOrder[0]));\r
-      ASSERT(NewOrder != NULL);\r
-      CopyMem(NewOrder, CurrentOrder, (OrderCount)*sizeof(NewOrder[0]));\r
+      NewOrder = AllocateZeroPool ((OrderCount + 1) * sizeof (NewOrder[0]));\r
+      if (NewOrder != NULL) {\r
+        CopyMem (NewOrder, CurrentOrder, (OrderCount) * sizeof (NewOrder[0]));\r
 \r
-      //\r
-      // Insert target into order list\r
-      //\r
-      for (Index=OrderCount; Index > Position; Index--) {\r
-        NewOrder[Index] = NewOrder[Index-1];\r
-      }\r
+        //\r
+        // Insert target into order list\r
+        //\r
+        for (Index = OrderCount; Index > Position; Index--) {\r
+          NewOrder[Index] = NewOrder[Index - 1];\r
+        }\r
 \r
-      NewOrder[Position] = (UINT16) TargetLocation;\r
-      Status = gRT->SetVariable (\r
-        Target == BcfgTargetBootOrder?L"BootOrder":L"DriverOrder",\r
-        &gEfiGlobalVariableGuid,\r
-        EFI_VARIABLE_NON_VOLATILE|EFI_VARIABLE_BOOTSERVICE_ACCESS|EFI_VARIABLE_RUNTIME_ACCESS,\r
-        (OrderCount+1) * sizeof(UINT16),\r
-        NewOrder\r
-       );\r
+        NewOrder[Position] = (UINT16) TargetLocation;\r
+        Status = gRT->SetVariable (\r
+          Target == BcfgTargetBootOrder ? L"BootOrder" : L"DriverOrder",\r
+          &gEfiGlobalVariableGuid,\r
+          EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,\r
+          (OrderCount + 1) * sizeof (UINT16),\r
+          NewOrder\r
+        );\r
 \r
-      FreePool(NewOrder);\r
+        FreePool (NewOrder);\r
 \r
-      if (EFI_ERROR(Status)) {\r
-        ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_BCFG_WRITE_FAIL), gShellBcfgHiiHandle, L"bcfg", Target == BcfgTargetBootOrder?L"BootOrder":L"DriverOrder");  \r
-        ShellStatus = SHELL_INVALID_PARAMETER;\r
-      } else {\r
-        Print (L"bcfg: Add %s as %x\n", OptionStr, Position);\r
+        if (EFI_ERROR (Status)) {\r
+          ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_BCFG_WRITE_FAIL), gShellBcfgHiiHandle, L"bcfg", Target == BcfgTargetBootOrder ? L"BootOrder" : L"DriverOrder");\r
+          ShellStatus = SHELL_INVALID_PARAMETER;\r
+        } else {\r
+          Print (L"bcfg: Add %s as %x\n", OptionStr, Position);\r
+        }\r
       }\r
     }\r
   }\r