check memory allocations for success.

[mirror_edk2.git] / ShellPkg / Library / UefiShellDebug1CommandsLib / SetVar.c

diff --git a/ShellPkg/Library/UefiShellDebug1CommandsLib/SetVar.c b/ShellPkg/Library/UefiShellDebug1CommandsLib/SetVar.c
index afccf7bd5821aa0a7740ae73225e5fb87387a1c5..34eb6af8ac80c77051fbcd3420b71f0999ecd694 100644 (file)
--- a/ShellPkg/Library/UefiShellDebug1CommandsLib/SetVar.c
+++ b/ShellPkg/Library/UefiShellDebug1CommandsLib/SetVar.c
@@ -205,15 +205,20 @@ ShellCommandRunSetVar (
           Data++;\r
           Data++;\r
           Buffer = AllocateZeroPool(StrSize(Data));\r
-          UnicodeSPrint(Buffer, StrSize(Data), L"%s", Data);\r
-          ((CHAR16*)Buffer)[StrLen(Buffer)-1] = CHAR_NULL;\r
-\r
-          Status = gRT->SetVariable((CHAR16*)VariableName, &Guid, Attributes, StrSize(Buffer)-sizeof(CHAR16), Buffer);\r
-          if (EFI_ERROR(Status)) {\r
-            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_SETVAR_ERROR_SET), gShellDebug1HiiHandle, &Guid, VariableName, Status);\r
-            ShellStatus = SHELL_ACCESS_DENIED;\r
+          if (Buffer == NULL) {\r
+            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellDebug1HiiHandle);\r
+            ShellStatus = SHELL_OUT_OF_RESOURCES;\r
           } else {\r
-            ASSERT(ShellStatus == SHELL_SUCCESS);\r
+            UnicodeSPrint(Buffer, StrSize(Data), L"%s", Data);\r
+            ((CHAR16*)Buffer)[StrLen(Buffer)-1] = CHAR_NULL;\r
+\r
+            Status = gRT->SetVariable((CHAR16*)VariableName, &Guid, Attributes, StrSize(Buffer)-sizeof(CHAR16), Buffer);\r
+            if (EFI_ERROR(Status)) {\r
+              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_SETVAR_ERROR_SET), gShellDebug1HiiHandle, &Guid, VariableName, Status);\r
+              ShellStatus = SHELL_ACCESS_DENIED;\r
+            } else {\r
+              ASSERT(ShellStatus == SHELL_SUCCESS);\r
+            }\r
           }\r
         } else if (StrnCmp(Data, L"--", 2) == 0) {\r
           //\r