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

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

diff --git a/ShellPkg/Application/Shell/ShellParametersProtocol.c b/ShellPkg/Application/Shell/ShellParametersProtocol.c
index 02606d615bedf86fb9947640e96b6a7881d31136..55cf3c279eb74401d6d62ed64b3095094f80a42d 100644 (file)
--- a/ShellPkg/Application/Shell/ShellParametersProtocol.c
+++ b/ShellPkg/Application/Shell/ShellParametersProtocol.c
@@ -470,6 +470,31 @@ StripQuotes (
   }\r
 }\r
 \r
+/**\r
+  Calcualte the 32-bit CRC in a EFI table using the service provided by the\r
+  gRuntime service.\r
+\r
+  @param  Hdr                    Pointer to an EFI standard header\r
+\r
+**/\r
+VOID\r
+CalculateEfiHdrCrc (\r
+  IN  OUT EFI_TABLE_HEADER    *Hdr\r
+  )\r
+{\r
+  UINT32 Crc;\r
+\r
+  Hdr->CRC32 = 0;\r
+\r
+  //\r
+  // If gBS->CalculateCrce32 () == CoreEfiNotAvailableYet () then\r
+  //  Crc will come back as zero if we set it to zero here\r
+  //\r
+  Crc = 0;\r
+  gBS->CalculateCrc32 ((UINT8 *)Hdr, Hdr->HeaderSize, &Crc);\r
+  Hdr->CRC32 = Crc;\r
+}\r
+\r
 /**\r
   Funcion will replace the current StdIn and StdOut in the ShellParameters protocol\r
   structure by parsing NewCommandLine.  The current values are returned to the\r
@@ -552,6 +577,9 @@ UpdateStdInStdOutStdErr(
   }\r
 \r
   CommandLineCopy = StrnCatGrow(&CommandLineCopy, NULL, NewCommandLine, 0);\r
+  if (CommandLineCopy == NULL) {\r
+    return (EFI_OUT_OF_RESOURCES);\r
+  }\r
   Status          = EFI_SUCCESS;\r
   Split           = NULL;\r
   FirstLocation   = CommandLineCopy + StrLen(CommandLineCopy);\r
@@ -1005,16 +1033,19 @@ UpdateStdInStdOutStdErr(
       //\r
       if (!EFI_ERROR(Status) && StdInVarName != NULL) {\r
         TempHandle = CreateFileInterfaceEnv(StdInVarName);\r
-        if (!InUnicode) {\r
-          TempHandle = CreateFileInterfaceFile(TempHandle, FALSE);\r
-        }\r
-        Size = 0;\r
-        ASSERT(TempHandle != NULL);\r
-        if (((EFI_FILE_PROTOCOL*)TempHandle)->Read(TempHandle, &Size, NULL) != EFI_BUFFER_TOO_SMALL) {\r
-          Status = EFI_INVALID_PARAMETER;\r
+        if (TempHandle == NULL) {\r
+          Status = EFI_OUT_OF_RESOURCES;\r
         } else {\r
-          ShellParameters->StdIn = TempHandle;\r
-          gST->ConIn = CreateSimpleTextInOnFile(TempHandle, &gST->ConsoleInHandle);\r
+          if (!InUnicode) {\r
+            TempHandle = CreateFileInterfaceFile(TempHandle, FALSE);\r
+          }\r
+          Size = 0;\r
+          if (((EFI_FILE_PROTOCOL*)TempHandle)->Read(TempHandle, &Size, NULL) != EFI_BUFFER_TOO_SMALL) {\r
+            Status = EFI_INVALID_PARAMETER;\r
+          } else {\r
+            ShellParameters->StdIn = TempHandle;\r
+            gST->ConIn = CreateSimpleTextInOnFile(TempHandle, &gST->ConsoleInHandle);\r
+          }\r
         }\r
       }\r
 \r
@@ -1039,6 +1070,8 @@ UpdateStdInStdOutStdErr(
   }\r
   FreePool(CommandLineCopy);\r
 \r
+  CalculateEfiHdrCrc(&gST->Hdr);\r
+\r
   if (gST->ConIn == NULL ||gST->ConOut == NULL) {\r
     return (EFI_OUT_OF_RESOURCES);\r
   }\r
@@ -1112,6 +1145,8 @@ RestoreStdInStdOutStdErr (
     gST->StandardErrorHandle  = SystemTableInfo->ConErrHandle;\r
   }\r
 \r
+  CalculateEfiHdrCrc(&gST->Hdr);\r
+\r
   return (EFI_SUCCESS);\r
 }\r
 /**\r