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 f8c6dd2e30d3c2383a369adaf7f92616893be9e5..34eb6af8ac80c77051fbcd3420b71f0999ecd694 100644 (file)
--- a/ShellPkg/Library/UefiShellDebug1CommandsLib/SetVar.c
+++ b/ShellPkg/Library/UefiShellDebug1CommandsLib/SetVar.c
@@ -1,7 +1,7 @@
 /** @file\r
   Main file for SetVar shell Debug1 function.\r
 \r
-  Copyright (c) 2010, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2010 - 2011, 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
@@ -22,6 +22,12 @@ STATIC CONST SHELL_PARAM_ITEM ParamList[] = {
   {NULL, TypeMax}\r
   };\r
 \r
+/**\r
+  Function for 'setvar' command.\r
+\r
+  @param[in] ImageHandle  Handle to the Image (NULL if Internal).\r
+  @param[in] SystemTable  Pointer to the System Table (NULL if Internal).\r
+**/\r
 SHELL_STATUS\r
 EFIAPI\r
 ShellCommandRunSetVar (\r
@@ -38,6 +44,7 @@ ShellCommandRunSetVar (
   EFI_GUID            Guid;\r
   CONST CHAR16        *StringGuid;\r
   UINT32              Attributes;\r
+  UINT32              Attributes2;\r
   VOID                *Buffer;\r
   UINTN               Size;\r
   UINTN               LoopVar;\r
@@ -88,7 +95,8 @@ ShellCommandRunSetVar (
         StringGuid = ShellCommandLineGetValue(Package, L"-guid");\r
         Status = ConvertStringToGuid(StringGuid, &Guid);\r
         if (EFI_ERROR(Status)) {\r
-          ShellStatus = SHELL_NOT_FOUND;\r
+          ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_PROBLEM), gShellDebug1HiiHandle, StringGuid);\r
+          ShellStatus = SHELL_INVALID_PARAMETER;\r
         }\r
       }\r
       if (Data == NULL) {\r
@@ -97,7 +105,7 @@ ShellCommandRunSetVar (
         //\r
         Status = gRT->GetVariable((CHAR16*)VariableName, &Guid, &Attributes, &Size, Buffer);\r
         if (Status == EFI_BUFFER_TOO_SMALL) {\r
-          Buffer = AllocatePool(Size);\r
+          Buffer = AllocateZeroPool(Size);\r
           Status = gRT->GetVariable((CHAR16*)VariableName, &Guid, &Attributes, &Size, Buffer);\r
         }\r
         if (!EFI_ERROR(Status)&& Buffer != NULL) {\r
@@ -138,31 +146,52 @@ ShellCommandRunSetVar (
           Attributes |= EFI_VARIABLE_NON_VOLATILE;\r
         }\r
         if (ShellIsHexOrDecimalNumber(Data, TRUE, FALSE)) {\r
-          //\r
-          // arbitrary buffer\r
-          //\r
-          Buffer = AllocateZeroPool((StrLen(Data) / 2));\r
-          for (LoopVar = 0 ; LoopVar < (StrLen(Data) / 2) ; LoopVar++) {\r
-            ((UINT8*)Buffer)[LoopVar] = (UINT8)(HexCharToUintn(Data[LoopVar*2]) * 16);\r
-            ((UINT8*)Buffer)[LoopVar] = (UINT8)(((UINT8*)Buffer)[LoopVar] + HexCharToUintn(Data[LoopVar*2+1]));\r
-          }\r
-          Status = gRT->SetVariable((CHAR16*)VariableName, &Guid, Attributes, StrLen(Data) / 2, 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 (StrLen(Data) % 2 != 0) {\r
+            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_PROBLEM_VAL), gShellDebug1HiiHandle, Data);\r
+            ShellStatus = SHELL_INVALID_PARAMETER;\r
           } else {\r
-            ASSERT(ShellStatus == SHELL_SUCCESS);\r
+            //\r
+            // arbitrary buffer\r
+            //\r
+            Buffer = AllocateZeroPool((StrLen(Data) / 2));\r
+            if (Buffer == NULL) {\r
+              Status = EFI_OUT_OF_RESOURCES;\r
+            } else {\r
+              for (LoopVar = 0 ; LoopVar < (StrLen(Data) / 2) ; LoopVar++) {\r
+                ((UINT8*)Buffer)[LoopVar] = (UINT8)(HexCharToUintn(Data[LoopVar*2]) * 16);\r
+                ((UINT8*)Buffer)[LoopVar] = (UINT8)(((UINT8*)Buffer)[LoopVar] + HexCharToUintn(Data[LoopVar*2+1]));\r
+              }\r
+              Status = gRT->SetVariable((CHAR16*)VariableName, &Guid, Attributes, StrLen(Data) / 2, Buffer);\r
+            }\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"\"", 1) == 0) {\r
+          Size = 0;\r
+          Attributes2 = 0;\r
+          Status = gRT->GetVariable((CHAR16*)VariableName, &Guid, &Attributes2, &Size, Buffer);\r
+          if (Status == EFI_BUFFER_TOO_SMALL) {\r
+            Buffer = AllocateZeroPool(Size);\r
+            Status = gRT->GetVariable((CHAR16*)VariableName, &Guid, &Attributes2, &Size, Buffer);\r
+            FreePool(Buffer);\r
+            Attributes = Attributes2;\r
+          }          \r
           //\r
           // ascii text\r
           //\r
           Data++;\r
           Buffer = AllocateZeroPool(StrSize(Data) / 2);\r
-          AsciiSPrint(Buffer, StrSize(Data) / 2, "%s", Data);\r
-          ((CHAR8*)Buffer)[AsciiStrLen(Buffer)-1] = CHAR_NULL;\r
-\r
-          Status = gRT->SetVariable((CHAR16*)VariableName, &Guid, Attributes, AsciiStrSize(Buffer)-sizeof(CHAR8), Buffer);\r
+          if (Buffer == NULL) {\r
+            Status = EFI_OUT_OF_RESOURCES;\r
+          } else {\r
+            AsciiSPrint(Buffer, StrSize(Data) / 2, "%s", Data);\r
+            ((CHAR8*)Buffer)[AsciiStrLen(Buffer)-1] = CHAR_NULL;\r
+            Status = gRT->SetVariable((CHAR16*)VariableName, &Guid, Attributes, AsciiStrSize(Buffer)-sizeof(CHAR8), Buffer);\r
+          }\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
@@ -176,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