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 51cf33f1ba788aaa8e137affd423cf87809aea3e..9300d5ec35a1f46cec9594b90ac5e7d92705d185 100644 (file)
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Vol.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Vol.c
@@ -113,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
@@ -154,22 +160,24 @@ HandleVol(
   \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