ShellPkg: Fixed Memory leak in UefiMain()

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Brendan Jackman <Brendan.Jackman@arm.com>
Reviewed-by: Olivier Martin <olivier.martin@arm.com>
Reviewed-by: Jaben Carsey <jaben.carsey@intel.com>
git-svn-id: https://svn.code.sf.net/p/edk2/code/trunk/edk2@15221 6f19259b-4bc3-4df7-8a09-765794883524

diff --git a/ShellPkg/Application/Shell/Shell.c b/ShellPkg/Application/Shell/Shell.c
index 87fd5e15edf87a4def0249193a5edf8f0658f6f3..df101a32af0157e16eeac13a085e67e62397796b 100644 (file)
--- a/ShellPkg/Application/Shell/Shell.c
+++ b/ShellPkg/Application/Shell/Shell.c
@@ -328,7 +328,8 @@ UefiMain (
         ///@todo Add our package into Framework HII\r
       }\r
       if (ShellInfoObject.HiiHandle == NULL) {\r
-        return (EFI_NOT_STARTED);\r
+        Status = EFI_NOT_STARTED;\r
+        goto FreeResources;\r
       }\r
     }\r
 \r
@@ -528,6 +529,7 @@ UefiMain (
     }\r
   }\r
 \r
+FreeResources:\r
   //\r
   // uninstall protocols / free memory / etc...\r
   //\r
@@ -589,29 +591,33 @@ UefiMain (
     DEBUG_CODE(ShellInfoObject.ConsoleInfo = NULL;);\r
   }\r
 \r
-  // If the command exited with an error, we pass this error out in the ExitData\r
-  // so that it can be retrieved by the EfiShellExecute function (which may\r
-  // start the shell with gBS->StartImage)\r
-  if (ExitStatus != SHELL_SUCCESS) {\r
-    // Allocate a buffer for exit data to pass to gBS->Exit().\r
-    // This buffer will contain the empty string immediately followed by\r
-    // the shell's exit status. (The empty string is required by the UEFI spec)\r
-    ExitDataSize = (sizeof (CHAR16) + sizeof (SHELL_STATUS));\r
-    ExitData = AllocatePool (ExitDataSize);\r
-    if (ExitData == NULL) {\r
-      return EFI_OUT_OF_RESOURCES;\r
-    }\r
-    ExitData[0] = '\0';\r
-    // Use CopyMem to avoid alignment faults\r
-    CopyMem ((ExitData + 1), &ExitStatus, sizeof (ExitStatus));\r
+  if (!EFI_ERROR (Status)) {\r
+    // If the command exited with an error, we pass this error out in the ExitData\r
+    // so that it can be retrieved by the EfiShellExecute function (which may\r
+    // start the shell with gBS->StartImage)\r
+    if (ExitStatus != SHELL_SUCCESS) {\r
+      // Allocate a buffer for exit data to pass to gBS->Exit().\r
+      // This buffer will contain the empty string immediately followed by\r
+      // the shell's exit status. (The empty string is required by the UEFI spec)\r
+      ExitDataSize = (sizeof (CHAR16) + sizeof (SHELL_STATUS));\r
+      ExitData = AllocatePool (ExitDataSize);\r
+      if (ExitData == NULL) {\r
+        return EFI_OUT_OF_RESOURCES;\r
+      }\r
+      ExitData[0] = '\0';\r
+      // Use CopyMem to avoid alignment faults\r
+      CopyMem ((ExitData + 1), &ExitStatus, sizeof (ExitStatus));\r
 \r
-    gBS->Exit (ImageHandle, EFI_ABORTED, ExitDataSize, ExitData);\r
+      gBS->Exit (ImageHandle, EFI_ABORTED, ExitDataSize, ExitData);\r
+\r
+      ASSERT (FALSE);\r
+      return EFI_SUCCESS;\r
+    } else {\r
+      return EFI_SUCCESS;\r
+    }\r
   } else {\r
-    return EFI_SUCCESS;\r
+    return Status;\r
   }\r
-\r
-  ASSERT (FALSE);\r
-  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r