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

1 files changed, 27 insertions, 21 deletions

diff --git a/ShellPkg/Application/Shell/Shell.c b/ShellPkg/Application/Shell/Shell.c
index 87fd5e15ed..df101a32af 100644
--- a/ShellPkg/Application/Shell/Shell.c
+++ b/ShellPkg/Application/Shell/Shell.c
@@ -328,7 +328,8 @@ UefiMain (
         ///@todo Add our package into Framework HII

       }

       if (ShellInfoObject.HiiHandle == NULL) {

-        return (EFI_NOT_STARTED);

+        Status = EFI_NOT_STARTED;

+        goto FreeResources;

       }

     }

 

@@ -528,6 +529,7 @@ UefiMain (
     }

   }

 

+FreeResources:

   //

   // uninstall protocols / free memory / etc...

   //

@@ -589,29 +591,33 @@ UefiMain (
     DEBUG_CODE(ShellInfoObject.ConsoleInfo = NULL;);

   }

 

-  // If the command exited with an error, we pass this error out in the ExitData

-  // so that it can be retrieved by the EfiShellExecute function (which may

-  // start the shell with gBS->StartImage)

-  if (ExitStatus != SHELL_SUCCESS) {

-    // Allocate a buffer for exit data to pass to gBS->Exit().

-    // This buffer will contain the empty string immediately followed by

-    // the shell's exit status. (The empty string is required by the UEFI spec)

-    ExitDataSize = (sizeof (CHAR16) + sizeof (SHELL_STATUS));

-    ExitData = AllocatePool (ExitDataSize);

-    if (ExitData == NULL) {

-      return EFI_OUT_OF_RESOURCES;

-    }

-    ExitData[0] = '\0';

-    // Use CopyMem to avoid alignment faults

-    CopyMem ((ExitData + 1), &ExitStatus, sizeof (ExitStatus));

+  if (!EFI_ERROR (Status)) {

+    // If the command exited with an error, we pass this error out in the ExitData

+    // so that it can be retrieved by the EfiShellExecute function (which may

+    // start the shell with gBS->StartImage)

+    if (ExitStatus != SHELL_SUCCESS) {

+      // Allocate a buffer for exit data to pass to gBS->Exit().

+      // This buffer will contain the empty string immediately followed by

+      // the shell's exit status. (The empty string is required by the UEFI spec)

+      ExitDataSize = (sizeof (CHAR16) + sizeof (SHELL_STATUS));

+      ExitData = AllocatePool (ExitDataSize);

+      if (ExitData == NULL) {

+        return EFI_OUT_OF_RESOURCES;

+      }

+      ExitData[0] = '\0';

+      // Use CopyMem to avoid alignment faults

+      CopyMem ((ExitData + 1), &ExitStatus, sizeof (ExitStatus));

 

-    gBS->Exit (ImageHandle, EFI_ABORTED, ExitDataSize, ExitData);

+      gBS->Exit (ImageHandle, EFI_ABORTED, ExitDataSize, ExitData);

+

+      ASSERT (FALSE);

+      return EFI_SUCCESS;

+    } else {

+      return EFI_SUCCESS;

+    }

   } else {

-    return EFI_SUCCESS;

+    return Status;

   }

-

-  ASSERT (FALSE);

-  return EFI_SUCCESS;

 }

 

 /**