Verify more memory allocations.

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

4 files changed, 64 insertions, 32 deletions

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

   DestPath    = AllocatePool(PathLen);

 

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

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

     SHELL_FREE_NON_NULL(DestPath);

     SHELL_FREE_NON_NULL(HiiOutput);

     SHELL_FREE_NON_NULL(HiiResultOk);

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Map.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Map.c
index f696523f82..ccfa0b0654 100644
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Map.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Map.c
@@ -664,6 +664,9 @@ PerformMappingDelete(
     HandleBuffer);

   if (Status == EFI_BUFFER_TOO_SMALL) {

     HandleBuffer = AllocatePool(BufferSize);

+    if (HandleBuffer == NULL) {

+      return (EFI_OUT_OF_RESOURCES);

+    }

     Status = gBS->LocateHandle(

       ByProtocol,

       &gEfiDevicePathProtocolGuid,

@@ -671,8 +674,10 @@ PerformMappingDelete(
       &BufferSize,

       HandleBuffer);

   }

-  ASSERT_EFI_ERROR(Status);

-  ASSERT(HandleBuffer != NULL);

+  if (EFI_ERROR(Status)) {

+    SHELL_FREE_NON_NULL(HandleBuffer);

+    return (Status);

+  }

 

   //

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

@@ -698,6 +703,9 @@ PerformMappingDelete(
   if (Status == EFI_BUFFER_TOO_SMALL) {

     FreePool(HandleBuffer);

     HandleBuffer = AllocatePool(BufferSize);

+    if (HandleBuffer == NULL) {

+      return (EFI_OUT_OF_RESOURCES);

+    }

     Status = gBS->LocateHandle(

       ByProtocol,

       &gEfiBlockIoProtocolGuid,

@@ -705,7 +713,10 @@ PerformMappingDelete(
       &BufferSize,

       HandleBuffer);

   }

-  ASSERT_EFI_ERROR(Status);

+  if (EFI_ERROR(Status)) {

+    SHELL_FREE_NON_NULL(HandleBuffer);

+    return (Status);

+  }

 

   //

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

diff --git a/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c b/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c
index a628781757..7530289a52 100644
--- a/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c
+++ b/ShellPkg/Library/UefiShellLevel2CommandsLib/Mv.c
@@ -44,7 +44,7 @@ IsValidMove(
   CHAR16  *Test;

   CHAR16  *Test1;

   CHAR16  *TestWalker;

-  UINTN   Result;

+  INTN    Result;

   UINTN   TempLen;

   if (Cwd != NULL && StrCmp(FullName, Cwd) == 0) {

     //

@@ -236,6 +236,7 @@ ValidateAndMoveFiles(
   EFI_FILE_INFO             *NewFileInfo;

   CHAR16                    *TempLocation;

   UINTN                     NewSize;

+  UINTN                     Length;

 

   ASSERT(FileList != NULL);

   ASSERT(DestDir  != NULL);

@@ -310,12 +311,16 @@ ValidateAndMoveFiles(
       } else {

         StrCpy(NewFileInfo->FileName, DestPath);

       }

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

+      Length = StrLen(NewFileInfo->FileName);

+      if (Length > 0) {

+        Length--;

+      }

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

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

           //

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

           //

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

+          NewFileInfo->FileName[Length] = CHAR_NULL;

         }

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

       }

diff --git a/ShellPkg/Library/UefiShellLib/UefiShellLib.c b/ShellPkg/Library/UefiShellLib/UefiShellLib.c
index 1694f46378..bcb6f2a7e4 100644
--- a/ShellPkg/Library/UefiShellLib/UefiShellLib.c
+++ b/ShellPkg/Library/UefiShellLib/UefiShellLib.c
@@ -1639,31 +1639,36 @@ ShellFindFilePath (
     Size = StrSize(Path)+sizeof(CHAR16);

     Size += StrSize(FileName);

     TestPath = AllocateZeroPool(Size);

+    if (TestPath == NULL) {

+      return (NULL);

+    }

     Walker = (CHAR16*)Path;

     do {

       CopyMem(TestPath, Walker, StrSize(Walker));

-      TempChar = StrStr(TestPath, L";");

-      if (TempChar != NULL) {

-        *TempChar = CHAR_NULL;

-      }

-      if (TestPath[StrLen(TestPath)-1] != L'\\') {

-        StrCat(TestPath, L"\\");

-      }

-      StrCat(TestPath, FileName);

-      if (StrStr(Walker, L";") != NULL) {

-        Walker = StrStr(Walker, L";") + 1;

-      } else {

-        Walker = NULL;

-      }

-      Status = ShellOpenFileByName(TestPath, &Handle, EFI_FILE_MODE_READ, 0);

-      if (!EFI_ERROR(Status)){

-        if (FileHandleIsDirectory(Handle) != EFI_SUCCESS) {

-          ASSERT(RetVal == NULL);

-          RetVal = StrnCatGrow(&RetVal, NULL, TestPath, 0);

-          ShellCloseFile(&Handle);

-          break;

+      if (TestPath != NULL) {

+        TempChar = StrStr(TestPath, L";");

+        if (TempChar != NULL) {

+          *TempChar = CHAR_NULL;

+        }

+        if (TestPath[StrLen(TestPath)-1] != L'\\') {

+          StrCat(TestPath, L"\\");

+        }

+        StrCat(TestPath, FileName);

+        if (StrStr(Walker, L";") != NULL) {

+          Walker = StrStr(Walker, L";") + 1;

         } else {

-          ShellCloseFile(&Handle);

+          Walker = NULL;

+        }

+        Status = ShellOpenFileByName(TestPath, &Handle, EFI_FILE_MODE_READ, 0);

+        if (!EFI_ERROR(Status)){

+          if (FileHandleIsDirectory(Handle) != EFI_SUCCESS) {

+            ASSERT(RetVal == NULL);

+            RetVal = StrnCatGrow(&RetVal, NULL, TestPath, 0);

+            ShellCloseFile(&Handle);

+            break;

+          } else {

+            ShellCloseFile(&Handle);

+          }

         }

       }

     } while (Walker != NULL && Walker[0] != CHAR_NULL);

@@ -2460,6 +2465,9 @@ ShellCopySearchAndReplace(
     Replace = AllocateZeroPool(StrSize(ReplaceWith) + 2*sizeof(CHAR16));

     UnicodeSPrint(Replace, StrSize(ReplaceWith) + 2*sizeof(CHAR16), L"\"%s\"", ReplaceWith);

   }

+  if (Replace == NULL) {

+    return (EFI_OUT_OF_RESOURCES);

+  }

   NewString = SetMem16(NewString, NewSize, CHAR_NULL);

   while (*SourceString != CHAR_NULL) {

     //

@@ -2702,6 +2710,9 @@ ShellPrintEx(
 {

   VA_LIST           Marker;

   EFI_STATUS        RetVal;

+  if (Format == NULL) {

+    return (EFI_INVALID_PARAMETER);

+  }

   VA_START (Marker, Format);

   RetVal = InternalShellPrintWorker(Col, Row, Format, Marker);

   VA_END(Marker);

@@ -2784,7 +2795,8 @@ ShellIsDirectory(
 {

   EFI_STATUS        Status;

   SHELL_FILE_HANDLE Handle;

-        CHAR16      *TempLocation;

+  CHAR16            *TempLocation;

+  CHAR16            *TempLocation2;

 

   ASSERT(DirName != NULL);

 

@@ -2797,9 +2809,10 @@ ShellIsDirectory(
     // try good logic first.

     //

     if (mEfiShellProtocol != NULL) {

-      TempLocation = StrnCatGrow(&TempLocation, NULL, DirName, 0);

-      if (StrStr(TempLocation, L":") != NULL && StrLen(StrStr(TempLocation, L":")) == 2) {

-        *(StrStr(TempLocation, L":")+1) = CHAR_NULL;

+      TempLocation  = StrnCatGrow(&TempLocation, NULL, DirName, 0);

+      TempLocation2 = StrStr(TempLocation, L":");

+      if (TempLocation2 != NULL && StrLen(StrStr(TempLocation, L":")) == 2) {

+        *(TempLocation2+1) = CHAR_NULL;

       }

       if (mEfiShellProtocol->GetDevicePathFromMap(TempLocation) != NULL) {

         FreePool(TempLocation);

@@ -3074,6 +3087,9 @@ ShellPromptForResponse (
   Size    = 0;

   if (Type != ShellPromptResponseTypeFreeform) {

     Resp = (SHELL_PROMPT_RESPONSE*)AllocatePool(sizeof(SHELL_PROMPT_RESPONSE));

+    if (Resp == NULL) {

+      return (EFI_OUT_OF_RESOURCES);

+    }

   }

 

   switch(Type) {