Force restoration of Chrome's shortcuts when Active Setup kicks in in response to a major OS upgrade.

BUG=502363
TEST=setup_unittests.exe --gtest_filter=SetupUtilTest.UpdateLastOSUpgradeHandledByActiveSetup

Review URL: https://codereview.chromium.org/1231973002

Cr-Commit-Position: refs/heads/master@{#338572}

5 files changed, 234 insertions, 25 deletions

diff --git a/chrome/installer/setup/install.cc b/chrome/installer/setup/install.cc
index 409fa25..37a1417 100644
--- a/chrome/installer/setup/install.cc
+++ b/chrome/installer/setup/install.cc
@@ -24,6 +24,7 @@
 #include "chrome/common/chrome_switches.h"
 #include "chrome/installer/setup/install_worker.h"
 #include "chrome/installer/setup/setup_constants.h"
+#include "chrome/installer/setup/setup_util.h"
 #include "chrome/installer/setup/update_active_setup_version_work_item.h"
 #include "chrome/installer/util/auto_launch_util.h"
 #include "chrome/installer/util/beacons.h"
@@ -655,6 +656,13 @@ void HandleActiveSetupForBrowser(const base::FilePath& installation_root,
                                  const installer::Product& chrome,
                                  bool force) {
   DCHECK(chrome.is_chrome());
+
+  // If the shortcuts are not being forcefully created we may want to forcefully
+  // create them anyways if this Active Setup trigger is in response to an OS
+  // update.
+  force = force || installer::UpdateLastOSUpgradeHandledByActiveSetup(
+                       chrome.distribution());
+
   // Only create shortcuts on Active Setup if the first run sentinel is not
   // present for this user (as some shortcuts used to be installed on first
   // run and this could otherwise re-install shortcuts for users that have
diff --git a/chrome/installer/setup/setup_util.cc b/chrome/installer/setup/setup_util.cc
index 7f0af71..b9575d0 100644
--- a/chrome/installer/setup/setup_util.cc
+++ b/chrome/installer/setup/setup_util.cc
@@ -7,6 +7,7 @@
 #include "chrome/installer/setup/setup_util.h"
 
 #include <windows.h>
+#include <stdint.h>
 
 #include "base/command_line.h"
 #include "base/cpu.h"
@@ -17,15 +18,21 @@
 #include "base/process/kill.h"
 #include "base/process/launch.h"
 #include "base/process/process_handle.h"
+#include "base/strings/string_number_conversions.h"
+#include "base/strings/string_split.h"
 #include "base/strings/string_util.h"
 #include "base/strings/utf_string_conversions.h"
 #include "base/version.h"
 #include "base/win/registry.h"
+#include "base/win/win_util.h"
 #include "base/win/windows_version.h"
 #include "chrome/installer/setup/setup_constants.h"
+#include "chrome/installer/setup/update_active_setup_version_work_item.h"
+#include "chrome/installer/util/app_registration_data.h"
 #include "chrome/installer/util/app_registration_data.h"
 #include "chrome/installer/util/copy_tree_work_item.h"
 #include "chrome/installer/util/google_update_constants.h"
+#include "chrome/installer/util/install_util.h"
 #include "chrome/installer/util/installation_state.h"
 #include "chrome/installer/util/installer_state.h"
 #include "chrome/installer/util/master_preferences.h"
@@ -95,6 +102,104 @@ bool SupportsSingleInstall(BrowserDistribution::Type type) {
 
 }  // namespace
 
+bool UpdateLastOSUpgradeHandledByActiveSetup(BrowserDistribution* dist) {
+  // FIRST: Find the value of the latest OS upgrade registered in the Active
+  // Setup version (bumped on every major OS upgrade), defaults to 0 if no OS
+  // upgrade was ever encountered by this install.
+  DWORD latest_os_upgrade = 0;
+
+  {
+    const base::string16 active_setup_key_path(
+        InstallUtil::GetActiveSetupPath(dist));
+
+    base::win::RegKey active_setup_key;
+    if (active_setup_key.Open(HKEY_LOCAL_MACHINE, active_setup_key_path.c_str(),
+                              KEY_QUERY_VALUE) == ERROR_SUCCESS) {
+      base::string16 existing_version;
+      if (active_setup_key.ReadValue(L"Version",
+                                     &existing_version) == ERROR_SUCCESS) {
+        std::vector<base::string16> version_components =
+            base::SplitString(existing_version, L",", base::TRIM_WHITESPACE,
+                              base::SPLIT_WANT_NONEMPTY);
+        uint32_t latest_os_upgrade_uint = 0;
+        if (version_components.size() == 4U &&
+            base::StringToUint(
+                version_components[UpdateActiveSetupVersionWorkItem::
+                                       VersionComponent::OS_UPGRADES],
+                &latest_os_upgrade_uint)) {
+          latest_os_upgrade = static_cast<DWORD>(latest_os_upgrade_uint);
+        } else {
+          LOG(ERROR) << "Failed to parse OS_UPGRADES component of "
+                     << existing_version;
+        }
+      }
+    }
+  }
+
+  // Whether the read failed or the existing value is 0, do not proceed.
+  if (latest_os_upgrade == 0U)
+    return false;
+
+  static const wchar_t kLastOSUpgradeHandledRegName[] = L"LastOSUpgradeHandled";
+
+  // SECOND: Find out the value of the last OS upgrade handled, defaults to 0 if
+  // none was ever handled.
+  DWORD last_os_upgrade_handled = 0;
+
+  base::string16 last_upgrade_handled_key_path =
+      dist->GetAppRegistrationData().GetStateMediumKey();
+  last_upgrade_handled_key_path.push_back(L'\\');
+  last_upgrade_handled_key_path.append(kLastOSUpgradeHandledRegName);
+
+  base::string16 user_specific_value;
+  // This should never fail. If it does, the beacon will be written in the key's
+  // default value, which is okay since the majority case is likely a machine
+  // with a single user.
+  if (!base::win::GetUserSidString(&user_specific_value))
+    NOTREACHED();
+
+  base::win::RegKey last_upgrade_key;
+  if (last_upgrade_key.Create(
+          HKEY_LOCAL_MACHINE, last_upgrade_handled_key_path.c_str(),
+          KEY_WOW64_32KEY | KEY_QUERY_VALUE | KEY_SET_VALUE) != ERROR_SUCCESS) {
+    LOG(ERROR) << "Failed to create LastOSUpgradeHandled value @ "
+               << last_upgrade_handled_key_path;
+    // If the key is not read/writeable, do not proceed as this could result in
+    // handling an OS upgrade twice.
+    return false;
+  }
+
+  // It's okay for this read to fail (i.e. there is an OS upgrade available but
+  // this user never handled one yet).
+  last_upgrade_key.ReadValueDW(user_specific_value.c_str(),
+                               &last_os_upgrade_handled);
+
+  // THIRD: Figure out whether the latest OS upgrade has been handled already.
+
+  if (last_os_upgrade_handled >= latest_os_upgrade) {
+    LOG_IF(ERROR, last_os_upgrade_handled > latest_os_upgrade)
+        << "Last OS upgrade handled is somehow ahead of the latest OS upgrade?";
+    VLOG_IF(1, last_os_upgrade_handled == latest_os_upgrade)
+        << "Latest OS upgrade already handled.";
+    return false;
+  }
+
+  // At this point |last_os_upgrade_handled < latest_os_upgrade| so,
+  // FOURTH: store the fact that the latest OS upgrade has been handled and
+  // return true for the caller to act accordingly.
+
+  if (last_upgrade_key.WriteValue(user_specific_value.c_str(),
+                                  latest_os_upgrade) != ERROR_SUCCESS) {
+    LOG(ERROR) << "Failed to save latest_os_upgrade value ("
+               << latest_os_upgrade << ") to " << last_upgrade_handled_key_path;
+    // Do not proceed if the write fails as this could otherwise result in
+    // handling this OS upgrade multiple times.
+    return false;
+  }
+
+  return true;
+}
+
 int CourgettePatchFiles(const base::FilePath& src,
                         const base::FilePath& patch,
                         const base::FilePath& dest) {
diff --git a/chrome/installer/setup/setup_util.h b/chrome/installer/setup/setup_util.h
index 0b04b7b..4d6303a 100644
--- a/chrome/installer/setup/setup_util.h
+++ b/chrome/installer/setup/setup_util.h
@@ -31,6 +31,13 @@ class InstallationState;
 class InstallerState;
 class ProductState;
 
+// Sets a bit in the registry to note that the latest OS upgrade notification
+// has been handled by this user. Returns true if the previous bit was
+// different or absent (i.e., the latest OS update wasn't handled yet), in
+// which case subsequent calls to this method will return false until the next
+// OS upgrade. This call is only valid on system-level installs.
+bool UpdateLastOSUpgradeHandledByActiveSetup(BrowserDistribution* dist);
+
 // Applies a patch file to source file using Courgette. Returns 0 in case of
 // success. In case of errors, it returns kCourgetteErrorOffset + a Courgette
 // status code, as defined in courgette/courgette.h
diff --git a/chrome/installer/setup/setup_util_unittest.cc b/chrome/installer/setup/setup_util_unittest.cc
index bebf65a..10b0936 100644
--- a/chrome/installer/setup/setup_util_unittest.cc
+++ b/chrome/installer/setup/setup_util_unittest.cc
@@ -11,6 +11,7 @@
 #include "base/command_line.h"
 #include "base/files/file_util.h"
 #include "base/files/scoped_temp_dir.h"
+#include "base/macros.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/process/kill.h"
 #include "base/process/launch.h"
@@ -20,11 +21,15 @@
 #include "base/threading/platform_thread.h"
 #include "base/time/time.h"
 #include "base/version.h"
+#include "base/win/registry.h"
 #include "base/win/scoped_handle.h"
 #include "base/win/windows_version.h"
 #include "chrome/installer/setup/setup_constants.h"
 #include "chrome/installer/setup/setup_util.h"
+#include "chrome/installer/setup/update_active_setup_version_work_item.h"
+#include "chrome/installer/util/browser_distribution.h"
 #include "chrome/installer/util/google_update_constants.h"
+#include "chrome/installer/util/install_util.h"
 #include "chrome/installer/util/installation_state.h"
 #include "chrome/installer/util/installer_state.h"
 #include "chrome/installer/util/updating_app_registration_data.h"
@@ -33,20 +38,22 @@
 
 namespace {
 
-class SetupUtilTestWithDir : public testing::Test {
+class SetupUtilTest : public testing::Test {
  protected:
+  SetupUtilTest() {}
+
   void SetUp() override {
-    // Create a temp directory for testing.
     ASSERT_TRUE(test_dir_.CreateUniqueTempDir());
+    registry_override_manager_.OverrideRegistry(HKEY_CURRENT_USER);
+    registry_override_manager_.OverrideRegistry(HKEY_LOCAL_MACHINE);
   }
 
-  void TearDown() override {
-    // Clean up test directory manually so we can fail if it leaks.
-    ASSERT_TRUE(test_dir_.Delete());
-  }
-
-  // The temporary directory used to contain the test operations.
   base::ScopedTempDir test_dir_;
+
+ private:
+  registry_util::RegistryOverrideManager registry_override_manager_;
+
+  DISALLOW_COPY_AND_ASSIGN(SetupUtilTest);
 };
 
 // The privilege tested in ScopeTokenPrivilege tests below.
@@ -100,8 +107,88 @@ bool CurrentProcessHasPrivilege(const wchar_t* privilege_name) {
 
 }  // namespace
 
+TEST_F(SetupUtilTest, UpdateLastOSUpgradeHandledByActiveSetup) {
+  BrowserDistribution* chrome_dist =
+      BrowserDistribution::GetSpecificDistribution(
+          BrowserDistribution::CHROME_BROWSER);
+  const base::string16 active_setup_path(
+      InstallUtil::GetActiveSetupPath(chrome_dist));
+
+  base::win::RegKey test_key;
+  base::string16 unused_tmp;
+
+  EXPECT_EQ(ERROR_FILE_NOT_FOUND,
+            test_key.Open(HKEY_LOCAL_MACHINE, active_setup_path.c_str(),
+                          KEY_QUERY_VALUE));
+  // The WorkItem assume the ActiveSetup key itself already exists and only
+  // handles the Version entry, create it now, but don't fill the "Version"
+  // entry just yet.
+  EXPECT_EQ(ERROR_SUCCESS,
+            test_key.Create(HKEY_LOCAL_MACHINE, active_setup_path.c_str(),
+                            KEY_QUERY_VALUE));
+  EXPECT_EQ(ERROR_FILE_NOT_FOUND, test_key.ReadValue(L"Version", &unused_tmp));
+
+  // Test returns false when no Active Setup version present (and doesn't alter
+  // that state).
+  EXPECT_FALSE(
+      installer::UpdateLastOSUpgradeHandledByActiveSetup(chrome_dist));
+  EXPECT_EQ(ERROR_FILE_NOT_FOUND, test_key.ReadValue(L"Version", &unused_tmp));
+
+  {
+    UpdateActiveSetupVersionWorkItem active_setup_work_item(
+        active_setup_path, UpdateActiveSetupVersionWorkItem::UPDATE);
+    active_setup_work_item.Do();
+    EXPECT_EQ(ERROR_SUCCESS, test_key.ReadValue(L"Version", &unused_tmp));
+  }
+
+  // Test returns false with default Active Setup version.
+  EXPECT_FALSE(
+      installer::UpdateLastOSUpgradeHandledByActiveSetup(chrome_dist));
+  EXPECT_EQ(ERROR_SUCCESS, test_key.ReadValue(L"Version", &unused_tmp));
+
+  // Run through |kIterations| sequences of bumping the OS upgrade version |i|
+  // times and simulating a regular update |kIterations-i| times, confirming
+  // that handling any number of OS upgrades only results in a single hit and
+  // that no amount of regular updates after that result in any hit.
+  const size_t kIterations = 4U;
+  for (size_t i = 0U; i < kIterations; ++i) {
+    SCOPED_TRACE(i);
+    // Bump the OS_UPGRADES component |i| times.
+    for (size_t j = 0; j < i; ++j) {
+      UpdateActiveSetupVersionWorkItem active_setup_work_item(
+          active_setup_path, UpdateActiveSetupVersionWorkItem::
+                                 UPDATE_AND_BUMP_OS_UPGRADES_COMPONENT);
+      active_setup_work_item.Do();
+    }
+
+    // There should be a single OS upgrade to handle if the OS_UPGRADES
+    // component was bumped at least once.
+    EXPECT_EQ(i > 0, installer::UpdateLastOSUpgradeHandledByActiveSetup(
+                         chrome_dist));
+
+    // We should only be told to handle the latest OS upgrade once above.
+    EXPECT_FALSE(
+        installer::UpdateLastOSUpgradeHandledByActiveSetup(chrome_dist));
+    EXPECT_FALSE(
+        installer::UpdateLastOSUpgradeHandledByActiveSetup(chrome_dist));
+
+    // Run |kIterations-i| regular updates.
+    for (size_t j = i; j < kIterations; ++j) {
+      UpdateActiveSetupVersionWorkItem active_setup_work_item(
+          active_setup_path, UpdateActiveSetupVersionWorkItem::UPDATE);
+      active_setup_work_item.Do();
+    }
+
+    // No amount of regular updates should trigger an OS upgrade to be handled.
+    EXPECT_FALSE(
+        installer::UpdateLastOSUpgradeHandledByActiveSetup(chrome_dist));
+    EXPECT_FALSE(
+        installer::UpdateLastOSUpgradeHandledByActiveSetup(chrome_dist));
+  }
+}
+
 // Test that we are parsing Chrome version correctly.
-TEST_F(SetupUtilTestWithDir, GetMaxVersionFromArchiveDirTest) {
+TEST_F(SetupUtilTest, GetMaxVersionFromArchiveDirTest) {
   // Create a version dir
   base::FilePath chrome_dir = test_dir_.path().AppendASCII("1.0.0.0");
   base::CreateDirectory(chrome_dir);
@@ -137,7 +224,7 @@ TEST_F(SetupUtilTestWithDir, GetMaxVersionFromArchiveDirTest) {
   ASSERT_EQ(version->GetString(), "9.9.9.9");
 }
 
-TEST_F(SetupUtilTestWithDir, DeleteFileFromTempProcess) {
+TEST_F(SetupUtilTest, DeleteFileFromTempProcess) {
   base::FilePath test_file;
   base::CreateTemporaryFileInDir(test_dir_.path(), &test_file);
   ASSERT_TRUE(base::PathExists(test_file));
@@ -267,7 +354,7 @@ namespace {
 
 // A test fixture that configures an InstallationState and an InstallerState
 // with a product being updated.
-class FindArchiveToPatchTest : public SetupUtilTestWithDir {
+class FindArchiveToPatchTest : public SetupUtilTest {
  protected:
   class FakeInstallationState : public installer::InstallationState {
   };
@@ -291,7 +378,7 @@ class FindArchiveToPatchTest : public SetupUtilTestWithDir {
   };
 
   void SetUp() override {
-    SetupUtilTestWithDir::SetUp();
+    SetupUtilTest::SetUp();
     product_version_ = Version("30.0.1559.0");
     max_version_ = Version("47.0.1559.0");
 
@@ -318,7 +405,7 @@ class FindArchiveToPatchTest : public SetupUtilTestWithDir {
 
   void TearDown() override {
     original_state_.reset();
-    SetupUtilTestWithDir::TearDown();
+    SetupUtilTest::TearDown();
   }
 
   base::FilePath GetArchivePath(const Version& version) const {
@@ -434,6 +521,8 @@ class MigrateMultiToSingleTest : public testing::Test {
   static const HKEY kRootKey;
   static const wchar_t kVersionString[];
   static const wchar_t kMultiChannel[];
+
+ private:
   registry_util::RegistryOverrideManager registry_override_manager_;
 };
 
diff --git a/chrome/installer/setup/update_active_setup_version_work_item.h b/chrome/installer/setup/update_active_setup_version_work_item.h
index 03aec44..bb962ac 100644
--- a/chrome/installer/setup/update_active_setup_version_work_item.h
+++ b/chrome/installer/setup/update_active_setup_version_work_item.h
@@ -15,6 +15,18 @@
 // on demand. This WorkItem is only viable on machine-wide installs.
 class UpdateActiveSetupVersionWorkItem : public WorkItem {
  public:
+  // The components of the Active Setup Version entry, in order.
+  enum VersionComponent {
+    // The major version.
+    MAJOR,
+    // Unused component, always 0 for now.
+    UNUSED1,
+    // Number of OS upgrades handled since original install.
+    OS_UPGRADES,
+    // Unused component, always 0 for now.
+    UNUSED2,
+  };
+
   // The operation to be performed by this UpdateActiveSetupVersionWorkItem.
   enum Operation {
     // Update (or install if not present) the Active Setup "Version" in the
@@ -36,18 +48,6 @@ class UpdateActiveSetupVersionWorkItem : public WorkItem {
   void Rollback() override;
 
  private:
-  // The components of the Active Setup Version entry, in order.
-  enum ActiveSetupVersionComponent {
-    // The major version.
-    MAJOR,
-    // Unused component, always 0 for now.
-    UNUSED1,
-    // Number of OS upgrades handled since original install.
-    OS_UPGRADES,
-    // Unused component, always 0 for now.
-    UNUSED2,
-  };
-
   // Returns the updated Active Setup version to be used based on the
   // |existing_version|.
   base::string16 GetUpdatedActiveSetupVersion(