Verify memory allocations were successful.

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

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c
index adb6f9975902d4e53ae4e8f0e211bc756e29b300..a62878175701e66c343a6b6d451b59f122f72dba 100644 (file)
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c
@@ -152,6 +152,10 @@ GetDestinationLocation(
       NewSize = StrSize(Cwd);\r
       NewSize += StrSize(DestDir);\r
       DestPath = AllocateZeroPool(NewSize);\r
+      if (DestPath == NULL) {\r
+        ShellCloseFileMetaArg(&DestList);\r
+        return (SHELL_OUT_OF_RESOURCES);\r
+      }\r
       StrCpy(DestPath, Cwd);\r
       if (DestPath[StrLen(DestPath)-1] != L'\\' && DestDir[0] != L'\\') {\r
         StrCat(DestPath, L"\\");\r
@@ -162,6 +166,10 @@ GetDestinationLocation(
     } else {\r
       ASSERT(DestPath == NULL);\r
       DestPath = StrnCatGrow(&DestPath, NULL, DestDir, 0);\r
+      if (DestPath == NULL) {\r
+        ShellCloseFileMetaArg(&DestList);\r
+        return (SHELL_OUT_OF_RESOURCES);\r
+      }\r
     }\r
   } else {\r
     Node = (EFI_SHELL_FILE_INFO*)GetFirstNode(&DestList->Link);\r
@@ -175,6 +183,10 @@ GetDestinationLocation(
     }\r
     if (ShellIsDirectory(Node->FullName)==EFI_SUCCESS) {\r
       DestPath = AllocateZeroPool(StrSize(Node->FullName)+sizeof(CHAR16));\r
+      if (DestPath == NULL) {\r
+        ShellCloseFileMetaArg(&DestList);\r
+        return (SHELL_OUT_OF_RESOURCES);\r
+      }\r
       StrCpy(DestPath, Node->FullName);\r
       StrCat(DestPath, L"\\");\r
     } else {\r
@@ -287,62 +299,66 @@ ValidateAndMoveFiles(
     NewSize = StrSize(DestPath);\r
     NewSize += StrSize(Node->FileName) + sizeof(EFI_FILE_INFO) + sizeof(CHAR16);\r
     NewFileInfo = AllocateZeroPool(NewSize);\r
-    ASSERT(NewFileInfo != NULL);\r
-    CopyMem(NewFileInfo, Node->Info, sizeof(EFI_FILE_INFO));\r
-    if (DestPath[0] != L'\\') {\r
-      StrCpy(NewFileInfo->FileName, L"\\");\r
-      StrCat(NewFileInfo->FileName, DestPath);\r
+    if (NewFileInfo == NULL) {\r
+      ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_MEM), gShellLevel2HiiHandle);\r
+      ShellStatus = SHELL_OUT_OF_RESOURCES;\r
     } else {\r
-      StrCpy(NewFileInfo->FileName, DestPath);\r
-    }\r
-    if (NewFileInfo->FileName[StrLen(NewFileInfo->FileName)-1] == L'\\') {\r
-      if (Node->FileName[0] == L'\\') {\r
-        //\r
-        // Don't allow for double slashes. Eliminate one of them.\r
-        //\r
-        NewFileInfo->FileName[StrLen(NewFileInfo->FileName)-1] = CHAR_NULL;\r
+      CopyMem(NewFileInfo, Node->Info, sizeof(EFI_FILE_INFO));\r
+      if (DestPath[0] != L'\\') {\r
+        StrCpy(NewFileInfo->FileName, L"\\");\r
+        StrCat(NewFileInfo->FileName, DestPath);\r
+      } else {\r
+        StrCpy(NewFileInfo->FileName, DestPath);\r
       }\r
-      StrCat(NewFileInfo->FileName, Node->FileName);\r
-    }\r
-    NewFileInfo->Size = sizeof(EFI_FILE_INFO) + StrSize(NewFileInfo->FileName);\r
+      if (NewFileInfo->FileName[StrLen(NewFileInfo->FileName)-1] == L'\\') {\r
+        if (Node->FileName[0] == L'\\') {\r
+          //\r
+          // Don't allow for double slashes. Eliminate one of them.\r
+          //\r
+          NewFileInfo->FileName[StrLen(NewFileInfo->FileName)-1] = CHAR_NULL;\r
+        }\r
+        StrCat(NewFileInfo->FileName, Node->FileName);\r
+      }\r
+      NewFileInfo->Size = sizeof(EFI_FILE_INFO) + StrSize(NewFileInfo->FileName);\r
 \r
-    ShellPrintEx(-1, -1, HiiOutput, Node->FullName, NewFileInfo->FileName);\r
+      ShellPrintEx(-1, -1, HiiOutput, Node->FullName, NewFileInfo->FileName);\r
 \r
-    //\r
-    // Perform the move operation\r
-    //\r
-    Status = ShellSetFileInfo(Node->Handle, NewFileInfo);\r
+      //\r
+      // Perform the move operation\r
+      //\r
+      Status = ShellSetFileInfo(Node->Handle, NewFileInfo);\r
 \r
-    //\r
-    // Free the info object we used...\r
-    //\r
-    ASSERT  (NewFileInfo != NULL);\r
-    FreePool(NewFileInfo);\r
+      //\r
+      // Free the info object we used...\r
+      //\r
+      ASSERT  (NewFileInfo != NULL);\r
+      FreePool(NewFileInfo);\r
 \r
-    //\r
-    // Check our result\r
-    //\r
-    if (EFI_ERROR(Status)) {\r
-      ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_ERR_UK), gShellLevel2HiiHandle, Status);\r
       //\r
-      // move failed\r
+      // Check our result\r
       //\r
-      switch(Status){\r
-        default:\r
-          ShellStatus = SHELL_INVALID_PARAMETER;\r
-        case EFI_SECURITY_VIOLATION:\r
-          ShellStatus = SHELL_SECURITY_VIOLATION;\r
-        case EFI_WRITE_PROTECTED:\r
-          ShellStatus = SHELL_WRITE_PROTECTED;\r
-        case EFI_OUT_OF_RESOURCES:\r
-          ShellStatus = SHELL_OUT_OF_RESOURCES;\r
-        case EFI_DEVICE_ERROR:\r
-          ShellStatus = SHELL_DEVICE_ERROR;\r
-        case EFI_ACCESS_DENIED:\r
-          ShellStatus = SHELL_ACCESS_DENIED;\r
-      } // switch\r
-    } else {\r
-      ShellPrintEx(-1, -1, L"%s", HiiResultOk);\r
+      if (EFI_ERROR(Status)) {\r
+        ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_ERR_UK), gShellLevel2HiiHandle, Status);\r
+        //\r
+        // move failed\r
+        //\r
+        switch(Status){\r
+          default:\r
+            ShellStatus = SHELL_INVALID_PARAMETER;\r
+          case EFI_SECURITY_VIOLATION:\r
+            ShellStatus = SHELL_SECURITY_VIOLATION;\r
+          case EFI_WRITE_PROTECTED:\r
+            ShellStatus = SHELL_WRITE_PROTECTED;\r
+          case EFI_OUT_OF_RESOURCES:\r
+            ShellStatus = SHELL_OUT_OF_RESOURCES;\r
+          case EFI_DEVICE_ERROR:\r
+            ShellStatus = SHELL_DEVICE_ERROR;\r
+          case EFI_ACCESS_DENIED:\r
+            ShellStatus = SHELL_ACCESS_DENIED;\r
+        } // switch\r
+      } else {\r
+        ShellPrintEx(-1, -1, L"%s", HiiResultOk);\r
+      }\r
     }\r
   } // for loop\r
 \r