gh-115554: Improved logic for handling multiple existing py.exe launcher installs (GH-115793)

3 files changed, 134 insertions, 74 deletions

diff --git a/Tools/msi/bundle/Default.wxl b/Tools/msi/bundle/Default.wxl
index 1540f05..0014204 100644
--- a/Tools/msi/bundle/Default.wxl
+++ b/Tools/msi/bundle/Default.wxl
@@ -88,6 +88,7 @@ Select Customize to review current options.</String>
   <String Id="InstallAllUsersLabel">Install Python [ShortVersion] for &amp;all users</String>
   <String Id="InstallLauncherAllUsersLabel">for &amp;all users (requires admin privileges)</String>
   <String Id="ShortInstallLauncherAllUsersLabel">Use admin privi&amp;leges when installing py.exe</String>
+  <String Id="ShortInstallLauncherBlockedLabel">Python Launcher is already installed</String>
   <String Id="PrecompileLabel">&amp;Precompile standard library</String>
   <String Id="Include_symbolsLabel">Download debugging &amp;symbols</String>
   <String Id="Include_debugLabel">Download debu&amp;g binaries (requires VS 2017 or later)</String>
diff --git a/Tools/msi/bundle/bootstrap/PythonBootstrapperApplication.cpp b/Tools/msi/bundle/bootstrap/PythonBootstrapperApplication.cpp
index 3a17ffb..e0e179e 100644
--- a/Tools/msi/bundle/bootstrap/PythonBootstrapperApplication.cpp
+++ b/Tools/msi/bundle/bootstrap/PythonBootstrapperApplication.cpp
@@ -442,6 +442,14 @@ class PythonBootstrapperApplication : public CBalBaseBootstrapperApplication {
         ThemeControlElevates(_theme, ID_INSTALL_BUTTON, elevated);
         ThemeControlElevates(_theme, ID_INSTALL_SIMPLE_BUTTON, elevated);
         ThemeControlElevates(_theme, ID_INSTALL_UPGRADE_BUTTON, elevated);
+
+        LONGLONG blockedLauncher;
+        if (SUCCEEDED(BalGetNumericVariable(L"BlockedLauncher", &blockedLauncher)) && blockedLauncher) {
+            LOC_STRING *pLocString = nullptr;
+            if (SUCCEEDED(LocGetString(_wixLoc, L"#(loc.ShortInstallLauncherBlockedLabel)", &pLocString)) && pLocString) {
+                ThemeSetTextControl(_theme, ID_INSTALL_LAUNCHER_ALL_USERS_CHECKBOX, pLocString->wzText);
+            }
+        }
     }
 
     void Custom1Page_Show() {
@@ -718,25 +726,67 @@ public: // IBootstrapperApplication
         __in DWORD64 /*dw64Version*/,
         __in BOOTSTRAPPER_RELATED_OPERATION operation
     ) {
-        if (BOOTSTRAPPER_RELATED_OPERATION_MAJOR_UPGRADE == operation && 
-            (CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, wzPackageId, -1, L"launcher_AllUsers", -1) ||
-             CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, wzPackageId, -1, L"launcher_JustForMe", -1))) {
-            auto hr = LoadAssociateFilesStateFromKey(_engine, fPerMachine ? HKEY_LOCAL_MACHINE : HKEY_CURRENT_USER);
-            if (hr == S_OK) {
-                _engine->SetVariableNumeric(L"AssociateFiles", 1);
-            } else if (hr == S_FALSE) {
-                _engine->SetVariableNumeric(L"AssociateFiles", 0);
-            } else if (FAILED(hr)) {
-                BalLog(BOOTSTRAPPER_LOG_LEVEL_ERROR, "Failed to load AssociateFiles state: error code 0x%08X", hr);
+        // Only check launcher_AllUsers because we'll find the same packages
+        // twice if we check launcher_JustForMe as well.
+        if (CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, wzPackageId, -1, L"launcher_AllUsers", -1)) {
+            BalLog(BOOTSTRAPPER_LOG_LEVEL_STANDARD, "Detected existing launcher install");
+
+            LONGLONG blockedLauncher, detectedLauncher;
+            if (FAILED(BalGetNumericVariable(L"BlockedLauncher", &blockedLauncher))) {
+                blockedLauncher = 0;
+            }
+
+            // Get the prior DetectedLauncher value so we can see if we've
+            // detected more than one, and then update the stored variable
+            // (we use the original value later on via the local).
+            if (FAILED(BalGetNumericVariable(L"DetectedLauncher", &detectedLauncher))) {
+                detectedLauncher = 0;
+            }
+            if (!detectedLauncher) {
+                _engine->SetVariableNumeric(L"DetectedLauncher", 1);
+            }
+
+            if (blockedLauncher) {
+                // Nothing else to do, we're already blocking
+            }
+            else if (BOOTSTRAPPER_RELATED_OPERATION_DOWNGRADE == operation) {
+                // Found a higher version, so we can't install ours.
+                BalLog(BOOTSTRAPPER_LOG_LEVEL_ERROR, "Higher version launcher has been detected.");
+                BalLog(BOOTSTRAPPER_LOG_LEVEL_ERROR, "Launcher will not be installed");
+                _engine->SetVariableNumeric(L"BlockedLauncher", 1);
+            }
+            else if (detectedLauncher) {
+                if (!blockedLauncher) {
+                    BalLog(BOOTSTRAPPER_LOG_LEVEL_ERROR, "Multiple launcher installs have been detected.");
+                    BalLog(BOOTSTRAPPER_LOG_LEVEL_ERROR, "No launcher will be installed or upgraded until one has been removed.");
+                    _engine->SetVariableNumeric(L"BlockedLauncher", 1);
+                }
             }
+            else if (BOOTSTRAPPER_RELATED_OPERATION_MAJOR_UPGRADE == operation) {
+                // Found an older version, so let's run the equivalent as an upgrade
+                // This overrides "unknown" all users options, but will leave alone
+                // any that have already been set/detected.
+                // User can deselect the option to include the launcher, but cannot
+                // change it from the current per user/machine setting.
+                LONGLONG includeLauncher, includeLauncherAllUsers;
+                if (FAILED(BalGetNumericVariable(L"Include_launcher", &includeLauncher))) {
+                    includeLauncher = -1;
+                }
+                if (FAILED(BalGetNumericVariable(L"InstallLauncherAllUsers", &includeLauncherAllUsers))) {
+                    includeLauncherAllUsers = -1;
+                }
 
-            LONGLONG includeLauncher;
-            if (FAILED(BalGetNumericVariable(L"Include_launcher", &includeLauncher))
-                || includeLauncher == -1) {
-                _engine->SetVariableNumeric(L"Include_launcher", 1);
-                _engine->SetVariableNumeric(L"InstallLauncherAllUsers", fPerMachine ? 1 : 0);
+                if (includeLauncher < 0) {
+                    _engine->SetVariableNumeric(L"Include_launcher", 1);
+                }
+                if (includeLauncherAllUsers < 0) {
+                    _engine->SetVariableNumeric(L"InstallLauncherAllUsers", fPerMachine ? 1 : 0);
+                } else if (includeLauncherAllUsers != fPerMachine ? 1 : 0) {
+                    // Requested AllUsers option is inconsistent, so block
+                    _engine->SetVariableNumeric(L"BlockedLauncher", 1);
+                }
+                _engine->SetVariableNumeric(L"DetectedOldLauncher", 1);
             }
-            _engine->SetVariableNumeric(L"DetectedOldLauncher", 1);
         }
         return CheckCanceled() ? IDCANCEL : IDNOACTION;
     }
@@ -784,48 +834,7 @@ public: // IBootstrapperApplication
         __in LPCWSTR wzPackageId,
         __in HRESULT hrStatus,
         __in BOOTSTRAPPER_PACKAGE_STATE state
-    ) {
-        if (FAILED(hrStatus)) {
-            return;
-        }
-
-        BOOL detectedLauncher = FALSE;
-        HKEY hkey = HKEY_LOCAL_MACHINE;
-        if (CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, wzPackageId, -1, L"launcher_AllUsers", -1)) {
-            if (BOOTSTRAPPER_PACKAGE_STATE_PRESENT == state || BOOTSTRAPPER_PACKAGE_STATE_OBSOLETE == state) {
-                detectedLauncher = TRUE;
-                _engine->SetVariableNumeric(L"InstallLauncherAllUsers", 1);
-            }
-        } else if (CSTR_EQUAL == ::CompareStringW(LOCALE_NEUTRAL, 0, wzPackageId, -1, L"launcher_JustForMe", -1)) {
-            if (BOOTSTRAPPER_PACKAGE_STATE_PRESENT == state || BOOTSTRAPPER_PACKAGE_STATE_OBSOLETE == state) {
-                detectedLauncher = TRUE;
-                _engine->SetVariableNumeric(L"InstallLauncherAllUsers", 0);
-            }
-        }
-
-        LONGLONG includeLauncher;
-        if (SUCCEEDED(BalGetNumericVariable(L"Include_launcher", &includeLauncher))
-            && includeLauncher != -1) {
-            detectedLauncher = FALSE;
-        }
-
-        if (detectedLauncher) {
-            /* When we detect the current version of the launcher. */
-            _engine->SetVariableNumeric(L"Include_launcher", 1);
-            _engine->SetVariableNumeric(L"DetectedLauncher", 1);
-            _engine->SetVariableString(L"Include_launcherState", L"disable");
-            _engine->SetVariableString(L"InstallLauncherAllUsersState", L"disable");
-
-            auto hr = LoadAssociateFilesStateFromKey(_engine, hkey);
-            if (hr == S_OK) {
-                _engine->SetVariableNumeric(L"AssociateFiles", 1);
-            } else if (hr == S_FALSE) {
-                _engine->SetVariableNumeric(L"AssociateFiles", 0);
-            } else if (FAILED(hr)) {
-                BalLog(BOOTSTRAPPER_LOG_LEVEL_ERROR, "Failed to load AssociateFiles state: error code 0x%08X", hr);
-            }
-        }
-    }
+    ) { }
 
 
     virtual STDMETHODIMP_(void) OnDetectComplete(__in HRESULT hrStatus) {
@@ -835,19 +844,67 @@ public: // IBootstrapperApplication
         }
 
         if (SUCCEEDED(hrStatus)) {
-            LONGLONG includeLauncher;
-            if (SUCCEEDED(BalGetNumericVariable(L"Include_launcher", &includeLauncher))
-                && includeLauncher == -1) {
-                if (BOOTSTRAPPER_ACTION_LAYOUT == _command.action ||
-                    (BOOTSTRAPPER_ACTION_INSTALL == _command.action && !_upgrading)) {
-                    // When installing/downloading, we want to include the launcher
-                    // by default.
-                    _engine->SetVariableNumeric(L"Include_launcher", 1);
-                } else {
-                    // Any other action, if we didn't detect the MSI then we want to
-                    // keep it excluded
-                    _engine->SetVariableNumeric(L"Include_launcher", 0);
-                    _engine->SetVariableNumeric(L"AssociateFiles", 0);
+            // Update launcher install states
+            // If we didn't detect any existing installs, Include_launcher and
+            // InstallLauncherAllUsers will both be -1, so we will set to their
+            // defaults and leave the options enabled.
+            // Otherwise, if we detected an existing install, we disable the
+            // options so they remain fixed.
+            // The code in OnDetectRelatedMsiPackage is responsible for figuring
+            // out whether existing installs are compatible with the settings in
+            // place during detection.
+            LONGLONG blockedLauncher;
+            if (SUCCEEDED(BalGetNumericVariable(L"BlockedLauncher", &blockedLauncher))
+                    && blockedLauncher) {
+                _engine->SetVariableNumeric(L"Include_launcher", 0);
+                _engine->SetVariableNumeric(L"InstallLauncherAllUsers", 0);
+                _engine->SetVariableString(L"InstallLauncherAllUsersState", L"disable");
+                _engine->SetVariableString(L"Include_launcherState", L"disable");
+            }
+            else {
+                LONGLONG includeLauncher, includeLauncherAllUsers, associateFiles;
+
+                if (FAILED(BalGetNumericVariable(L"Include_launcher", &includeLauncher))) {
+                    includeLauncher = -1;
+                }
+                if (FAILED(BalGetNumericVariable(L"InstallLauncherAllUsers", &includeLauncherAllUsers))) {
+                    includeLauncherAllUsers = -1;
+                }
+                if (FAILED(BalGetNumericVariable(L"AssociateFiles", &associateFiles))) {
+                    associateFiles = -1;
+                }
+
+                if (includeLauncherAllUsers < 0) {
+                    includeLauncherAllUsers = 0;
+                    _engine->SetVariableNumeric(L"InstallLauncherAllUsers", includeLauncherAllUsers);
+                }
+
+                if (includeLauncher < 0) {
+                    if (BOOTSTRAPPER_ACTION_LAYOUT == _command.action ||
+                        (BOOTSTRAPPER_ACTION_INSTALL == _command.action && !_upgrading)) {
+                        // When installing/downloading, we include the launcher
+                        // (though downloads should ignore this setting anyway)
+                        _engine->SetVariableNumeric(L"Include_launcher", 1);
+                    } else {
+                        // Any other action, we should have detected an existing
+                        // install (e.g. on remove/modify), so if we didn't, we
+                        // assume it's not selected.
+                        _engine->SetVariableNumeric(L"Include_launcher", 0);
+                        _engine->SetVariableNumeric(L"AssociateFiles", 0);
+                    }
+                }
+
+                if (associateFiles < 0) {
+                    auto hr = LoadAssociateFilesStateFromKey(
+                        _engine,
+                        includeLauncherAllUsers ? HKEY_LOCAL_MACHINE : HKEY_CURRENT_USER
+                    );
+                    if (FAILED(hr)) {
+                        BalLog(BOOTSTRAPPER_LOG_LEVEL_ERROR, "Failed to load AssociateFiles state: error code 0x%08X", hr);
+                    } else if (hr == S_OK) {
+                        associateFiles = 1;
+                    }
+                    _engine->SetVariableNumeric(L"AssociateFiles", associateFiles);
                 }
             }
         }
diff --git a/Tools/msi/bundle/bundle.wxs b/Tools/msi/bundle/bundle.wxs
index 9b4f072..abfeb88 100644
--- a/Tools/msi/bundle/bundle.wxs
+++ b/Tools/msi/bundle/bundle.wxs
@@ -28,10 +28,11 @@
 
     <Variable Name="InstallAllUsers" Value="0" bal:Overridable="yes" />
     <?if "$(var.PyTestExt)"="" ?>
-    <Variable Name="InstallLauncherAllUsers" Value="0" bal:Overridable="yes" />
+    <Variable Name="InstallLauncherAllUsers" Value="-1" bal:Overridable="yes" />
     <?else ?>
-    <Variable Name="InstallLauncherAllUsers" Value="0" />
+    <Variable Name="InstallLauncherAllUsers" Value="-1" />
     <?endif ?>
+
     <Variable Name="TargetDir" Value="" bal:Overridable="yes" />
     <?if $(var.Platform)~="x64" ?>
     <Variable Name="DefaultAllUsersTargetDir" Value="[ProgramFiles64Folder]Python[WinVerNoDot]" bal:Overridable="yes" />
@@ -91,10 +92,11 @@
     <?endif ?>
 
     <Variable Name="LauncherOnly" Value="0" bal:Overridable="yes" />
+    <Variable Name="BlockedLauncher" Value="0" />
     <Variable Name="DetectedLauncher" Value="0" />
     <Variable Name="DetectedOldLauncher" Value="0" />
 
-    <Variable Name="AssociateFiles" Value="1" bal:Overridable="yes" />
+    <Variable Name="AssociateFiles" Value="-1" bal:Overridable="yes" />
     <Variable Name="Shortcuts" Value="1" bal:Overridable="yes" />
     <Variable Name="PrependPath" Value="0" bal:Overridable="yes" />
     <Variable Name="AppendPath" Value="0" bal:Overridable="yes" />