Make the BA responsible for parsing restart prompt behavior.

Fixes #4975

19 files changed, 180 insertions, 143 deletions

diff --git a/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h b/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
index 56f6b361..8301d45f 100644
--- a/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
+++ b/src/api/burn/WixToolset.BootstrapperCore.Native/inc/BootstrapperApplication.h
@@ -11,15 +11,6 @@ enum BOOTSTRAPPER_DISPLAY
     BOOTSTRAPPER_DISPLAY_FULL,
 };
 
-enum BOOTSTRAPPER_RESTART
-{
-    BOOTSTRAPPER_RESTART_UNKNOWN,
-    BOOTSTRAPPER_RESTART_NEVER,
-    BOOTSTRAPPER_RESTART_PROMPT,
-    BOOTSTRAPPER_RESTART_AUTOMATIC,
-    BOOTSTRAPPER_RESTART_ALWAYS,
-};
-
 enum BOOTSTRAPPER_REGISTRATION_TYPE
 {
     BOOTSTRAPPER_REGISTRATION_TYPE_NONE, // The engine will ignore NONE if it recommended INPROGRESS or FULL.
@@ -286,7 +277,6 @@ struct BOOTSTRAPPER_COMMAND
     DWORD cbSize;
     BOOTSTRAPPER_ACTION action;
     BOOTSTRAPPER_DISPLAY display;
-    BOOTSTRAPPER_RESTART restart;
 
     LPWSTR wzCommandLine;
     int nCmdShow;
diff --git a/src/api/burn/WixToolset.Mba.Core/BootstrapperCommand.cs b/src/api/burn/WixToolset.Mba.Core/BootstrapperCommand.cs
index 345e0448..88a9b9bb 100644
--- a/src/api/burn/WixToolset.Mba.Core/BootstrapperCommand.cs
+++ b/src/api/burn/WixToolset.Mba.Core/BootstrapperCommand.cs
@@ -17,7 +17,6 @@ namespace WixToolset.Mba.Core
         /// </summary>
         /// <param name="action"></param>
         /// <param name="display"></param>
-        /// <param name="restart"></param>
         /// <param name="commandLine"></param>
         /// <param name="cmdShow"></param>
         /// <param name="resume"></param>
@@ -30,7 +29,6 @@ namespace WixToolset.Mba.Core
         public BootstrapperCommand(
             LaunchAction action,
             Display display,
-            Restart restart,
             string commandLine,
             int cmdShow,
             ResumeType resume,
@@ -43,7 +41,6 @@ namespace WixToolset.Mba.Core
         {
             this.Action = action;
             this.Display = display;
-            this.Restart = restart;
             this.CommandLine = commandLine;
             this.CmdShow = cmdShow;
             this.Resume = resume;
@@ -62,9 +59,6 @@ namespace WixToolset.Mba.Core
         public Display Display { get; }
 
         /// <inheritdoc/>
-        public Restart Restart { get; }
-
-        /// <inheritdoc/>
         public string CommandLine { get; }
 
         /// <inheritdoc/>
@@ -97,6 +91,7 @@ namespace WixToolset.Mba.Core
             var args = ParseCommandLineToArgs(this.CommandLine);
             var unknownArgs = new List<string>();
             var variables = new List<KeyValuePair<string, string>>();
+            var restart = Restart.Unknown;
 
             foreach (var arg in args)
             {
@@ -104,7 +99,25 @@ namespace WixToolset.Mba.Core
 
                 if (arg[0] == '-' || arg[0] == '/')
                 {
-                    unknownArg = true;
+                    var parameter = arg.Substring(1).ToLowerInvariant();
+                    switch (parameter)
+                    {
+                        case "norestart":
+                            if (restart == Restart.Unknown)
+                            {
+                                restart = Restart.Never;
+                            }
+                            break;
+                        case "forcerestart":
+                            if (restart == Restart.Unknown)
+                            {
+                                restart = Restart.Always;
+                            }
+                            break;
+                        default:
+                            unknownArg = true;
+                            break;
+                    }
                 }
                 else
                 {
@@ -127,8 +140,14 @@ namespace WixToolset.Mba.Core
                 }
             }
 
+            if (restart == Restart.Unknown)
+            {
+                restart = this.Display < Display.Full ? Restart.Automatic : Restart.Prompt;
+            }
+
             return new MbaCommand
             {
+                Restart = restart,
                 UnknownCommandLineArgs = unknownArgs.ToArray(),
                 Variables = variables.ToArray(),
             };
diff --git a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplicationFactory.cs b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplicationFactory.cs
index 0f9193d0..64e25ff4 100644
--- a/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplicationFactory.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IBootstrapperApplicationFactory.cs
@@ -35,7 +35,6 @@ namespace WixToolset.Mba.Core
         [MarshalAs(UnmanagedType.I4)] internal int cbSize;
         [MarshalAs(UnmanagedType.U4)] private readonly LaunchAction action;
         [MarshalAs(UnmanagedType.U4)] private readonly Display display;
-        [MarshalAs(UnmanagedType.U4)] private readonly Restart restart;
         private readonly IntPtr wzCommandLine;
         [MarshalAs(UnmanagedType.I4)] private readonly int nCmdShow;
         [MarshalAs(UnmanagedType.U4)] private readonly ResumeType resume;
@@ -51,7 +50,6 @@ namespace WixToolset.Mba.Core
             return new BootstrapperCommand(
                 this.action,
                 this.display,
-                this.restart,
                 Marshal.PtrToStringUni(this.wzCommandLine),
                 this.nCmdShow,
                 this.resume,
diff --git a/src/api/burn/WixToolset.Mba.Core/IBootstrapperCommand.cs b/src/api/burn/WixToolset.Mba.Core/IBootstrapperCommand.cs
index b7baa55c..f583e619 100644
--- a/src/api/burn/WixToolset.Mba.Core/IBootstrapperCommand.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IBootstrapperCommand.cs
@@ -20,11 +20,6 @@ namespace WixToolset.Mba.Core
         Display Display { get; }
 
         /// <summary>
-        /// Gets the action to perform if a reboot is required.
-        /// </summary>
-        Restart Restart { get; }
-
-        /// <summary>
         /// Gets the command line arguments.
         /// </summary>
         /// <returns>
diff --git a/src/api/burn/WixToolset.Mba.Core/IMbaCommand.cs b/src/api/burn/WixToolset.Mba.Core/IMbaCommand.cs
index a3ad68d8..495b2f44 100644
--- a/src/api/burn/WixToolset.Mba.Core/IMbaCommand.cs
+++ b/src/api/burn/WixToolset.Mba.Core/IMbaCommand.cs
@@ -10,6 +10,11 @@ namespace WixToolset.Mba.Core
     public interface IMbaCommand
     {
         /// <summary>
+        /// Gets the action to perform if a reboot is required.
+        /// </summary>
+        Restart Restart { get; }
+
+        /// <summary>
         /// The command line arguments not parsed into <see cref="IBootstrapperCommand"/> or <see cref="IMbaCommand"/>.
         /// </summary>
         string[] UnknownCommandLineArgs { get; }
diff --git a/src/api/burn/WixToolset.Mba.Core/MbaCommand.cs b/src/api/burn/WixToolset.Mba.Core/MbaCommand.cs
index 424cde63..8917a334 100644
--- a/src/api/burn/WixToolset.Mba.Core/MbaCommand.cs
+++ b/src/api/burn/WixToolset.Mba.Core/MbaCommand.cs
@@ -9,6 +9,8 @@ namespace WixToolset.Mba.Core
     /// </summary>
     internal sealed class MbaCommand : IMbaCommand
     {
+        public Restart Restart { get; internal set; }
+
         public string[] UnknownCommandLineArgs { get; internal set; }
 
         public KeyValuePair<string, string>[] Variables { get; internal set; }
diff --git a/src/api/burn/balutil/balinfo.cpp b/src/api/burn/balutil/balinfo.cpp
index f71784a5..2746f49e 100644
--- a/src/api/burn/balutil/balinfo.cpp
+++ b/src/api/burn/balutil/balinfo.cpp
@@ -19,7 +19,7 @@ static HRESULT ParseOverridableVariablesFromXml(
 
 DAPI_(HRESULT) BalInfoParseCommandLine(
     __in BAL_INFO_COMMAND* pCommand,
-    __in LPCWSTR wzCommandLine
+    __in const BOOTSTRAPPER_COMMAND* pBootstrapperCommand
     )
 {
     HRESULT hr = S_OK;
@@ -29,13 +29,13 @@ DAPI_(HRESULT) BalInfoParseCommandLine(
 
     BalInfoUninitializeCommandLine(pCommand);
 
-    if (!wzCommandLine || !*wzCommandLine)
+    if (!pBootstrapperCommand->wzCommandLine || !*pBootstrapperCommand->wzCommandLine)
     {
         ExitFunction();
     }
 
-    hr = AppParseCommandLine(wzCommandLine, &argc, &argv);
-    ExitOnFailure(hr, "Failed to parse command line.");
+    hr = AppParseCommandLine(pBootstrapperCommand->wzCommandLine, &argc, &argv);
+    BalExitOnFailure(hr, "Failed to parse command line.");
 
     for (int i = 0; i < argc; ++i)
     {
@@ -43,7 +43,24 @@ DAPI_(HRESULT) BalInfoParseCommandLine(
 
         if (argv[i][0] == L'-' || argv[i][0] == L'/')
         {
-            fUnknownArg = TRUE;
+            if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, &argv[i][1], -1, L"norestart", -1))
+            {
+                if (BAL_INFO_RESTART_UNKNOWN == pCommand->restart)
+                {
+                    pCommand->restart = BAL_INFO_RESTART_NEVER;
+                }
+            }
+            else if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, &argv[i][1], -1, L"forcerestart", -1))
+            {
+                if (BAL_INFO_RESTART_UNKNOWN == pCommand->restart)
+                {
+                    pCommand->restart = BAL_INFO_RESTART_ALWAYS;
+                }
+            }
+            else
+            {
+                fUnknownArg = TRUE;
+            }
         }
         else
         {
@@ -55,10 +72,10 @@ DAPI_(HRESULT) BalInfoParseCommandLine(
             else
             {
                 hr = MemEnsureArraySizeForNewItems(reinterpret_cast<LPVOID*>(&pCommand->rgVariableNames), pCommand->cVariables, 1, sizeof(LPWSTR), 5);
-                ExitOnFailure(hr, "Failed to ensure size for variable names.");
+                BalExitOnFailure(hr, "Failed to ensure size for variable names.");
 
                 hr = MemEnsureArraySizeForNewItems(reinterpret_cast<LPVOID*>(&pCommand->rgVariableValues), pCommand->cVariables, 1, sizeof(LPWSTR), 5);
-                ExitOnFailure(hr, "Failed to ensure size for variable values.");
+                BalExitOnFailure(hr, "Failed to ensure size for variable values.");
 
                 LPWSTR* psczVariableName = pCommand->rgVariableNames + pCommand->cVariables;
                 LPWSTR* psczVariableValue = pCommand->rgVariableValues + pCommand->cVariables;
@@ -86,6 +103,11 @@ DAPI_(HRESULT) BalInfoParseCommandLine(
     }
 
 LExit:
+    if (BAL_INFO_RESTART_UNKNOWN == pCommand->restart)
+    {
+        pCommand->restart = BOOTSTRAPPER_DISPLAY_FULL > pBootstrapperCommand->display ? BAL_INFO_RESTART_AUTOMATIC : BAL_INFO_RESTART_PROMPT;
+    }
+
     if (argv)
     {
         AppFreeCommandLineArgs(argv);
diff --git a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
index 393987ba..53fa369b 100644
--- a/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
+++ b/src/api/burn/balutil/inc/BalBaseBootstrapperApplication.h
@@ -9,6 +9,7 @@
 #include "IBootstrapperApplication.h"
 
 #include "balutil.h"
+#include "balinfo.h"
 #include "balretry.h"
 
 class CBalBaseBootstrapperApplication : public IBootstrapperApplication
@@ -794,7 +795,7 @@ public: // IBootstrapperApplication
     {
         HRESULT hr = S_OK;
         BOOL fRestartRequired = BOOTSTRAPPER_APPLY_RESTART_REQUIRED == restart;
-        BOOL fShouldBlockRestart = BOOTSTRAPPER_DISPLAY_FULL <= m_display && BOOTSTRAPPER_RESTART_PROMPT >= m_restart;
+        BOOL fShouldBlockRestart = BOOTSTRAPPER_DISPLAY_FULL <= m_display && BAL_INFO_RESTART_PROMPT >= m_BalInfoCommand.restart;
 
         if (fRestartRequired && !fShouldBlockRestart)
         {
@@ -976,6 +977,22 @@ public: // IBootstrapperApplication
         return S_OK;
     }
 
+public: //CBalBaseBootstrapperApplication
+    virtual STDMETHODIMP Initialize(
+        __in const BOOTSTRAPPER_CREATE_ARGS* pCreateArgs
+        )
+    {
+        HRESULT hr = S_OK;
+
+        m_display = pCreateArgs->pCommand->display;
+
+        hr = BalInfoParseCommandLine(&m_BalInfoCommand, pCreateArgs->pCommand);
+        BalExitOnFailure(hr, "Failed to parse command line with balutil.");
+
+    LExit:
+        return hr;
+    }
+
 protected:
     //
     // PromptCancel - prompts the user to close (if not forced).
@@ -1029,20 +1046,19 @@ protected:
 
     CBalBaseBootstrapperApplication(
         __in IBootstrapperEngine* pEngine,
-        __in const BOOTSTRAPPER_CREATE_ARGS* pArgs,
         __in DWORD dwRetryCount = 0,
         __in DWORD dwRetryTimeout = 1000
         )
     {
         m_cReferences = 1;
-        m_display = pArgs->pCommand->display;
-        m_restart = pArgs->pCommand->restart;
+        m_display = BOOTSTRAPPER_DISPLAY_UNKNOWN;
 
         pEngine->AddRef();
         m_pEngine = pEngine;
 
         ::InitializeCriticalSection(&m_csCanceled);
         m_fCanceled = FALSE;
+        m_BalInfoCommand = { };
         m_fApplying = FALSE;
         m_fRollingBack = FALSE;
 
@@ -1054,6 +1070,7 @@ protected:
 
     virtual ~CBalBaseBootstrapperApplication()
     {
+        BalInfoUninitializeCommandLine(&m_BalInfoCommand);
         BalRetryUninitialize();
         ::DeleteCriticalSection(&m_csCanceled);
 
@@ -1064,10 +1081,11 @@ protected:
     CRITICAL_SECTION m_csCanceled;
     BOOL m_fCanceled;
 
+    BAL_INFO_COMMAND m_BalInfoCommand;
+
 private:
     long m_cReferences;
     BOOTSTRAPPER_DISPLAY m_display;
-    BOOTSTRAPPER_RESTART m_restart;
     IBootstrapperEngine* m_pEngine;
 
     BOOL m_fApplying;
diff --git a/src/api/burn/balutil/inc/balinfo.h b/src/api/burn/balutil/inc/balinfo.h
index 07a1cbb7..769becb2 100644
--- a/src/api/burn/balutil/inc/balinfo.h
+++ b/src/api/burn/balutil/inc/balinfo.h
@@ -18,6 +18,15 @@ typedef enum BAL_INFO_PACKAGE_TYPE
     BAL_INFO_PACKAGE_TYPE_BUNDLE_PATCH,
 } BAL_INFO_PACKAGE_TYPE;
 
+typedef enum _BAL_INFO_RESTART
+{
+    BAL_INFO_RESTART_UNKNOWN,
+    BAL_INFO_RESTART_NEVER,
+    BAL_INFO_RESTART_PROMPT,
+    BAL_INFO_RESTART_AUTOMATIC,
+    BAL_INFO_RESTART_ALWAYS,
+} BAL_INFO_RESTART;
+
 typedef enum _BAL_INFO_VARIABLE_COMMAND_LINE_TYPE
 {
     BAL_INFO_VARIABLE_COMMAND_LINE_TYPE_UPPER_CASE,
@@ -85,6 +94,7 @@ typedef struct _BAL_INFO_COMMAND
     DWORD cVariables;
     LPWSTR* rgVariableNames;
     LPWSTR* rgVariableValues;
+    BAL_INFO_RESTART restart;
 } BAL_INFO_COMMAND;
 
 
@@ -94,7 +104,7 @@ typedef struct _BAL_INFO_COMMAND
 ********************************************************************/
 HRESULT DAPI BalInfoParseCommandLine(
     __in BAL_INFO_COMMAND* pCommand,
-    __in LPCWSTR wzCommandLine
+    __in const BOOTSTRAPPER_COMMAND* pBootstrapperCommand
     );
 
 
diff --git a/src/api/burn/bextutil/inc/BextBaseBundleExtension.h b/src/api/burn/bextutil/inc/BextBaseBundleExtension.h
index 69c338e4..a302702e 100644
--- a/src/api/burn/bextutil/inc/BextBaseBundleExtension.h
+++ b/src/api/burn/bextutil/inc/BextBaseBundleExtension.h
@@ -80,7 +80,7 @@ public: // IBundleExtension
 public: //CBextBaseBundleExtension
     virtual STDMETHODIMP Initialize(
         __in const BUNDLE_EXTENSION_CREATE_ARGS* pCreateArgs
-    )
+        )
     {
         HRESULT hr = S_OK;
 
diff --git a/src/api/burn/test/BalUtilUnitTest/TestBootstrapperApplication.cpp b/src/api/burn/test/BalUtilUnitTest/TestBootstrapperApplication.cpp
index 13d22e72..daa1d690 100644
--- a/src/api/burn/test/BalUtilUnitTest/TestBootstrapperApplication.cpp
+++ b/src/api/burn/test/BalUtilUnitTest/TestBootstrapperApplication.cpp
@@ -8,9 +8,8 @@ class CTestBootstrapperApplication : public CBalBaseBootstrapperApplication
 {
 public:
     CTestBootstrapperApplication(
-        __in IBootstrapperEngine* pEngine,
-        __in const BOOTSTRAPPER_CREATE_ARGS* pArgs
-        ) : CBalBaseBootstrapperApplication(pEngine, pArgs)
+        __in IBootstrapperEngine* pEngine
+        ) : CBalBaseBootstrapperApplication(pEngine)
     {
     }
 };
@@ -25,9 +24,12 @@ HRESULT CreateBootstrapperApplication(
     HRESULT hr = S_OK;
     CTestBootstrapperApplication* pApplication = NULL;
 
-    pApplication = new CTestBootstrapperApplication(pEngine, pArgs);
+    pApplication = new CTestBootstrapperApplication(pEngine);
     ExitOnNull(pApplication, hr, E_OUTOFMEMORY, "Failed to create new test bootstrapper application object.");
 
+    hr = pApplication->Initialize(pArgs);
+    ExitOnFailure(hr, "CTestBootstrapperApplication initialization failed.");
+
     pResults->pfnBootstrapperApplicationProc = BalBaseBootstrapperApplicationProc;
     pResults->pvBootstrapperApplicationProcContext = pApplication;
     *ppApplication = pApplication;
diff --git a/src/api/burn/test/WixToolsetTest.Mba.Core/BaseBootstrapperApplicationFactoryFixture.cs b/src/api/burn/test/WixToolsetTest.Mba.Core/BaseBootstrapperApplicationFactoryFixture.cs
index b49b3927..8705ed13 100644
--- a/src/api/burn/test/WixToolsetTest.Mba.Core/BaseBootstrapperApplicationFactoryFixture.cs
+++ b/src/api/burn/test/WixToolsetTest.Mba.Core/BaseBootstrapperApplicationFactoryFixture.cs
@@ -101,7 +101,6 @@ namespace WixToolsetTest.Mba.Core
             public int cbSize;
             public LaunchAction action;
             public Display display;
-            public Restart restart;
             [MarshalAs(UnmanagedType.LPWStr)] public string wzCommandLine;
             public int nCmdShow;
             public ResumeType resume;
diff --git a/src/burn/engine/core.cpp b/src/burn/engine/core.cpp
index 8b97952c..f71103a8 100644
--- a/src/burn/engine/core.cpp
+++ b/src/burn/engine/core.cpp
@@ -924,7 +924,6 @@ extern "C" HRESULT CoreRecreateCommandLine(
     __deref_inout_z LPWSTR* psczCommandLine,
     __in BOOTSTRAPPER_ACTION action,
     __in BOOTSTRAPPER_DISPLAY display,
-    __in BOOTSTRAPPER_RESTART restart,
     __in BOOTSTRAPPER_RELATION_TYPE relationType,
     __in BOOL fPassthrough,
     __in_z_opt LPCWSTR wzActiveParent,
@@ -965,17 +964,6 @@ extern "C" HRESULT CoreRecreateCommandLine(
     }
     ExitOnFailure(hr, "Failed to append action state to command-line");
 
-    switch (restart)
-    {
-    case BOOTSTRAPPER_RESTART_ALWAYS:
-        hr = StrAllocConcat(psczCommandLine, L" /forcerestart", 0);
-        break;
-    case BOOTSTRAPPER_RESTART_NEVER:
-        hr = StrAllocConcat(psczCommandLine, L" /norestart", 0);
-        break;
-    }
-    ExitOnFailure(hr, "Failed to append restart state to command-line");
-
     if (wzActiveParent)
     {
         if (*wzActiveParent)
@@ -1246,32 +1234,10 @@ extern "C" HRESULT CoreParseCommandLine(
                      CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, &argv[i][1], -1, L"silent", -1))
             {
                 pCommand->display = BOOTSTRAPPER_DISPLAY_NONE;
-
-                if (BOOTSTRAPPER_RESTART_UNKNOWN == pCommand->restart)
-                {
-                    pCommand->restart = BOOTSTRAPPER_RESTART_AUTOMATIC;
-                }
             }
             else if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, &argv[i][1], -1, L"passive", -1))
             {
                 pCommand->display = BOOTSTRAPPER_DISPLAY_PASSIVE;
-
-                if (BOOTSTRAPPER_RESTART_UNKNOWN == pCommand->restart)
-                {
-                    pCommand->restart = BOOTSTRAPPER_RESTART_AUTOMATIC;
-                }
-            }
-            else if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, &argv[i][1], -1, L"norestart", -1))
-            {
-                pCommand->restart = BOOTSTRAPPER_RESTART_NEVER;
-            }
-            else if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, &argv[i][1], -1, L"forcerestart", -1))
-            {
-                pCommand->restart = BOOTSTRAPPER_RESTART_ALWAYS;
-            }
-            else if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, &argv[i][1], -1, L"promptrestart", -1))
-            {
-                pCommand->restart = BOOTSTRAPPER_RESTART_PROMPT;
             }
             else if (CSTR_EQUAL == ::CompareStringW(LOCALE_INVARIANT, NORM_IGNORECASE, &argv[i][1], -1, L"layout", -1))
             {
@@ -1653,11 +1619,6 @@ extern "C" HRESULT CoreParseCommandLine(
         pCommand->display = BOOTSTRAPPER_DISPLAY_FULL;
     }
 
-    if (BOOTSTRAPPER_RESTART_UNKNOWN == pCommand->restart)
-    {
-        pCommand->restart = BOOTSTRAPPER_RESTART_PROMPT;
-    }
-
 LExit:
     if (fInvalidCommandLine)
     {
diff --git a/src/burn/engine/core.h b/src/burn/engine/core.h
index 0a92d64f..3fa1f04a 100644
--- a/src/burn/engine/core.h
+++ b/src/burn/engine/core.h
@@ -210,7 +210,6 @@ HRESULT CoreRecreateCommandLine(
     __deref_inout_z LPWSTR* psczCommandLine,
     __in BOOTSTRAPPER_ACTION action,
     __in BOOTSTRAPPER_DISPLAY display,
-    __in BOOTSTRAPPER_RESTART restart,
     __in BOOTSTRAPPER_RELATION_TYPE relationType,
     __in BOOL fPassthrough,
     __in_z_opt LPCWSTR wzActiveParent,
diff --git a/src/burn/engine/externalengine.cpp b/src/burn/engine/externalengine.cpp
index 409353e4..a5f75bd4 100644
--- a/src/burn/engine/externalengine.cpp
+++ b/src/burn/engine/externalengine.cpp
@@ -295,7 +295,7 @@ HRESULT ExternalEngineSetUpdate(
     {
         UpdateUninitialize(&pEngineState->update);
 
-        hr = CoreRecreateCommandLine(&sczCommandline, BOOTSTRAPPER_ACTION_INSTALL, pEngineState->command.display, pEngineState->command.restart, BOOTSTRAPPER_RELATION_NONE, FALSE, pEngineState->registration.sczActiveParent, pEngineState->registration.sczAncestors, NULL, pEngineState->command.wzCommandLine);
+        hr = CoreRecreateCommandLine(&sczCommandline, BOOTSTRAPPER_ACTION_INSTALL, pEngineState->command.display, BOOTSTRAPPER_RELATION_NONE, FALSE, pEngineState->registration.sczActiveParent, pEngineState->registration.sczAncestors, NULL, pEngineState->command.wzCommandLine);
         ExitOnFailure(hr, "Failed to recreate command-line for update bundle.");
 
         // Bundles would fail to use the downloaded update bundle, as the running bundle would be one of the search paths.
diff --git a/src/burn/engine/plan.cpp b/src/burn/engine/plan.cpp
index d1f07508..3d6fc65f 100644
--- a/src/burn/engine/plan.cpp
+++ b/src/burn/engine/plan.cpp
@@ -1783,7 +1783,7 @@ extern "C" HRESULT PlanSetResumeCommand(
     HRESULT hr = S_OK;
 
     // build the resume command-line.
-    hr = CoreRecreateCommandLine(&pRegistration->sczResumeCommandLine, action, pCommand->display, pCommand->restart, pCommand->relationType, pCommand->fPassthrough, pRegistration->sczActiveParent, pRegistration->sczAncestors, pLog->sczPath, pCommand->wzCommandLine);
+    hr = CoreRecreateCommandLine(&pRegistration->sczResumeCommandLine, action, pCommand->display, pCommand->relationType, pCommand->fPassthrough, pRegistration->sczActiveParent, pRegistration->sczAncestors, pLog->sczPath, pCommand->wzCommandLine);
     ExitOnFailure(hr, "Failed to recreate resume command-line.");
 
 LExit:
diff --git a/src/burn/engine/pseudobundle.cpp b/src/burn/engine/pseudobundle.cpp
index 1b2eae75..351fe044 100644
--- a/src/burn/engine/pseudobundle.cpp
+++ b/src/burn/engine/pseudobundle.cpp
@@ -205,7 +205,7 @@ extern "C" HRESULT PseudoBundleInitializePassthrough(
 
     // No matter the operation, we're passing the same command-line. That's what makes
     // this a passthrough bundle.
-    hr = CoreRecreateCommandLine(&sczArguments, pCommand->action, pCommand->display, pCommand->restart, pCommand->relationType, TRUE, wzActiveParent, wzAncestors, wzAppendLogPath, pCommand->wzCommandLine);
+    hr = CoreRecreateCommandLine(&sczArguments, pCommand->action, pCommand->display, pCommand->relationType, TRUE, wzActiveParent, wzAncestors, wzAppendLogPath, pCommand->wzCommandLine);
     ExitOnFailure(hr, "Failed to recreate command-line arguments.");
 
     hr = StrAllocString(&pPassthroughPackage->Exe.sczInstallArguments, sczArguments, 0);
diff --git a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
index 9f0b0082..bbe926f1 100644
--- a/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
+++ b/src/ext/Bal/wixstdba/WixStandardBootstrapperApplication.cpp
@@ -206,7 +206,7 @@ public: // IBootstrapperApplication
         m_hUiThread = ::CreateThread(NULL, 0, UiThreadProc, this, 0, &dwUIThreadId);
         if (!m_hUiThread)
         {
-            ExitWithLastError(hr, "Failed to create UI thread.");
+            BalExitWithLastError(hr, "Failed to create UI thread.");
         }
 
     LExit:
@@ -1100,10 +1100,10 @@ public: // IBootstrapperApplication
         m_fRestartRequired = BOOTSTRAPPER_APPLY_RESTART_NONE != restart;
         BalSetStringVariable(WIXSTDBA_VARIABLE_RESTART_REQUIRED, m_fRestartRequired ? L"1" : NULL, FALSE);
 
-        m_fShouldRestart = m_fRestartRequired && BOOTSTRAPPER_RESTART_NEVER < m_command.restart;
+        m_fShouldRestart = m_fRestartRequired && BAL_INFO_RESTART_NEVER < m_BalInfoCommand.restart;
 
         // Automatically restart if we're not displaying a UI or the command line said to always allow restarts.
-        m_fAllowRestart = m_fShouldRestart && (BOOTSTRAPPER_DISPLAY_FULL > m_command.display || BOOTSTRAPPER_RESTART_PROMPT < m_command.restart);
+        m_fAllowRestart = m_fShouldRestart && (BOOTSTRAPPER_DISPLAY_FULL > m_command.display || BAL_INFO_RESTART_PROMPT < m_BalInfoCommand.restart);
 
         if (m_fPrereq)
         {
@@ -1986,6 +1986,62 @@ private: // privates
         m_pfnBAFunctionsProc(BA_FUNCTIONS_MESSAGE_ONCACHEPAYLOADEXTRACTPROGRESS, pArgs, pResults, m_pvBAFunctionsProcContext);
     }
 
+
+public: //CBalBaseBootstrapperApplication
+    virtual STDMETHODIMP Initialize(
+        __in const BOOTSTRAPPER_CREATE_ARGS* pCreateArgs
+        )
+    {
+        HRESULT hr = S_OK;
+        LONGLONG llInstalled = 0;
+
+        hr = __super::Initialize(pCreateArgs);
+        BalExitOnFailure(hr, "CBalBaseBootstrapperApplication initialization failed.");
+
+        memcpy_s(&m_command, sizeof(m_command), pCreateArgs->pCommand, sizeof(BOOTSTRAPPER_COMMAND));
+        memcpy_s(&m_createArgs, sizeof(m_createArgs), pCreateArgs, sizeof(BOOTSTRAPPER_CREATE_ARGS));
+        m_createArgs.pCommand = &m_command;
+
+        if (m_fPrereq)
+        {
+            // Pre-req BA should only show help or do an install (to launch the Managed BA which can then do the right action).
+            if (BOOTSTRAPPER_ACTION_HELP != m_command.action)
+            {
+                m_command.action = BOOTSTRAPPER_ACTION_INSTALL;
+            }
+        }
+        else // maybe modify the action state if the bundle is or is not already installed.
+        {
+            hr = BalGetNumericVariable(L"WixBundleInstalled", &llInstalled);
+            if (SUCCEEDED(hr) && BOOTSTRAPPER_RESUME_TYPE_REBOOT != m_command.resumeType && llInstalled && BOOTSTRAPPER_ACTION_INSTALL == m_command.action)
+            {
+                m_command.action = BOOTSTRAPPER_ACTION_MODIFY;
+            }
+            else if (!llInstalled && (BOOTSTRAPPER_ACTION_MODIFY == m_command.action || BOOTSTRAPPER_ACTION_REPAIR == m_command.action))
+            {
+                m_command.action = BOOTSTRAPPER_ACTION_INSTALL;
+            }
+        }
+
+        // When resuming from restart doing some install-like operation, try to find the package that forced the
+        // restart. We'll use this information during planning.
+        if (BOOTSTRAPPER_RESUME_TYPE_REBOOT == m_command.resumeType && BOOTSTRAPPER_ACTION_UNINSTALL < m_command.action)
+        {
+            // Ensure the forced restart package variable is null when it is an empty string.
+            hr = BalGetStringVariable(L"WixBundleForcedRestartPackage", &m_sczAfterForcedRestartPackage);
+            if (FAILED(hr) || !m_sczAfterForcedRestartPackage || !*m_sczAfterForcedRestartPackage)
+            {
+                ReleaseNullStr(m_sczAfterForcedRestartPackage);
+            }
+        }
+
+        hr = S_OK;
+
+    LExit:
+        return hr;
+    }
+
+private:
     //
     // UiThreadProc - entrypoint for UI thread.
     //
@@ -2172,9 +2228,6 @@ private: // privates
         LPWSTR* argv = NULL;
         BOOL fUnknownArg = FALSE;
 
-        hr = BalInfoParseCommandLine(&m_BalInfoCommand, m_command.wzCommandLine);
-        BalExitOnFailure(hr, "Failed to parse command line with balutil.");
-
         argc = m_BalInfoCommand.cUnknownArgs;
         argv = m_BalInfoCommand.rgUnknownArgs;
 
@@ -3139,7 +3192,7 @@ private: // privates
                     BOOL fLaunchTargetExists = FALSE;
                     if (m_fShouldRestart)
                     {
-                        if (BOOTSTRAPPER_RESTART_PROMPT == m_command.restart)
+                        if (BAL_INFO_RESTART_PROMPT == m_BalInfoCommand.restart)
                         {
                             fEnableRestartButton = TRUE;
                         }
@@ -3240,7 +3293,7 @@ private: // privates
 
                     if (m_fShouldRestart)
                     {
-                        if (BOOTSTRAPPER_RESTART_PROMPT == m_command.restart)
+                        if (BAL_INFO_RESTART_PROMPT == m_BalInfoCommand.restart)
                         {
                             fEnableRestartButton = TRUE;
                         }
@@ -3832,57 +3885,20 @@ public:
         __in HMODULE hModule,
         __in BOOL fPrereq,
         __in HRESULT hrHostInitialization,
-        __in IBootstrapperEngine* pEngine,
-        __in const BOOTSTRAPPER_CREATE_ARGS* pArgs
-        ) : CBalBaseBootstrapperApplication(pEngine, pArgs, 3, 3000)
+        __in IBootstrapperEngine* pEngine
+        ) : CBalBaseBootstrapperApplication(pEngine, 3, 3000)
     {
         m_hModule = hModule;
-        memcpy_s(&m_command, sizeof(m_command), pArgs->pCommand, sizeof(BOOTSTRAPPER_COMMAND));
-        memcpy_s(&m_createArgs, sizeof(m_createArgs), pArgs, sizeof(BOOTSTRAPPER_CREATE_ARGS));
-        m_createArgs.pCommand = &m_command;
-
-        if (fPrereq)
-        {
-            // Pre-req BA should only show help or do an install (to launch the Managed BA which can then do the right action).
-            if (BOOTSTRAPPER_ACTION_HELP != m_command.action)
-            {
-                m_command.action = BOOTSTRAPPER_ACTION_INSTALL;
-            }
-        }
-        else // maybe modify the action state if the bundle is or is not already installed.
-        {
-            LONGLONG llInstalled = 0;
-            HRESULT hr = BalGetNumericVariable(L"WixBundleInstalled", &llInstalled);
-            if (SUCCEEDED(hr) && BOOTSTRAPPER_RESUME_TYPE_REBOOT != m_command.resumeType && 0 < llInstalled && BOOTSTRAPPER_ACTION_INSTALL == m_command.action)
-            {
-                m_command.action = BOOTSTRAPPER_ACTION_MODIFY;
-            }
-            else if (0 == llInstalled && (BOOTSTRAPPER_ACTION_MODIFY == m_command.action || BOOTSTRAPPER_ACTION_REPAIR == m_command.action))
-            {
-                m_command.action = BOOTSTRAPPER_ACTION_INSTALL;
-            }
-        }
+        m_command = { };
+        m_createArgs = { };
 
         m_plannedAction = BOOTSTRAPPER_ACTION_UNKNOWN;
 
-        // When resuming from restart doing some install-like operation, try to find the package that forced the
-        // restart. We'll use this information during planning.
         m_sczAfterForcedRestartPackage = NULL;
 
-        if (BOOTSTRAPPER_RESUME_TYPE_REBOOT == m_command.resumeType && BOOTSTRAPPER_ACTION_UNINSTALL < m_command.action)
-        {
-            // Ensure the forced restart package variable is null when it is an empty string.
-            HRESULT hr = BalGetStringVariable(L"WixBundleForcedRestartPackage", &m_sczAfterForcedRestartPackage);
-            if (FAILED(hr) || !m_sczAfterForcedRestartPackage || !*m_sczAfterForcedRestartPackage)
-            {
-                ReleaseNullStr(m_sczAfterForcedRestartPackage);
-            }
-        }
-
         m_pWixLoc = NULL;
-        memset(&m_Bundle, 0, sizeof(m_Bundle));
-        memset(&m_BalInfoCommand, 0, sizeof(m_BalInfoCommand));
-        memset(&m_Conditions, 0, sizeof(m_Conditions));
+        m_Bundle = { };
+        m_Conditions = { };
         m_sczConfirmCloseMessage = NULL;
         m_sczFailedMessage = NULL;
 
@@ -3947,7 +3963,6 @@ public:
         ReleaseStr(m_sczConfirmCloseMessage);
         BalConditionsUninitialize(&m_Conditions);
         BalInfoUninitialize(&m_Bundle);
-        BalInfoUninitializeCommandLine(&m_BalInfoCommand);
         LocFree(m_pWixLoc);
 
         ReleaseStr(m_sczLanguage);
@@ -3980,7 +3995,6 @@ private:
 
     WIX_LOCALIZATION* m_pWixLoc;
     BAL_INFO_BUNDLE m_Bundle;
-    BAL_INFO_COMMAND m_BalInfoCommand;
     BAL_CONDITIONS m_Conditions;
     LPWSTR m_sczFailedMessage;
     LPWSTR m_sczConfirmCloseMessage;
@@ -4049,8 +4063,11 @@ HRESULT CreateBootstrapperApplication(
         BalExitOnFailure(hr = E_INVALIDARG, "Engine requested Unknown display type.");
     }
 
-    pApplication = new CWixStandardBootstrapperApplication(hModule, fPrereq, hrHostInitialization, pEngine, pArgs);
-    ExitOnNull(pApplication, hr, E_OUTOFMEMORY, "Failed to create new standard bootstrapper application object.");
+    pApplication = new CWixStandardBootstrapperApplication(hModule, fPrereq, hrHostInitialization, pEngine);
+    BalExitOnNull(pApplication, hr, E_OUTOFMEMORY, "Failed to create new standard bootstrapper application object.");
+
+    hr = pApplication->Initialize(pArgs);
+    ExitOnFailure(hr, "CWixStandardBootstrapperApplication initialization failed.");
 
     pResults->pfnBootstrapperApplicationProc = BalBaseBootstrapperApplicationProc;
     pResults->pvBootstrapperApplicationProcContext = pApplication;
diff --git a/src/ext/Util/be/UtilBundleExtension.cpp b/src/ext/Util/be/UtilBundleExtension.cpp
index 2ac842a5..110599fa 100644
--- a/src/ext/Util/be/UtilBundleExtension.cpp
+++ b/src/ext/Util/be/UtilBundleExtension.cpp
@@ -27,7 +27,7 @@ public: //CBextBaseBundleExtension
         IXMLDOMDocument* pixdManifest = NULL;
         IXMLDOMNode* pixnBundleExtension = NULL;
 
-        hr = CBextBaseBundleExtension::Initialize(pCreateArgs);
+        hr = __super::Initialize(pCreateArgs);
         ExitOnFailure(hr, "CBextBaseBundleExtension initialization failed.");
 
         hr = XmlLoadDocumentFromFile(m_sczBundleExtensionDataPath, &pixdManifest);
@@ -76,7 +76,7 @@ HRESULT UtilBundleExtensionCreate(
     ExitOnNull(pExtension, hr, E_OUTOFMEMORY, "Failed to create new CWixUtilBundleExtension.");
 
     hr = pExtension->Initialize(pArgs);
-    ExitOnFailure(hr, "CWixUtilBundleExtension initialization failed");
+    ExitOnFailure(hr, "CWixUtilBundleExtension initialization failed.");
 
     *ppBundleExtension = pExtension;
     pExtension = NULL;