check memory allocations for success.

[mirror_edk2.git] / ShellPkg / Library / UefiShellLevel2CommandsLib / Vol.c

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Vol.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Vol.c
index 53b7770dfe12073e9a2d05e1e5ac3f0d2ce43179..9300d5ec35a1f46cec9594b90ac5e7d92705d185 100644 (file)
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Vol.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Vol.c
@@ -30,7 +30,7 @@ EFIAPI
 HandleVol(\r
   IN CONST CHAR16  *Path,\r
   IN CONST BOOLEAN Delete,\r
-  IN CONST CHAR16  *Name\r
+  IN CONST CHAR16  *Name OPTIONAL\r
   )\r
 {\r
   EFI_STATUS            Status;\r
@@ -45,6 +45,7 @@ HandleVol(
   ShellStatus   = SHELL_SUCCESS;\r
 \r
   if (\r
+      Name != NULL && (\r
       StrStr(Name, L"%") != NULL ||\r
       StrStr(Name, L"^") != NULL ||\r
       StrStr(Name, L"*") != NULL ||\r
@@ -60,7 +61,7 @@ HandleVol(
       StrStr(Name, L">") != NULL ||\r
       StrStr(Name, L"?") != NULL ||\r
       StrStr(Name, L"/") != NULL ||\r
-      StrStr(Name, L" ") != NULL\r
+      StrStr(Name, L" ") != NULL )\r
       ){\r
     ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_PROBLEM), gShellLevel2HiiHandle, Name);\r
     return (SHELL_INVALID_PARAMETER);\r
@@ -112,14 +113,20 @@ HandleVol(
     Size2 = StrSize(SysInfo->VolumeLabel);\r
     if (Size1 > Size2) {\r
       SysInfo = ReallocatePool((UINTN)SysInfo->Size, (UINTN)SysInfo->Size + Size1 - Size2, SysInfo);\r
+      if (SysInfo == NULL) {\r
+        ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellLevel2HiiHandle);\r
+        ShellStatus = SHELL_OUT_OF_RESOURCES;\r
+      } \r
+    }\r
+    if (SysInfo != NULL) {\r
+      StrCpy ((CHAR16 *) SysInfo->VolumeLabel, Name);\r
+      SysInfo->Size = SIZE_OF_EFI_FILE_SYSTEM_INFO + Size1;\r
+      Status = EfiFpHandle->SetInfo(\r
+        EfiFpHandle,\r
+        &gEfiFileSystemInfoGuid,\r
+        (UINTN)SysInfo->Size,\r
+        SysInfo);\r
     }\r
-    StrCpy ((CHAR16 *) SysInfo->VolumeLabel, Name);\r
-    SysInfo->Size = SIZE_OF_EFI_FILE_SYSTEM_INFO + Size1;\r
-    Status = EfiFpHandle->SetInfo(\r
-      EfiFpHandle,\r
-      &gEfiFileSystemInfoGuid,\r
-      (UINTN)SysInfo->Size,\r
-      SysInfo);\r
   }  \r
 \r
   FreePool(SysInfo);\r
@@ -150,23 +157,27 @@ HandleVol(
   }\r
 \r
   gEfiShellProtocol->CloseFile(ShellFileHandle);\r
+  \r
+  ASSERT(SysInfo != NULL);\r
 \r
-  //\r
-  // print VolumeInfo table\r
-  //\r
-  ShellPrintHiiEx (\r
-    0,\r
-    gST->ConOut->Mode->CursorRow,\r
-    NULL,\r
-    STRING_TOKEN (STR_VOL_VOLINFO),\r
-    gShellLevel2HiiHandle,\r
-    SysInfo->VolumeLabel,\r
-    SysInfo->ReadOnly?L"r":L"rw",\r
-    SysInfo->VolumeSize,\r
-    SysInfo->FreeSpace,\r
-    SysInfo->BlockSize\r
-   );\r
-  SHELL_FREE_NON_NULL(SysInfo);\r
+  if (SysInfo != NULL) {\r
+    //\r
+    // print VolumeInfo table\r
+    //\r
+    ShellPrintHiiEx (\r
+      0,\r
+      gST->ConOut->Mode->CursorRow,\r
+      NULL,\r
+      STRING_TOKEN (STR_VOL_VOLINFO),\r
+      gShellLevel2HiiHandle,\r
+      SysInfo->VolumeLabel,\r
+      SysInfo->ReadOnly?L"r":L"rw",\r
+      SysInfo->VolumeSize,\r
+      SysInfo->FreeSpace,\r
+      SysInfo->BlockSize\r
+     );\r
+    SHELL_FREE_NON_NULL(SysInfo);\r
+  }\r
 \r
   return (ShellStatus);\r
 }\r
@@ -200,6 +211,7 @@ ShellCommandRunVol (
   CHAR16        *FullPath;\r
   CHAR16        *TempSpot;\r
   UINTN         Length;\r
+  CONST CHAR16  *NewName;\r
 \r
   Length              = 0;\r
   ProblemParam        = NULL;\r
@@ -266,20 +278,21 @@ ShellCommandRunVol (
         StrnCatGrow(&FullPath, &Length, PathName, 0);\r
         StrnCatGrow(&FullPath, &Length, L":\\", 0);\r
         DeleteMode = ShellCommandLineGetFlag(Package, L"-d");\r
+        NewName    = ShellCommandLineGetValue(Package, L"-n");\r
         if (DeleteMode && ShellCommandLineGetFlag(Package, L"-n")) {\r
           ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_PARAM_CON), gShellLevel2HiiHandle);\r
           ShellStatus = SHELL_INVALID_PARAMETER;\r
-        } else if (ShellCommandLineGetFlag(Package, L"-n") && ShellCommandLineGetValue(Package, L"-n") == NULL) {\r
+        } else if (ShellCommandLineGetFlag(Package, L"-n") && NewName == NULL) {\r
           ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_VALUE), gShellLevel2HiiHandle, L"-n");\r
           ShellStatus = SHELL_INVALID_PARAMETER;\r
-        } else if (ShellCommandLineGetValue(Package, L"-n") != NULL && StrLen(ShellCommandLineGetValue(Package, L"-n")) > 11) {\r
+        } else if (NewName != NULL && StrLen(NewName) > 11) {\r
           ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_PROBLEM_VAL), gShellLevel2HiiHandle, L"-n");\r
           ShellStatus = SHELL_INVALID_PARAMETER;\r
         } else if (ShellStatus == SHELL_SUCCESS) {\r
           ShellStatus = HandleVol(\r
             FullPath,\r
             DeleteMode,\r
-            ShellCommandLineGetValue(Package, L"-n")\r
+            NewName\r
            );\r
         }\r
       }\r