ShellPkg/Mv: Handle memory allocation failure

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Ruiyu Ni <ruiyu.ni@intel.com>
Reviewed-by: Jaben Carsey <jaben.carsey@intel.com>

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c
index d02a6ae5f5c0560c7bf987b46f26be12bb8a9534..f93772c6f8d144bc869ecf307aad4f0fe231d8f2 100644 (file)
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c
@@ -736,11 +736,15 @@ ShellCommandRunMv (
             //\r
             CwdSize = StrSize(ShellGetCurrentDir(NULL)) + sizeof(CHAR16);\r
             Cwd = AllocateZeroPool(CwdSize);\r
-            ASSERT (Cwd != NULL);\r
-            StrCpyS(Cwd, CwdSize/sizeof(CHAR16), ShellGetCurrentDir(NULL));\r
-            StrCatS(Cwd, CwdSize/sizeof(CHAR16), L"\\");\r
-            ShellStatus = ValidateAndMoveFiles(FileList, &Response, Cwd);\r
-            FreePool(Cwd);\r
+            if (Cwd == NULL) {\r
+              ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_OUT_MEM), gShellLevel2HiiHandle, L"mv");\r
+              ShellStatus = SHELL_OUT_OF_RESOURCES;\r
+            } else {\r
+              StrCpyS (Cwd, CwdSize / sizeof (CHAR16), ShellGetCurrentDir (NULL));\r
+              StrCatS (Cwd, CwdSize / sizeof (CHAR16), L"\\");\r
+              ShellStatus = ValidateAndMoveFiles (FileList, &Response, Cwd);\r
+              FreePool (Cwd);\r
+            }\r
           }\r
         }\r
 \r