From fe6c94d2e158e13fb6452979fc730fdd7c8f6447 Mon Sep 17 00:00:00 2001
From: Brendan Jackman <Brendan.Jackman@arm.com>
Date: Tue, 11 Feb 2014 22:46:56 +0000
Subject: ShellPkg: InternalShellExecuteDevicePath: avoid memory leaks

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@15224 6f19259b-4bc3-4df7-8a09-765794883524
---
 ShellPkg/Application/Shell/ShellProtocol.c | 42 ++++++++++++++++--------------
 1 file changed, 22 insertions(+), 20 deletions(-)

diff --git a/ShellPkg/Application/Shell/ShellProtocol.c b/ShellPkg/Application/Shell/ShellProtocol.c
index 5027137705..0a19793a1b 100644
--- a/ShellPkg/Application/Shell/ShellProtocol.c
+++ b/ShellPkg/Application/Shell/ShellProtocol.c
@@ -1397,6 +1397,7 @@ InternalShellExecuteDevicePath(
   UINTN                         InternalExitDataSize;
   UINTN                         *ExitDataSizePtr;
   CHAR16                        *ImagePath;
+  UINTN                         Index;
 
   // ExitDataSize is not OPTIONAL for gBS->BootServices, provide somewhere for
   // it to be dumped if the caller doesn't want it.
@@ -1482,7 +1483,7 @@ InternalShellExecuteDevicePath(
         ShellParamsProtocol.Argv = AllocatePool (sizeof (CHAR16 *));
         if (ShellParamsProtocol.Argv == NULL) {
           Status = EFI_OUT_OF_RESOURCES;
-          goto Cleanup;
+          goto UnloadImage;
         }
         ShellParamsProtocol.Argc = 1;
       } else {
@@ -1506,36 +1507,37 @@ InternalShellExecuteDevicePath(
                           ExitDataSizePtr,
                           ExitData
                           );
-    }
 
-    //
-    // Cleanup (and dont overwrite errors)
-    //
-    if (EFI_ERROR(Status)) {
-      CleanupStatus = gBS->UninstallProtocolInterface(
-                            NewHandle,
-                            &gEfiShellParametersProtocolGuid,
-                            &ShellParamsProtocol
-                            );
-      ASSERT_EFI_ERROR(CleanupStatus);
-    } else {
       CleanupStatus = gBS->UninstallProtocolInterface(
                             NewHandle,
                             &gEfiShellParametersProtocolGuid,
                             &ShellParamsProtocol
                             );
       ASSERT_EFI_ERROR(CleanupStatus);
+
+      goto FreeAlloc;
+    }
+
+UnloadImage:
+    // Unload image - We should only get here if we didn't call StartImage
+    gBS->UnloadImage (NewHandle);
+
+FreeAlloc:
+    // Free Argv (Allocated in UpdateArgcArgv)
+    if (ShellParamsProtocol.Argv != NULL) {
+      for (Index = 0; Index < ShellParamsProtocol.Argc; Index++) {
+        if (ShellParamsProtocol.Argv[Index] != NULL) {
+          FreePool (ShellParamsProtocol.Argv[Index]);
+        }
+      }
+      FreePool (ShellParamsProtocol.Argv);
     }
   }
 
+  // Restore environment variables
   if (!IsListEmpty(&OrigEnvs)) {
-    if (EFI_ERROR(Status)) {
-      CleanupStatus = SetEnvironmentVariableList(&OrigEnvs);
-      ASSERT_EFI_ERROR(CleanupStatus);
-    } else {
-      CleanupStatus = SetEnvironmentVariableList(&OrigEnvs);
-      ASSERT_EFI_ERROR (CleanupStatus);
-    }
+    CleanupStatus = SetEnvironmentVariableList(&OrigEnvs);
+    ASSERT_EFI_ERROR (CleanupStatus);
   }
 
   return(Status);
-- 
cgit v1.2.3