Issue #25715: Python 3.5.1 installer shows wrong upgrade path and incorrect logic for launcher detection.

5 files changed, 132 insertions, 103 deletions

diff --git a/Tools/msi/bundle/Default.thm b/Tools/msi/bundle/Default.thm
index 8ff4c3b9f3..4d9c97a194 100644
--- a/Tools/msi/bundle/Default.thm
+++ b/Tools/msi/bundle/Default.thm
@@ -66,7 +66,7 @@
 
         <Checkbox Name="Include_launcher" X="185" Y="251" Width="100" Height="24" TabStop="yes" FontId="3" HideWhenDisabled="no">#(loc.Include_launcherLabel)</Checkbox>
         <Checkbox Name="CustomInstallLauncherAllUsers" X="285" Y="251" Width="-11" Height="24" TabStop="yes" FontId="3">#(loc.InstallLauncherAllUsersLabel)</Checkbox>
-        <Text X="205" Y="276" Width="-11" Height="24" TabStop="no" FontId="5">#(loc.Include_launcherHelpLabel)</Text>
+        <Text Name="Include_launcherHelp" X="205" Y="276" Width="-11" Height="24" TabStop="no" FontId="5"></Text>
 
         <Button Name="Custom1BackButton" X="185" Y="-11" Width="85" Height="27" TabStop="yes" FontId="0">#(loc.CustomBackButton)</Button>
         <Button Name="CustomNextButton" X="-101" Y="-11" Width="85" Height="27" TabStop="yes" FontId="0">#(loc.CustomNextButton)</Button>
diff --git a/Tools/msi/bundle/Default.wxl b/Tools/msi/bundle/Default.wxl
index 1f3e77c59d..697066b2e1 100644
--- a/Tools/msi/bundle/Default.wxl
+++ b/Tools/msi/bundle/Default.wxl
@@ -76,7 +76,9 @@ Select Customize to review current options.</String>
   <String Id="Include_testLabel">Python &amp;test suite</String>
   <String Id="Include_testHelpLabel">Installs the standard library test suite.</String>
   <String Id="Include_launcherLabel">py &amp;launcher</String>
-  <String Id="Include_launcherHelpLabel">Installs the global 'py' launcher to make it easier to start Python.</String>
+  <String Id="Include_launcherHelp">Installs the global 'py' launcher to make it easier to start Python.</String>
+  <String Id="Include_launcherRemove">Use Programs and Features to remove the 'py' launcher.</String>
+  <String Id="Include_launcherUpgrade">Upgrades the global 'py' launcher from the previous version.</String>
   
   <String Id="AssociateFilesLabel">Associate &amp;files with Python (requires the py launcher)</String>
   <String Id="ShortcutsLabel">Create shortcuts for installed applications</String>
diff --git a/Tools/msi/bundle/bootstrap/PythonBootstrapperApplication.cpp b/Tools/msi/bundle/bootstrap/PythonBootstrapperApplication.cpp
index 9d94e3b46f..1462d7b67e 100644
--- a/Tools/msi/bundle/bootstrap/PythonBootstrapperApplication.cpp
+++ b/Tools/msi/bundle/bootstrap/PythonBootstrapperApplication.cpp
@@ -94,6 +94,7 @@ enum CONTROL_ID {
     ID_CUSTOM_ASSOCIATE_FILES_CHECKBOX,
     ID_CUSTOM_INSTALL_ALL_USERS_CHECKBOX,
     ID_CUSTOM_INSTALL_LAUNCHER_ALL_USERS_CHECKBOX,
+    ID_CUSTOM_INCLUDE_LAUNCHER_HELP_LABEL,
     ID_CUSTOM_COMPILE_ALL_CHECKBOX,
     ID_CUSTOM_BROWSE_BUTTON,
     ID_CUSTOM_BROWSE_BUTTON_LABEL,
@@ -158,6 +159,7 @@ static THEME_ASSIGN_CONTROL_ID CONTROL_ID_NAMES[] = {
     { ID_CUSTOM_ASSOCIATE_FILES_CHECKBOX, L"AssociateFiles" },
     { ID_CUSTOM_INSTALL_ALL_USERS_CHECKBOX, L"InstallAllUsers" },
     { ID_CUSTOM_INSTALL_LAUNCHER_ALL_USERS_CHECKBOX, L"CustomInstallLauncherAllUsers" },
+    { ID_CUSTOM_INCLUDE_LAUNCHER_HELP_LABEL, L"Include_launcherHelp" },
     { ID_CUSTOM_COMPILE_ALL_CHECKBOX, L"CompileAll" },
     { ID_CUSTOM_BROWSE_BUTTON, L"CustomBrowseButton" },
     { ID_CUSTOM_BROWSE_BUTTON_LABEL, L"CustomBrowseButtonLabel" },
@@ -454,6 +456,20 @@ class PythonBootstrapperApplication : public CBalBaseBootstrapperApplication {
 
         ThemeSendControlMessage(_theme, ID_CUSTOM_INSTALL_LAUNCHER_ALL_USERS_CHECKBOX, BM_SETCHECK,
             installLauncherAllUsers ? BST_CHECKED : BST_UNCHECKED, 0);
+
+        LOC_STRING *pLocString = nullptr;
+        LPCWSTR locKey = L"#(loc.Include_launcherHelp)";
+        LONGLONG detectedLauncher;
+
+        if (SUCCEEDED(BalGetNumericVariable(L"DetectedLauncher", &detectedLauncher)) && detectedLauncher) {
+            locKey = L"#(loc.Include_launcherRemove)";
+        } else if (SUCCEEDED(BalGetNumericVariable(L"DetectedOldLauncher", &detectedLauncher)) && detectedLauncher) {
+            locKey = L"#(loc.Include_launcherUpgrade)";
+        }
+
+        if (SUCCEEDED(LocGetString(_wixLoc, locKey, &pLocString)) && pLocString) {
+            ThemeSetTextControl(_theme, ID_CUSTOM_INCLUDE_LAUNCHER_HELP_LABEL, pLocString->wzText);
+        }
     }
 
     void Custom2Page_Show() {
@@ -641,6 +657,29 @@ public: // IBootstrapperApplication
         return nResult;
     }
 
+    virtual STDMETHODIMP_(int) OnDetectRelatedMsiPackage(
+        __in_z LPCWSTR wzPackageId,
+        __in_z LPCWSTR /*wzProductCode*/,
+        __in BOOL fPerMachine,
+        __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 (FAILED(hr)) {
+                BalLog(BOOTSTRAPPER_LOG_LEVEL_ERROR, "Failed to load AssociateFiles state: error code 0x%08X", hr);
+            }
+
+            _engine->SetVariableNumeric(L"Include_launcher", 1);
+            _engine->SetVariableNumeric(L"DetectedOldLauncher", 1);
+            _engine->SetVariableNumeric(L"InstallLauncherAllUsers", fPerMachine ? 1 : 0);
+        }
+        return CheckCanceled() ? IDCANCEL : IDNOACTION;
+    }
 
     virtual STDMETHODIMP_(int) OnDetectRelatedBundle(
         __in LPCWSTR wzBundleId,
@@ -656,24 +695,8 @@ public: // IBootstrapperApplication
         if (BOOTSTRAPPER_RELATED_OPERATION_DOWNGRADE == operation) {
             _downgradingOtherVersion = TRUE;
         } else if (BOOTSTRAPPER_RELATED_OPERATION_MAJOR_UPGRADE == operation) {
-            _upgradingOldVersion = TRUE;
-
-            // Assume we don't want the launcher or file associations, and if
-            // they have already been installed then loading the state will
-            // reactivate these settings.
-            _engine->SetVariableNumeric(L"Include_launcher", 0);
-            _engine->SetVariableNumeric(L"AssociateFiles", 0);
-            auto hr = LoadLauncherStateFromKey(_engine, HKEY_CURRENT_USER);
-            if (hr == S_FALSE) {
-                hr = LoadLauncherStateFromKey(_engine, HKEY_LOCAL_MACHINE);
-            }
-            if (FAILED(hr)) {
-                BalLog(
-                    BOOTSTRAPPER_LOG_LEVEL_ERROR,
-                    "Failed to load launcher state: error code 0x%08X",
-                    hr
-                );
-            }
+            BalLog(BOOTSTRAPPER_LOG_LEVEL_STANDARD, "Detected previous version - planning upgrade");
+            _upgrading = TRUE;
 
             LoadOptionalFeatureStates(_engine);
         } else if (BOOTSTRAPPER_RELATED_OPERATION_NONE == operation) {
@@ -699,9 +722,42 @@ public: // IBootstrapperApplication
 
     virtual STDMETHODIMP_(void) OnDetectPackageComplete(
         __in LPCWSTR wzPackageId,
-        __in HRESULT /*hrStatus*/,
+        __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);
+            }
+        }
+
+        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 (FAILED(hr)) {
+                BalLog(BOOTSTRAPPER_LOG_LEVEL_ERROR, "Failed to load AssociateFiles state: error code 0x%08X", hr);
+            }
+        }
+    }
 
 
     virtual STDMETHODIMP_(void) OnDetectComplete(__in HRESULT hrStatus) {
@@ -716,32 +772,7 @@ public: // IBootstrapperApplication
 
         if (SUCCEEDED(hrStatus)) {
             // Ensure the default path has been set
-            LONGLONG installAll;
-            LPWSTR targetDir = nullptr;
-            LPWSTR defaultTargetDir = nullptr;
-
-            hrStatus = BalGetStringVariable(L"TargetDir", &targetDir);
-            if (FAILED(hrStatus) || !targetDir || !targetDir[0]) {
-                ReleaseStr(targetDir);
-                targetDir = nullptr;
-
-                if (FAILED(BalGetNumericVariable(L"InstallAllUsers", &installAll))) {
-                    installAll = 0;
-                }
-
-                hrStatus = BalGetStringVariable(
-                    installAll ? L"DefaultAllUsersTargetDir" : L"DefaultJustForMeTargetDir",
-                    &defaultTargetDir
-                );
-
-                if (SUCCEEDED(hrStatus) && defaultTargetDir) {
-                    if (defaultTargetDir[0] && SUCCEEDED(BalFormatString(defaultTargetDir, &targetDir))) {
-                        hrStatus = _engine->SetVariableString(L"TargetDir", targetDir);
-                        ReleaseStr(targetDir);
-                    }
-                    ReleaseStr(defaultTargetDir);
-                }
-            }
+            hrStatus = EnsureTargetDir();
         }
 
         SetState(PYBA_STATE_DETECTED, hrStatus);
@@ -1396,9 +1427,14 @@ private:
 
         hr = LoadBootstrapperBAFunctions();
         BalExitOnFailure(hr, "Failed to load bootstrapper functions.");
+
         hr = UpdateUIStrings(_command.action);
         BalExitOnFailure(hr, "Failed to load UI strings.");
 
+        if (_command.action == BOOTSTRAPPER_ACTION_MODIFY) {
+            LoadOptionalFeatureStates(_engine);
+        }
+
         GetBundleFileVersion();
         // don't fail if we couldn't get the version info; best-effort only
     LExit:
@@ -1906,27 +1942,6 @@ private:
         for (DWORD i = 0; i < _theme->cControls; ++i) {
             THEME_CONTROL* pControl = _theme->rgControls + i;
             LPWSTR text = nullptr;
-            LPWSTR name = nullptr;
-            LOC_STRING *locText = nullptr;
-
-            // If a command link has a note, then add it.
-            if ((pControl->dwStyle & BS_TYPEMASK) == BS_COMMANDLINK ||
-                (pControl->dwStyle & BS_TYPEMASK) == BS_DEFCOMMANDLINK) {
-                hr = StrAllocFormatted(&name, L"#(loc.%lsNote)", pControl->sczName);
-                if (SUCCEEDED(hr)) {
-                    hr = LocGetString(_wixLoc, name, &locText);
-                    ReleaseStr(name);
-                    if (SUCCEEDED(hr) && locText && locText->wzText && locText->wzText[0]) {
-                        hr = BalFormatString(locText->wzText, &text);
-                        if (SUCCEEDED(hr) && text && text[0]) {
-                            ThemeSendControlMessage(_theme, pControl->wId, BCM_SETNOTE, 0, (LPARAM)text);
-                            ReleaseStr(text);
-                            text = nullptr;
-                        }
-                    }
-                }
-                hr = S_OK;
-            }
 
             if (!pControl->wPageId && pControl->sczText && *pControl->sczText) {
                 HRESULT hrFormat;
@@ -2048,6 +2063,7 @@ private:
 
         return;
     }
+
     HRESULT UpdateUIStrings(__in BOOTSTRAPPER_ACTION action) {
         HRESULT hr = S_OK;
         LPCWSTR likeInstalling = nullptr;
@@ -2270,6 +2286,30 @@ private:
                     StrFree(controlState);
                 }
                 StrFree(controlName);
+                controlName = nullptr;
+
+
+                // If a command link has a note, then add it.
+                if ((pControl->dwStyle & BS_TYPEMASK) == BS_COMMANDLINK ||
+                    (pControl->dwStyle & BS_TYPEMASK) == BS_DEFCOMMANDLINK) {
+                    hr = StrAllocFormatted(&controlName, L"#(loc.%lsNote)", pControl->sczName);
+                    if (SUCCEEDED(hr)) {
+                        LOC_STRING *locText = nullptr;
+                        hr = LocGetString(_wixLoc, controlName, &locText);
+                        if (SUCCEEDED(hr) && locText && locText->wzText && locText->wzText[0]) {
+                            LPWSTR text = nullptr;
+                            hr = BalFormatString(locText->wzText, &text);
+                            if (SUCCEEDED(hr) && text && text[0]) {
+                                ThemeSendControlMessage(_theme, pControl->wId, BCM_SETNOTE, 0, (LPARAM)text);
+                                ReleaseStr(text);
+                                text = nullptr;
+                            }
+                        }
+                        ReleaseStr(controlName);
+                        controlName = nullptr;
+                    }
+                    hr = S_OK;
+                }
             }
 
             ThemeControlEnable(_theme, pControl->wId, enableControl);
@@ -2515,9 +2555,8 @@ private:
                 if (_installPage == PAGE_LOADING) {
                     switch (_command.action) {
                     case BOOTSTRAPPER_ACTION_INSTALL:
-                        if (_upgradingOldVersion) {
+                        if (_upgrading) {
                             _installPage = PAGE_UPGRADE;
-                            _upgrading = TRUE;
                         } else if (SUCCEEDED(BalGetNumericVariable(L"SimpleInstall", &simple)) && simple) {
                             _installPage = PAGE_SIMPLE_INSTALL;
                         } else {
@@ -2560,11 +2599,9 @@ private:
         static BAL_CONDITION WILL_ELEVATE_CONDITION = {
             L"not WixBundleElevated and ("
                 /*Elevate when installing for all users*/
-                L"InstallAllUsers or"
+                L"InstallAllUsers or "
                 /*Elevate when installing the launcher for all users and it was not detected*/
-                L"(InstallLauncherAllUsers and Include_launcher and not DetectedLauncher) or"
-                /*Elevate when the launcher was installed for all users and it is being removed*/
-                L"(DetectedLauncher and DetectedLauncherAllUsers and not Include_launcher)"
+                L"(Include_launcher and InstallLauncherAllUsers and not DetectedLauncher)"
             L")",
             L""
         };
@@ -2867,19 +2904,16 @@ private:
         return HRESULT_FROM_WIN32(res);
     }
 
-    static HRESULT LoadLauncherStateFromKey(
+    static HRESULT LoadAssociateFilesStateFromKey(
         __in IBootstrapperEngine* pEngine,
         __in HKEY hkHive
     ) {
         const LPCWSTR subkey = L"Software\\Python\\PyLauncher";
         HKEY hKey;
         LRESULT res;
+        HRESULT hr;
 
-        if (IsTargetPlatformx64(pEngine)) {
-            res = RegOpenKeyExW(hkHive, subkey, 0, KEY_READ | KEY_WOW64_64KEY, &hKey);
-        } else {
-            res = RegOpenKeyExW(hkHive, subkey, 0, KEY_READ | KEY_WOW64_32KEY, &hKey);
-        }
+        res = RegOpenKeyExW(hkHive, subkey, 0, KEY_READ | KEY_WOW64_32KEY, &hKey);
 
         if (res == ERROR_FILE_NOT_FOUND) {
             return S_FALSE;
@@ -2888,26 +2922,17 @@ private:
             return HRESULT_FROM_WIN32(res);
         }
 
-        res = RegQueryValueExW(hKey, nullptr, nullptr, nullptr, nullptr, nullptr);
-        if (res == ERROR_FILE_NOT_FOUND) {
-            pEngine->SetVariableNumeric(L"Include_launcher", 0);
-        } else if (res == ERROR_SUCCESS) {
-            pEngine->SetVariableNumeric(L"Include_launcher", 1);
-            pEngine->SetVariableNumeric(L"DetectedLauncher", 1);
-            pEngine->SetVariableNumeric(L"InstallLauncherAllUsers", (hkHive == HKEY_LOCAL_MACHINE) ? 1 : 0);
-            pEngine->SetVariableNumeric(L"DetectedLauncherAllUsers", (hkHive == HKEY_LOCAL_MACHINE) ? 1 : 0);
-            pEngine->SetVariableString(L"InstallLauncherAllUsersState", L"disable");
-        }
-
         res = RegQueryValueExW(hKey, L"AssociateFiles", nullptr, nullptr, nullptr, nullptr);
         if (res == ERROR_FILE_NOT_FOUND) {
-            pEngine->SetVariableNumeric(L"AssociateFiles", 0);
+            hr = S_FALSE;
         } else if (res == ERROR_SUCCESS) {
-            pEngine->SetVariableNumeric(L"AssociateFiles", 1);
+            hr = S_OK;
+        } else {
+            hr = HRESULT_FROM_WIN32(res);
         }
 
         RegCloseKey(hKey);
-        return S_OK;
+        return hr;
     }
 
     static void LoadOptionalFeatureStates(__in IBootstrapperEngine* pEngine) {
@@ -2918,7 +2943,11 @@ private:
         HKEY hkHive;
 
         // The launcher installation is separate from the Python install, so we
-        // check its state later. This also checks the file association option.
+        // check its state later. For now, assume we don't want the launcher or
+        // file associations, and if they have already been installed then
+        // loading the state will reactivate these settings.
+        pEngine->SetVariableNumeric(L"Include_launcher", 0);
+        pEngine->SetVariableNumeric(L"AssociateFiles", 0);
 
         // Get the registry key from the bundle, to save having to duplicate it
         // in multiple places.
@@ -3089,7 +3118,6 @@ public:
         _hrFinal = hrHostInitialization;
 
         _downgradingOtherVersion = FALSE;
-        _upgradingOldVersion = FALSE;
         _restartResult = BOOTSTRAPPER_APPLY_RESTART_NONE;
         _restartRequired = FALSE;
         _allowRestart = FALSE;
@@ -3113,8 +3141,6 @@ public:
 
         _hBAFModule = nullptr;
         _baFunction = nullptr;
-
-        EnsureTargetDir();
     }
 
 
@@ -3174,7 +3200,6 @@ private:
     DWORD _calculatedExecuteProgress;
 
     BOOL _downgradingOtherVersion;
-    BOOL _upgradingOldVersion;
     BOOTSTRAPPER_APPLY_RESTART _restartResult;
     BOOL _restartRequired;
     BOOL _allowRestart;
diff --git a/Tools/msi/bundle/bundle.wxs b/Tools/msi/bundle/bundle.wxs
index 978efc0a21..38307e063c 100644
--- a/Tools/msi/bundle/bundle.wxs
+++ b/Tools/msi/bundle/bundle.wxs
@@ -51,7 +51,7 @@
     -->
     <Variable Name="DefaultCustomTargetDir" Value="" bal:Overridable="yes" />
 
-    <Variable Name="InstallAllUsersState" Value="enabled" />
+    <Variable Name="InstallAllUsersState" Value="enabled" bal:Overridable="yes" />
     <?if "$(var.PyTestExt)"="" ?>
     <Variable Name="InstallLauncherAllUsersState" Value="enabled" bal:Overridable="yes" />
     <?else ?>
@@ -72,6 +72,7 @@
     <Variable Name="Include_pip" Value="1" bal:Overridable="yes" />
     <?if "$(var.PyTestExt)"="" ?>
     <Variable Name="Include_launcher" Value="1" bal:Overridable="yes" />
+    <Variable Name="Include_launcherState" Value="enabled" bal:Overridable="yes" />
     <?else ?>
     <Variable Name="Include_launcher" Value="0" />
     <Variable Name="Include_launcherState" Value="disable" />
@@ -81,6 +82,7 @@
     
     <Variable Name="LauncherOnly" Value="0" bal:Overridable="yes" />
     <Variable Name="DetectedLauncher" Value="0" />
+    <Variable Name="DetectedOldLauncher" Value="0" />
     
     <Variable Name="AssociateFiles" Value="1" bal:Overridable="yes" />
     <Variable Name="Shortcuts" Value="1" bal:Overridable="yes" />
diff --git a/Tools/msi/bundle/packagegroups/launcher.wxs b/Tools/msi/bundle/packagegroups/launcher.wxs
index 4b03fd29b6..4444f45a98 100644
--- a/Tools/msi/bundle/packagegroups/launcher.wxs
+++ b/Tools/msi/bundle/packagegroups/launcher.wxs
@@ -11,7 +11,7 @@
                         EnableFeatureSelection="yes"
                         Permanent="yes"
                         Visible="yes"
-                        InstallCondition="(InstallAllUsers or InstallLauncherAllUsers) and Include_launcher" />
+                        InstallCondition="(InstallAllUsers or InstallLauncherAllUsers) and Include_launcher and not DetectedLauncher" />
 
             <MsiPackage Id="launcher_JustForMe"
                         SourceFile="launcher.msi"
@@ -21,7 +21,7 @@
                         EnableFeatureSelection="yes"
                         Permanent="yes"
                         Visible="yes"
-                        InstallCondition="not (InstallAllUsers or InstallLauncherAllUsers) and Include_launcher" />
+                        InstallCondition="not (InstallAllUsers or InstallLauncherAllUsers) and Include_launcher and not DetectedLauncher" />
         </PackageGroup>
     </Fragment>
 </Wix>
\ No newline at end of file