ShellPkg/for: Fix potential null pointer deference

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

diff --git a/ShellPkg/Library/UefiShellDebug1CommandsLib/SetSize.c b/ShellPkg/Library/UefiShellDebug1CommandsLib/SetSize.c
index 30a382fb049ba13a19b2b45dac3f838e82a5b47c..a33aec12a515e6a6f3d6aec85c79b776cb7d3dbf 100644 (file)
--- a/ShellPkg/Library/UefiShellDebug1CommandsLib/SetSize.c
+++ b/ShellPkg/Library/UefiShellDebug1CommandsLib/SetSize.c
@@ -1,7 +1,8 @@
 /** @file\r
   Main file for SetSize shell Debug1 function.\r
 \r
-  Copyright (c) 2010, Intel Corporation. All rights reserved.<BR>\r
+  (C) Copyright 2015 Hewlett-Packard Development Company, L.P.<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
@@ -14,6 +15,12 @@
 \r
 #include "UefiShellDebug1CommandsLib.h"\r
 \r
+/**\r
+  Function for 'setsize' 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
 ShellCommandRunSetSize (\r
@@ -48,36 +55,49 @@ ShellCommandRunSetSize (
   Status = ShellCommandLineParse (EmptyParamList, &Package, &ProblemParam, TRUE);\r
   if (EFI_ERROR(Status)) {\r
     if (Status == EFI_VOLUME_CORRUPTED && ProblemParam != NULL) {\r
-      ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_PROBLEM), gShellDebug1HiiHandle, ProblemParam);\r
+      ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_PROBLEM), gShellDebug1HiiHandle, L"setsize", ProblemParam);  \r
       FreePool(ProblemParam);\r
       ShellStatus = SHELL_INVALID_PARAMETER;\r
     } else {\r
       ASSERT(FALSE);\r
     }\r
   } else {\r
-    Temp1 = ShellCommandLineGetRawValue(Package, 1);\r
-    if (Temp1 == NULL) {\r
-      ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_SIZE_NOT_SPEC), gShellDebug1HiiHandle);\r
+    if (ShellCommandLineGetCount(Package) < 3) {\r
+      ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_TOO_FEW), gShellDebug1HiiHandle, L"setsize");  \r
       ShellStatus = SHELL_INVALID_PARAMETER;\r
       NewSize = 0;\r
     } else {\r
-      NewSize = ShellStrToUintn(Temp1);\r
+      Temp1 = ShellCommandLineGetRawValue(Package, 1);\r
+      if (!ShellIsHexOrDecimalNumber(Temp1, FALSE, FALSE)) {\r
+        ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_SIZE_NOT_SPEC), gShellDebug1HiiHandle, L"setsize");  \r
+        ShellStatus = SHELL_INVALID_PARAMETER;\r
+        NewSize = 0;\r
+      } else {\r
+        NewSize = ShellStrToUintn(Temp1);\r
+      }\r
     }\r
     for (LoopVar = 2 ; LoopVar < ShellCommandLineGetCount(Package) && ShellStatus == SHELL_SUCCESS ; LoopVar++) {\r
-      Status = ShellOpenFileByName(ShellCommandLineGetRawValue(Package, LoopVar), &FileHandle, EFI_FILE_MODE_READ|EFI_FILE_MODE_WRITE|EFI_FILE_MODE_CREATE, 0);\r
+      Status = ShellOpenFileByName(ShellCommandLineGetRawValue(Package, LoopVar), &FileHandle, EFI_FILE_MODE_READ|EFI_FILE_MODE_WRITE, 0);\r
+      if (EFI_ERROR(Status)) {\r
+        Status = ShellOpenFileByName(ShellCommandLineGetRawValue(Package, LoopVar), &FileHandle, EFI_FILE_MODE_READ|EFI_FILE_MODE_WRITE|EFI_FILE_MODE_CREATE, 0);\r
+      }\r
       if (EFI_ERROR(Status) && LoopVar == 2) {\r
-        ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_FILE_NOT_SPEC), gShellDebug1HiiHandle);\r
+        ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_FILE_NOT_SPEC), gShellDebug1HiiHandle, L"setsize");  \r
         ShellStatus = SHELL_INVALID_PARAMETER;\r
       } else if (EFI_ERROR(Status)) {\r
+        ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_FILE_OPEN_FAIL), gShellDebug1HiiHandle, L"setsize", ShellCommandLineGetRawValue(Package, LoopVar));  \r
+        ShellStatus = SHELL_INVALID_PARAMETER;\r
         break;\r
       } else {\r
         Status = FileHandleSetSize(FileHandle, NewSize);\r
         if (Status == EFI_VOLUME_FULL) {\r
-          ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_VOLUME_FULL), gShellDebug1HiiHandle);\r
+          ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_VOLUME_FULL), gShellDebug1HiiHandle, L"setsize");  \r
           ShellStatus = SHELL_VOLUME_FULL;\r
         } else if (EFI_ERROR(Status)) {\r
-          ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_SET_SIZE_FAIL), gShellDebug1HiiHandle, ShellCommandLineGetRawValue(Package, LoopVar), Status);\r
+          ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_SET_SIZE_FAIL), gShellDebug1HiiHandle, L"setsize", ShellCommandLineGetRawValue(Package, LoopVar));  \r
           ShellStatus = SHELL_INVALID_PARAMETER;\r
+        } else {\r
+          ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_SET_SIZE_DONE), gShellDebug1HiiHandle, ShellCommandLineGetRawValue(Package, LoopVar));\r
         }\r
         ShellCloseFile(&FileHandle);\r
       }\r