ShellPkg: InternalShellExecuteDevicePath: avoid memory leaks

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

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