Verify more memory allocations.

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

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);

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);
+  if (EFI_ERROR(Status)) {
-  ASSERT(HandleBuffer != NULL);
+    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.

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);
       }

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 (TestPath != NULL) {
-      if (TempChar != NULL) {
+        TempChar = StrStr(TestPath, L";");
-        *TempChar = CHAR_NULL;
+        if (TempChar != NULL) {
-      }
+          *TempChar = CHAR_NULL;
-      if (TestPath[StrLen(TestPath)-1] != L'\\') {
+        }
-        StrCat(TestPath, L"\\");
+        if (TestPath[StrLen(TestPath)-1] != L'\\') {
-      }
+          StrCat(TestPath, L"\\");
-      StrCat(TestPath, FileName);
+        }
-      if (StrStr(Walker, L";") != NULL) {
+        StrCat(TestPath, FileName);
-        Walker = StrStr(Walker, L";") + 1;
+        if (StrStr(Walker, L";") != NULL) {
-      } else {
+          Walker = StrStr(Walker, L";") + 1;
-        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);
+          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);
+      TempLocation  = StrnCatGrow(&TempLocation, NULL, DirName, 0);
-      if (StrStr(TempLocation, L":") != NULL && StrLen(StrStr(TempLocation, L":")) == 2) {
+      TempLocation2 = StrStr(TempLocation, L":");
-        *(StrStr(TempLocation, L":")+1) = CHAR_NULL;
+      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) {