Verify memory allocations were successful.
[mirror_edk2.git] / ShellPkg / Library / UefiShellLevel2CommandsLib / Cd.c
index 78403c2..99680b4 100644 (file)
@@ -186,17 +186,27 @@ ShellCommandRunCd (
         // 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
-        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