Verify memory allocations were successful.

git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@10909 6f19259b-4bc3-4df7-8a09-765794883524

7 files changed, 163 insertions, 109 deletions

diff --git a/ShellPkg/Library/UefiShellLevel1CommandsLib/For.c b/ShellPkg/Library/UefiShellLevel1CommandsLib/For.c
index 0e68119b8d..297a137664 100644
--- a/ShellPkg/Library/UefiShellLevel1CommandsLib/For.c
+++ b/ShellPkg/Library/UefiShellLevel1CommandsLib/For.c
@@ -103,8 +103,11 @@ typedef struct {
   @param[in] Alias              The alias to test for.

   @param[in] CommandString      The updated command string.

   @param[in,out] List           The list to search.

+

+  @retval EFI_SUCCESS           The operation was completed successfully.

+  @retval EFI_OUT_OF_RESOURCES  There was not enough free memory.

 **/

-VOID

+EFI_STATUS

 EFIAPI

 InternalUpdateAliasOnList(

   IN CONST CHAR16       *Alias,

@@ -139,12 +142,16 @@ InternalUpdateAliasOnList(
   }

   if (!Found) {

     Node = AllocateZeroPool(sizeof(ALIAS_LIST));

+    if (Node == NULL) {

+      return (EFI_OUT_OF_RESOURCES);

+    }

     ASSERT(Node->Alias == NULL);

     Node->Alias         = StrnCatGrow(&Node->Alias, NULL, Alias, 0);

     ASSERT(Node->CommandString == NULL);

     Node->CommandString = StrnCatGrow(&Node->CommandString, NULL, CommandString, 0);

     InsertTailList(List, &Node->Link);

   }

+  return (EFI_SUCCESS);

 }

 

 /**

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Cd.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Cd.c
index 78403c2274..99680b49d9 100644
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Cd.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Cd.c
@@ -186,17 +186,27 @@ ShellCommandRunCd (
         // change directory on other drive letter

         //

         Drive = AllocateZeroPool(StrSize(Param1));

-        Drive = StrCpy(Drive, Param1);

-        Path = StrStr(Drive, L":");

-        *(++Path) = CHAR_NULL;

-        Status = gEfiShellProtocol->SetCurDir(Drive, ++Path);

+        if (Drive == NULL) {

+          ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_MEM), gShellLevel2HiiHandle);

+          ShellStatus = SHELL_OUT_OF_RESOURCES;

+        } else {

+          Drive = StrCpy(Drive, Param1);

+          Path = StrStr(Drive, L":");

+          *(++Path) = CHAR_NULL;

+          if (Path == Drive + StrLen(Drive)) {

+            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle);

+            ShellStatus = SHELL_NOT_FOUND;

+          } else {

+            Status = gEfiShellProtocol->SetCurDir(Drive, ++Path);

+          }

 

-        if (Status == EFI_NOT_FOUND) {

-          ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle);

-          Status = SHELL_NOT_FOUND;

-        } else if (EFI_ERROR(Status)) {

-          ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_DIR_NF), gShellLevel2HiiHandle, Param1);

-          Status = SHELL_NOT_FOUND;

+          if (Status == EFI_NOT_FOUND) {

+            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_CD_NF), gShellLevel2HiiHandle);

+            Status = SHELL_NOT_FOUND;

+          } else if (EFI_ERROR(Status)) {

+            ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_DIR_NF), gShellLevel2HiiHandle, Param1);

+            Status = SHELL_NOT_FOUND;

+          }

         }

       }

     }

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Cp.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Cp.c
index 94161cfecb..9344c331e5 100644
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Cp.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Cp.c
@@ -287,6 +287,13 @@ ValidateAndCopyFiles(
   HiiResultOk = HiiGetString (gShellLevel2HiiHandle, STRING_TOKEN (STR_GEN_RES_OK), NULL);

   DestPath    = AllocatePool(PathLen);

 

+  if (HiiOutput == NULL || HiiOutput == NULL || HiiResultOk == NULL) {

+    SHELL_FREE_NON_NULL(DestPath);

+    SHELL_FREE_NON_NULL(HiiOutput);

+    SHELL_FREE_NON_NULL(HiiResultOk);

+    return (SHELL_OUT_OF_RESOURCES);

+  }

+

   //

   // Go through the list of files to copy...

   //

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Ls.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Ls.c
index 61d48977f1..440d2450f5 100644
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Ls.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Ls.c
@@ -349,30 +349,35 @@ PrintLsOutput(
 

   if (Rec){

     DirectoryName = AllocatePool(LongestPath + 2*sizeof(CHAR16));

-    for ( Node = (EFI_SHELL_FILE_INFO *)GetFirstNode(&ListHead->Link)

-        ; !IsNull(&ListHead->Link, &Node->Link)

-        ; Node = (EFI_SHELL_FILE_INFO *)GetNextNode(&ListHead->Link, &Node->Link)

-       ){

-      //

-      // recurse on any directory except the traversing ones...

-      //

-      if (((Node->Info->Attribute & EFI_FILE_DIRECTORY) == EFI_FILE_DIRECTORY)

-        && StrCmp(Node->FileName, L".") != 0

-        && StrCmp(Node->FileName, L"..") != 0

-       ){

-        StrCpy(DirectoryName, Node->FullName);

-        StrCat(DirectoryName, L"\\*");

-        PrintLsOutput(

-          Rec,

-          Attribs,

-          Sfo,

-          DirectoryName,

-          FALSE,

-          Count,

-          TimeZone);

+    if (DirectoryName == NULL) {

+      ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_MEM), gShellLevel2HiiHandle);

+      ShellStatus = SHELL_OUT_OF_RESOURCES;

+    } else {

+      for ( Node = (EFI_SHELL_FILE_INFO *)GetFirstNode(&ListHead->Link)

+          ; !IsNull(&ListHead->Link, &Node->Link)

+          ; Node = (EFI_SHELL_FILE_INFO *)GetNextNode(&ListHead->Link, &Node->Link)

+         ){

+        //

+        // recurse on any directory except the traversing ones...

+        //

+        if (((Node->Info->Attribute & EFI_FILE_DIRECTORY) == EFI_FILE_DIRECTORY)

+          && StrCmp(Node->FileName, L".") != 0

+          && StrCmp(Node->FileName, L"..") != 0

+         ){

+          StrCpy(DirectoryName, Node->FullName);

+          StrCat(DirectoryName, L"\\*");

+          PrintLsOutput(

+            Rec,

+            Attribs,

+            Sfo,

+            DirectoryName,

+            FALSE,

+            Count,

+            TimeZone);

+        }

       }

+      FreePool(DirectoryName);

     }

-    FreePool(DirectoryName);

   }

 

   FreePool(CorrectedPath);

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Map.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Map.c
index c4c8c87ff4..f696523f82 100644
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Map.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Map.c
@@ -99,7 +99,7 @@ UpdateMapping (
   //

   // Find each handle with Simple File System

   //

-  HandleList = GetHandleListByPotocol(&gEfiSimpleFileSystemProtocolGuid);

+  HandleList = GetHandleListByProtocol(&gEfiSimpleFileSystemProtocolGuid);

   if (HandleList != NULL) {

     //

     // Do a count of the handles

@@ -503,6 +503,9 @@ PerformMappingDisplay(
     HandleBuffer);

   if (Status == EFI_BUFFER_TOO_SMALL) {

     HandleBuffer = AllocatePool(BufferSize);

+    if (HandleBuffer == NULL) {

+      return (SHELL_OUT_OF_RESOURCES);

+    }

     Status = gBS->LocateHandle(

       ByProtocol,

       &gEfiDevicePathProtocolGuid,

@@ -542,6 +545,9 @@ PerformMappingDisplay(
   if (Status == EFI_BUFFER_TOO_SMALL) {

     FreePool(HandleBuffer);

     HandleBuffer = AllocatePool(BufferSize);

+    if (HandleBuffer == NULL) {

+      return (SHELL_OUT_OF_RESOURCES);

+    }

     Status = gBS->LocateHandle(

       ByProtocol,

       &gEfiBlockIoProtocolGuid,

@@ -549,37 +555,37 @@ PerformMappingDisplay(
       &BufferSize,

       HandleBuffer);

   }

-  ASSERT_EFI_ERROR(Status);

-

-  //

-  // Get the map name(s) for each one.

-  //

-  for ( LoopVar = 0

-      ; LoopVar < BufferSize / sizeof(EFI_HANDLE)

-      ; LoopVar ++

-     ){

+  if (!EFI_ERROR(Status)) {

     //

-    // Skip any that were already done...

+    // Get the map name(s) for each one.

     //

-    if (gBS->OpenProtocol(

-      HandleBuffer[LoopVar],

-      &gEfiDevicePathProtocolGuid,

-      NULL,

-      gImageHandle,

-      NULL,

-      EFI_OPEN_PROTOCOL_TEST_PROTOCOL) == EFI_SUCCESS) {

-        continue;

+    for ( LoopVar = 0

+        ; LoopVar < BufferSize / sizeof(EFI_HANDLE)

+        ; LoopVar ++

+       ){

+      //

+      // Skip any that were already done...

+      //

+      if (gBS->OpenProtocol(

+        HandleBuffer[LoopVar],

+        &gEfiDevicePathProtocolGuid,

+        NULL,

+        gImageHandle,

+        NULL,

+        EFI_OPEN_PROTOCOL_TEST_PROTOCOL) == EFI_SUCCESS) {

+          continue;

+      }

+      PerformSingleMappingDisplay(

+        Verbose,

+        Consist,

+        Normal,

+        Test,

+        SFO,

+        Specific,

+        HandleBuffer[LoopVar]);

     }

-    PerformSingleMappingDisplay(

-      Verbose,

-      Consist,

-      Normal,

-      Test,

-      SFO,

-      Specific,

-      HandleBuffer[LoopVar]);

+    FreePool(HandleBuffer);

   }

-  FreePool(HandleBuffer);

   return (SHELL_SUCCESS);

 }

 

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c
index adb6f99759..a628781757 100644
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c
@@ -152,6 +152,10 @@ GetDestinationLocation(
       NewSize = StrSize(Cwd);

       NewSize += StrSize(DestDir);

       DestPath = AllocateZeroPool(NewSize);

+      if (DestPath == NULL) {

+        ShellCloseFileMetaArg(&DestList);

+        return (SHELL_OUT_OF_RESOURCES);

+      }

       StrCpy(DestPath, Cwd);

       if (DestPath[StrLen(DestPath)-1] != L'\\' && DestDir[0] != L'\\') {

         StrCat(DestPath, L"\\");

@@ -162,6 +166,10 @@ GetDestinationLocation(
     } else {

       ASSERT(DestPath == NULL);

       DestPath = StrnCatGrow(&DestPath, NULL, DestDir, 0);

+      if (DestPath == NULL) {

+        ShellCloseFileMetaArg(&DestList);

+        return (SHELL_OUT_OF_RESOURCES);

+      }

     }

   } else {

     Node = (EFI_SHELL_FILE_INFO*)GetFirstNode(&DestList->Link);

@@ -175,6 +183,10 @@ GetDestinationLocation(
     }

     if (ShellIsDirectory(Node->FullName)==EFI_SUCCESS) {

       DestPath = AllocateZeroPool(StrSize(Node->FullName)+sizeof(CHAR16));

+      if (DestPath == NULL) {

+        ShellCloseFileMetaArg(&DestList);

+        return (SHELL_OUT_OF_RESOURCES);

+      }

       StrCpy(DestPath, Node->FullName);

       StrCat(DestPath, L"\\");

     } else {

@@ -287,62 +299,66 @@ ValidateAndMoveFiles(
     NewSize = StrSize(DestPath);

     NewSize += StrSize(Node->FileName) + sizeof(EFI_FILE_INFO) + sizeof(CHAR16);

     NewFileInfo = AllocateZeroPool(NewSize);

-    ASSERT(NewFileInfo != NULL);

-    CopyMem(NewFileInfo, Node->Info, sizeof(EFI_FILE_INFO));

-    if (DestPath[0] != L'\\') {

-      StrCpy(NewFileInfo->FileName, L"\\");

-      StrCat(NewFileInfo->FileName, DestPath);

+    if (NewFileInfo == NULL) {

+      ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_NO_MEM), gShellLevel2HiiHandle);

+      ShellStatus = SHELL_OUT_OF_RESOURCES;

     } else {

-      StrCpy(NewFileInfo->FileName, DestPath);

-    }

-    if (NewFileInfo->FileName[StrLen(NewFileInfo->FileName)-1] == L'\\') {

-      if (Node->FileName[0] == L'\\') {

-        //

-        // Don't allow for double slashes. Eliminate one of them.

-        //

-        NewFileInfo->FileName[StrLen(NewFileInfo->FileName)-1] = CHAR_NULL;

+      CopyMem(NewFileInfo, Node->Info, sizeof(EFI_FILE_INFO));

+      if (DestPath[0] != L'\\') {

+        StrCpy(NewFileInfo->FileName, L"\\");

+        StrCat(NewFileInfo->FileName, DestPath);

+      } else {

+        StrCpy(NewFileInfo->FileName, DestPath);

       }

-      StrCat(NewFileInfo->FileName, Node->FileName);

-    }

-    NewFileInfo->Size = sizeof(EFI_FILE_INFO) + StrSize(NewFileInfo->FileName);

+      if (NewFileInfo->FileName[StrLen(NewFileInfo->FileName)-1] == L'\\') {

+        if (Node->FileName[0] == L'\\') {

+          //

+          // Don't allow for double slashes. Eliminate one of them.

+          //

+          NewFileInfo->FileName[StrLen(NewFileInfo->FileName)-1] = CHAR_NULL;

+        }

+        StrCat(NewFileInfo->FileName, Node->FileName);

+      }

+      NewFileInfo->Size = sizeof(EFI_FILE_INFO) + StrSize(NewFileInfo->FileName);

 

-    ShellPrintEx(-1, -1, HiiOutput, Node->FullName, NewFileInfo->FileName);

+      ShellPrintEx(-1, -1, HiiOutput, Node->FullName, NewFileInfo->FileName);

 

-    //

-    // Perform the move operation

-    //

-    Status = ShellSetFileInfo(Node->Handle, NewFileInfo);

+      //

+      // Perform the move operation

+      //

+      Status = ShellSetFileInfo(Node->Handle, NewFileInfo);

 

-    //

-    // Free the info object we used...

-    //

-    ASSERT  (NewFileInfo != NULL);

-    FreePool(NewFileInfo);

+      //

+      // Free the info object we used...

+      //

+      ASSERT  (NewFileInfo != NULL);

+      FreePool(NewFileInfo);

 

-    //

-    // Check our result

-    //

-    if (EFI_ERROR(Status)) {

-      ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_ERR_UK), gShellLevel2HiiHandle, Status);

       //

-      // move failed

+      // Check our result

       //

-      switch(Status){

-        default:

-          ShellStatus = SHELL_INVALID_PARAMETER;

-        case EFI_SECURITY_VIOLATION:

-          ShellStatus = SHELL_SECURITY_VIOLATION;

-        case EFI_WRITE_PROTECTED:

-          ShellStatus = SHELL_WRITE_PROTECTED;

-        case EFI_OUT_OF_RESOURCES:

-          ShellStatus = SHELL_OUT_OF_RESOURCES;

-        case EFI_DEVICE_ERROR:

-          ShellStatus = SHELL_DEVICE_ERROR;

-        case EFI_ACCESS_DENIED:

-          ShellStatus = SHELL_ACCESS_DENIED;

-      } // switch

-    } else {

-      ShellPrintEx(-1, -1, L"%s", HiiResultOk);

+      if (EFI_ERROR(Status)) {

+        ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN (STR_GEN_ERR_UK), gShellLevel2HiiHandle, Status);

+        //

+        // move failed

+        //

+        switch(Status){

+          default:

+            ShellStatus = SHELL_INVALID_PARAMETER;

+          case EFI_SECURITY_VIOLATION:

+            ShellStatus = SHELL_SECURITY_VIOLATION;

+          case EFI_WRITE_PROTECTED:

+            ShellStatus = SHELL_WRITE_PROTECTED;

+          case EFI_OUT_OF_RESOURCES:

+            ShellStatus = SHELL_OUT_OF_RESOURCES;

+          case EFI_DEVICE_ERROR:

+            ShellStatus = SHELL_DEVICE_ERROR;

+          case EFI_ACCESS_DENIED:

+            ShellStatus = SHELL_ACCESS_DENIED;

+        } // switch

+      } else {

+        ShellPrintEx(-1, -1, L"%s", HiiResultOk);

+      }

     }

   } // for loop

 

diff --git a/ShellPkg/Library/UefiShellLevel3CommandsLib/Alias.c b/ShellPkg/Library/UefiShellLevel3CommandsLib/Alias.c
index da53af4c5c..1428d37a64 100644
--- a/ShellPkg/Library/UefiShellLevel3CommandsLib/Alias.c
+++ b/ShellPkg/Library/UefiShellLevel3CommandsLib/Alias.c
@@ -41,6 +41,9 @@ PrintAllShellAlias(
     return (SHELL_SUCCESS);

   }

   Alias = AllocateZeroPool(StrSize(ConstAllAliasList));

+  if (Alias == NULL) {

+    return (SHELL_OUT_OF_RESOURCES);

+  }

   Walker = (CHAR16*)ConstAllAliasList;

 

   do {