Make the BA responsible for parsing restart prompt behavior.

Fixes #4975

3 files changed, 63 insertions, 13 deletions

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
     );