ShellPkg/Cp: Handle memory allocation failure

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

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Cd.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Cd.c
index c0d3c920f492fb70f332c59a26a14b2022c6b37f..2e51b4cb4dd152e18e6daeda11e0fe83a65ff488 100644 (file)
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Cd.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Cd.c
@@ -1,7 +1,8 @@
 /** @file\r
   Main file for attrib shell level 2 function.\r
 \r
-  Copyright (c) 2009 - 2013, Intel Corporation. All rights reserved.<BR>\r
+  (C) Copyright 2015 Hewlett-Packard Development Company, L.P.<BR>\r
+  Copyright (c) 2009 - 2016, 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
@@ -30,14 +31,17 @@ ShellCommandRunCd (
   EFI_STATUS        Status;\r
   LIST_ENTRY        *Package;\r
   CONST CHAR16      *Directory;\r
+  CHAR16            *Cwd;\r
   CHAR16            *Path;\r
   CHAR16            *Drive;\r
+  UINTN             CwdSize;\r
   UINTN             DriveSize;\r
   CHAR16            *ProblemParam;\r
   SHELL_STATUS      ShellStatus;\r
   SHELL_FILE_HANDLE Handle;\r
   CONST CHAR16      *Param1;\r
   CHAR16            *Param1Copy;\r
+  CHAR16*           Walker;\r
 \r
   ProblemParam = NULL;\r
   ShellStatus = SHELL_SUCCESS;\r
@@ -59,7 +63,7 @@ ShellCommandRunCd (
   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), gShellLevel2HiiHandle, ProblemParam);\r
+      ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_PROBLEM), gShellLevel2HiiHandle, L"cd", ProblemParam);  \r
       FreePool(ProblemParam);\r
       ShellStatus = SHELL_INVALID_PARAMETER;\r
     } else {\r
@@ -73,7 +77,7 @@ ShellCommandRunCd (
   if (ShellCommandLineGetFlag(Package, L"-?")) {\r
     ASSERT(FALSE);\r
   } else if (ShellCommandLineGetRawValue(Package, 2) != NULL) {\r
-    ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_TOO_MANY), gShellLevel2HiiHandle);\r
+    ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_TOO_MANY), gShellLevel2HiiHandle, L"cd");  \r
     ShellStatus = SHELL_INVALID_PARAMETER;\r
   } else {\r
     //\r
@@ -91,11 +95,17 @@ ShellCommandRunCd (
       if (Directory != NULL) {\r
         ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_PRINT), gShellLevel2HiiHandle, Directory);\r
       } else {\r
-        ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_CWD), gShellLevel2HiiHandle);\r
+        ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_CWD), gShellLevel2HiiHandle, L"cd");  \r
         ShellStatus = SHELL_NOT_FOUND;\r
       }\r
     } else {\r
       Param1Copy = CatSPrint(NULL, L"%s", Param1, NULL);\r
+      for (Walker = Param1Copy; Walker != NULL && *Walker != CHAR_NULL ; Walker++) {\r
+        if (*Walker == L'\"') {\r
+          CopyMem(Walker, Walker+1, StrSize(Walker) - sizeof(Walker[0]));\r
+        }\r
+      }\r
+      \r
       if (Param1Copy != NULL) {\r
         Param1Copy = PathCleanUpDirectories(Param1Copy);\r
       }\r
@@ -110,11 +120,21 @@ ShellCommandRunCd (
           //\r
           Directory = ShellGetCurrentDir(NULL);\r
           if (Directory == NULL) {\r
-            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_CWD), gShellLevel2HiiHandle);\r
+            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_CWD), gShellLevel2HiiHandle, L"cd");  \r
             ShellStatus = SHELL_NOT_FOUND;\r
           } else {\r
-            Drive = GetFullyQualifiedPath(Directory);\r
-            PathRemoveLastItem(Drive);\r
+            CwdSize = StrSize(Directory) + sizeof(CHAR16);\r
+            Cwd = AllocateZeroPool(CwdSize);\r
+            if (Cwd == NULL) {\r
+              ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellLevel2HiiHandle, L"cd");\r
+              ShellStatus = SHELL_OUT_OF_RESOURCES;\r
+            } else {\r
+              StrCpyS (Cwd, StrSize (Directory) / sizeof (CHAR16) + 1, Directory);\r
+              StrCatS (Cwd, StrSize (Directory) / sizeof (CHAR16) + 1, L"\\");\r
+              Drive = GetFullyQualifiedPath (Cwd);\r
+              PathRemoveLastItem (Drive);\r
+              FreePool (Cwd);\r
+            }\r
           }\r
           if (ShellStatus == SHELL_SUCCESS && Drive != NULL) {\r
             //\r
@@ -122,7 +142,7 @@ ShellCommandRunCd (
             //\r
             Status = gEfiShellProtocol->SetCurDir(NULL, Drive);\r
             if (Status == EFI_NOT_FOUND) {\r
-              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle);\r
+              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle, L"cd");  \r
               ShellStatus = SHELL_NOT_FOUND;\r
             }\r
           }\r
@@ -132,11 +152,21 @@ ShellCommandRunCd (
           //\r
           Directory = ShellGetCurrentDir(NULL);\r
           if (Directory == NULL) {\r
-            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_CWD), gShellLevel2HiiHandle);\r
+            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_CWD), gShellLevel2HiiHandle, L"cd");  \r
             ShellStatus = SHELL_NOT_FOUND;\r
           } else {\r
-            Drive = GetFullyQualifiedPath(Directory);\r
-            while (PathRemoveLastItem(Drive)) ;\r
+            CwdSize = StrSize(Directory) + sizeof(CHAR16);\r
+            Cwd = AllocateZeroPool(CwdSize);\r
+            if (Cwd == NULL) {\r
+              ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellLevel2HiiHandle, L"cd");\r
+              ShellStatus = SHELL_OUT_OF_RESOURCES;\r
+            } else {\r
+              StrCpyS (Cwd, StrSize (Directory) / sizeof (CHAR16) + 1, Directory);\r
+              StrCatS (Cwd, StrSize (Directory) / sizeof (CHAR16) + 1, L"\\");\r
+              Drive = GetFullyQualifiedPath (Cwd);\r
+              while (PathRemoveLastItem (Drive));\r
+              FreePool (Cwd);\r
+            }\r
           }\r
           if (ShellStatus == SHELL_SUCCESS && Drive != NULL) {\r
             //\r
@@ -144,7 +174,7 @@ ShellCommandRunCd (
             //\r
             Status = gEfiShellProtocol->SetCurDir(NULL, Drive);\r
             if (Status == EFI_NOT_FOUND) {\r
-              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle);\r
+              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle, L"cd");  \r
               ShellStatus = SHELL_NOT_FOUND;\r
             }\r
           }\r
@@ -153,11 +183,12 @@ ShellCommandRunCd (
           // change directory without a drive identifier\r
           //\r
           if (ShellGetCurrentDir(NULL) == NULL) {\r
-            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_CWD), gShellLevel2HiiHandle);\r
+            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_CWD), gShellLevel2HiiHandle, L"cd");  \r
             ShellStatus = SHELL_NOT_FOUND;\r
           } else {\r
             ASSERT((Drive == NULL && DriveSize == 0) || (Drive != NULL));\r
             Drive = StrnCatGrow(&Drive, &DriveSize, ShellGetCurrentDir(NULL), 0);\r
+            Drive = StrnCatGrow(&Drive, &DriveSize, L"\\", 0);\r
             if (*Param1Copy == L'\\') {\r
               while (PathRemoveLastItem(Drive)) ;\r
               Drive = StrnCatGrow(&Drive, &DriveSize, Param1Copy+1, 0);\r
@@ -169,10 +200,10 @@ ShellCommandRunCd (
             //\r
             Status = gEfiShellProtocol->OpenFileByName(Drive, &Handle, EFI_FILE_MODE_READ);\r
             if (EFI_ERROR(Status)) {\r
-              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_DIR_NF), gShellLevel2HiiHandle, Drive);\r
+              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_DIR_NF), gShellLevel2HiiHandle, L"cd", Drive);  \r
               ShellStatus = SHELL_NOT_FOUND;\r
             } else if (EFI_ERROR(FileHandleIsDirectory(Handle))) {\r
-              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NOT_DIR), gShellLevel2HiiHandle, Drive);\r
+              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NOT_DIR), gShellLevel2HiiHandle, L"cd", Drive);  \r
               ShellStatus = SHELL_NOT_FOUND;\r
             }\r
             if (ShellStatus == SHELL_SUCCESS && Drive != NULL) {\r
@@ -181,7 +212,7 @@ ShellCommandRunCd (
               //\r
               Status = gEfiShellProtocol->SetCurDir(NULL, Drive);\r
               if (Status == EFI_NOT_FOUND) {\r
-                ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle);\r
+                ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle, L"cd");  \r
                 ShellStatus = SHELL_NOT_FOUND;\r
               }\r
             }\r
@@ -194,24 +225,23 @@ ShellCommandRunCd (
           //\r
           // change directory with a drive letter\r
           //\r
-          Drive = AllocateZeroPool(StrSize(Param1Copy));\r
+          Drive = AllocateCopyPool(StrSize(Param1Copy), Param1Copy);\r
           if (Drive == NULL) {\r
-            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_MEM), gShellLevel2HiiHandle);\r
+            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_MEM), gShellLevel2HiiHandle, L"cd");  \r
             ShellStatus = SHELL_OUT_OF_RESOURCES;\r
           } else {\r
-            Drive = StrCpy(Drive, Param1Copy);\r
             Path = StrStr(Drive, L":");\r
             ASSERT(Path != NULL);\r
             if (EFI_ERROR(ShellIsDirectory(Param1Copy))) {\r
-              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NOT_DIR), gShellLevel2HiiHandle, Param1Copy);\r
+              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NOT_DIR), gShellLevel2HiiHandle, L"cd", Param1Copy);  \r
               ShellStatus = SHELL_NOT_FOUND;\r
             } else if (*(Path+1) == CHAR_NULL) {\r
-              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle);\r
+              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle, L"cd");  \r
               ShellStatus = SHELL_NOT_FOUND;\r
             } else {\r
               *(Path+1) = CHAR_NULL;\r
               if (Path == Drive + StrLen(Drive)) {\r
-                ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle);\r
+                ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle, L"cd");  \r
                 ShellStatus = SHELL_NOT_FOUND;\r
               } else {\r
                 Status = gEfiShellProtocol->SetCurDir(Drive, Path+2);\r
@@ -219,10 +249,10 @@ ShellCommandRunCd (
               }\r
             }\r
             if (Status == EFI_NOT_FOUND) {\r
-              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle);\r
+              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle, L"cd");  \r
               Status = SHELL_NOT_FOUND;\r
             } else if (EFI_ERROR(Status)) {\r
-              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_DIR_NF), gShellLevel2HiiHandle, Param1Copy);\r
+              ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_DIR_NF), gShellLevel2HiiHandle, L"cd", Param1Copy);  \r
               Status = SHELL_NOT_FOUND;\r
             }\r
           }\r