Verify memory allocations were successful.

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

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Cd.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Cd.c
index 78403c22748bae815907aed756d1dd99a02a9047..99680b49d9abf8f318e3e168f9e7161407ebd61a 100644 (file)
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Cd.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Cd.c
@@ -186,17 +186,27 @@ ShellCommandRunCd (
         // change directory on other drive letter\r
         //\r
         Drive = AllocateZeroPool(StrSize(Param1));\r
         // change directory on other drive letter\r
         //\r
         Drive = AllocateZeroPool(StrSize(Param1));\r

-        Drive = StrCpy(Drive, Param1);\r
-        Path = StrStr(Drive, L":");\r
-        *(++Path) = CHAR_NULL;\r
-        Status = gEfiShellProtocol->SetCurDir(Drive, ++Path);\r
+        if (Drive == NULL) {\r
+          ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_MEM), gShellLevel2HiiHandle);\r
+          ShellStatus = SHELL_OUT_OF_RESOURCES;\r
+        } else {\r
+          Drive = StrCpy(Drive, Param1);\r
+          Path = StrStr(Drive, L":");\r
+          *(++Path) = CHAR_NULL;\r
+          if (Path == Drive + StrLen(Drive)) {\r
+            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle);\r
+            ShellStatus = SHELL_NOT_FOUND;\r
+          } else {\r
+            Status = gEfiShellProtocol->SetCurDir(Drive, ++Path);\r
+          }\r

 \r
 \r

-        if (Status == EFI_NOT_FOUND) {\r
-          ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle);\r
-          Status = SHELL_NOT_FOUND;\r
-        } else if (EFI_ERROR(Status)) {\r
-          ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_DIR_NF), gShellLevel2HiiHandle, Param1);\r
-          Status = SHELL_NOT_FOUND;\r
+          if (Status == EFI_NOT_FOUND) {\r
+            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle);\r
+            Status = SHELL_NOT_FOUND;\r
+          } else if (EFI_ERROR(Status)) {\r
+            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_DIR_NF), gShellLevel2HiiHandle, Param1);\r
+            Status = SHELL_NOT_FOUND;\r
+          }\r

         }\r
       }\r
     }\r
         }\r
       }\r
     }\r