ShellPkg: Add checking for memory allocation and pointer returns from functions.

[mirror_edk2.git] / ShellPkg / Application / Shell / Shell.c

diff --git a/ShellPkg/Application/Shell/Shell.c b/ShellPkg/Application/Shell/Shell.c
index 614c1e9595d3d02132750a1c3550a82c02db47d8..d266c826e9572d3f9f8caa152b942d82a6f660a4 100644 (file)
--- a/ShellPkg/Application/Shell/Shell.c
+++ b/ShellPkg/Application/Shell/Shell.c
@@ -857,15 +857,19 @@ DoStartupScript(
     FileStringPath = NULL;\r
     NewSize = 0;\r
     FileStringPath = StrnCatGrow(&FileStringPath, &NewSize, MapName, 0);\r
-    TempSpot = StrStr(FileStringPath, L";");\r
-    if (TempSpot != NULL) {\r
-      *TempSpot = CHAR_NULL;\r
+    if (FileStringPath == NULL) {\r
+      Status = EFI_OUT_OF_RESOURCES;\r
+    } else {\r
+      TempSpot = StrStr(FileStringPath, L";");\r
+      if (TempSpot != NULL) {\r
+        *TempSpot = CHAR_NULL;\r
+      }\r
+      FileStringPath = StrnCatGrow(&FileStringPath, &NewSize, ((FILEPATH_DEVICE_PATH*)FilePath)->PathName, 0);\r
+      PathRemoveLastItem(FileStringPath);\r
+      FileStringPath = StrnCatGrow(&FileStringPath, &NewSize, mStartupScript, 0);\r
+      Status = ShellInfoObject.NewEfiShellProtocol->OpenFileByName(FileStringPath, &FileHandle, EFI_FILE_MODE_READ);\r
+      FreePool(FileStringPath);\r
     }\r
-    FileStringPath = StrnCatGrow(&FileStringPath, &NewSize, ((FILEPATH_DEVICE_PATH*)FilePath)->PathName, 0);\r
-    PathRemoveLastItem(FileStringPath);\r
-    FileStringPath = StrnCatGrow(&FileStringPath, &NewSize, mStartupScript, 0);\r
-    Status = ShellInfoObject.NewEfiShellProtocol->OpenFileByName(FileStringPath, &FileHandle, EFI_FILE_MODE_READ);\r
-    FreePool(FileStringPath);\r
   }\r
   if (EFI_ERROR(Status)) {\r
     NamePath = FileDevicePath (NULL, mStartupScript);\r
@@ -1215,6 +1219,12 @@ RunSplitCommand(
 \r
   NextCommandLine = StrnCatGrow(&NextCommandLine, &Size1, StrStr(CmdLine, L"|")+1, 0);\r
   OurCommandLine  = StrnCatGrow(&OurCommandLine , &Size2, CmdLine                , StrStr(CmdLine, L"|") - CmdLine);\r
+\r
+  if (NextCommandLine == NULL || OurCommandLine == NULL) {\r
+    SHELL_FREE_NON_NULL(OurCommandLine);\r
+    SHELL_FREE_NON_NULL(NextCommandLine);\r
+    return (EFI_OUT_OF_RESOURCES);\r
+  }\r
   if (NextCommandLine[0] != CHAR_NULL &&\r
       NextCommandLine[0] == L'a' &&\r
       NextCommandLine[1] == L' '\r
@@ -1334,6 +1344,9 @@ RunCommand(
   Split               = NULL;\r
 \r
   CleanOriginal = StrnCatGrow(&CleanOriginal, NULL, CmdLine, 0);\r
+  if (CleanOriginal == NULL) {\r
+    return (EFI_OUT_OF_RESOURCES);\r
+  }\r
   while (CleanOriginal[StrLen(CleanOriginal)-1] == L' ') {\r
     CleanOriginal[StrLen(CleanOriginal)-1] = CHAR_NULL;\r
   }\r