Can now adjust the number of retries before the blacklist is disabled.

Also added UMA recording of the blacklist disabled and retry succeeded events.  

Deprecating the blacklist start up events 'BLACKLIST_THUNK_SETUP_FAILED' and 
'BLACKLIST_INTERCEPTION_FAILED' as we are switching over to a more coarse 
grained recording of the set up status (this information can be gathered from 
crash reporting now anyways).

BUG=

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

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@277614 0039d316-1c4b-4281-b951-d872f2087c98

5 files changed, 200 insertions, 97 deletions

diff --git a/chrome_elf/blacklist/blacklist.cc b/chrome_elf/blacklist/blacklist.cc
index 8107ffc..055b0ab 100644
--- a/chrome_elf/blacklist/blacklist.cc
+++ b/chrome_elf/blacklist/blacklist.cc
@@ -59,20 +59,54 @@ namespace {
 // determine if the blacklist is enabled for them.
 bool g_blacklist_initialized = false;
 
-// Record that the thunk setup completed succesfully and close the registry
-// key handle since it is no longer needed.
-void RecordSuccessfulThunkSetup(HKEY* key) {
-  if (key != NULL) {
-    DWORD blacklist_state = blacklist::BLACKLIST_SETUP_RUNNING;
-    ::RegSetValueEx(*key,
-                    blacklist::kBeaconState,
-                    0,
-                    REG_DWORD,
-                    reinterpret_cast<LPBYTE>(&blacklist_state),
-                    sizeof(blacklist_state));
-    ::RegCloseKey(*key);
-    key = NULL;
+// Helper to set DWORD registry values.
+DWORD SetDWValue(HKEY* key, const wchar_t* property, DWORD value) {
+  return ::RegSetValueEx(*key,
+                         property,
+                         0,
+                         REG_DWORD,
+                         reinterpret_cast<LPBYTE>(&value),
+                         sizeof(value));
+}
+
+bool GenerateStateFromBeaconAndAttemptCount(HKEY* key, DWORD blacklist_state) {
+  LONG result = 0;
+  if (blacklist_state == blacklist::BLACKLIST_SETUP_RUNNING) {
+    // Some part of the blacklist setup failed last time.  If this has occured
+    // blacklist::kBeaconMaxAttempts times in a row we switch the state to
+    // failed and skip setting up the blacklist.
+    DWORD attempt_count = 0;
+    DWORD attempt_count_size = sizeof(attempt_count);
+    result = ::RegQueryValueEx(*key,
+                               blacklist::kBeaconAttemptCount,
+                               0,
+                               NULL,
+                               reinterpret_cast<LPBYTE>(&attempt_count),
+                               &attempt_count_size);
+
+    if (result == ERROR_FILE_NOT_FOUND)
+      attempt_count = 0;
+    else if (result != ERROR_SUCCESS)
+      return false;
+
+    ++attempt_count;
+    SetDWValue(key, blacklist::kBeaconAttemptCount, attempt_count);
+
+    if (attempt_count >= blacklist::kBeaconMaxAttempts) {
+      blacklist_state = blacklist::BLACKLIST_SETUP_FAILED;
+      SetDWValue(key, blacklist::kBeaconState, blacklist_state);
+      return false;
+    }
+  } else if (blacklist_state == blacklist::BLACKLIST_ENABLED) {
+    // If the blacklist succeeded on the previous run reset the failure
+    // counter.
+    result =
+        SetDWValue(key, blacklist::kBeaconAttemptCount, static_cast<DWORD>(0));
+    if (result != ERROR_SUCCESS) {
+      return false;
+    }
   }
+  return true;
 }
 
 }  // namespace
@@ -102,7 +136,7 @@ bool LeaveSetupBeacon() {
     return false;
 
   // Retrieve the current blacklist state.
-  DWORD blacklist_state = BLACKLIST_DISABLED;
+  DWORD blacklist_state = BLACKLIST_STATE_MAX;
   DWORD blacklist_state_size = sizeof(blacklist_state);
   DWORD type = 0;
   result = ::RegQueryValueEx(key,
@@ -112,21 +146,18 @@ bool LeaveSetupBeacon() {
                              reinterpret_cast<LPBYTE>(&blacklist_state),
                              &blacklist_state_size);
 
-  if (blacklist_state != BLACKLIST_ENABLED ||
-      result != ERROR_SUCCESS || type != REG_DWORD) {
+  if (blacklist_state == BLACKLIST_DISABLED || result != ERROR_SUCCESS ||
+      type != REG_DWORD) {
+    ::RegCloseKey(key);
+    return false;
+  }
+
+  if (!GenerateStateFromBeaconAndAttemptCount(&key, blacklist_state)) {
     ::RegCloseKey(key);
     return false;
   }
 
-  // Mark the blacklist setup code as running so if it crashes the blacklist
-  // won't be enabled for the next run.
-  blacklist_state = BLACKLIST_SETUP_RUNNING;
-  result = ::RegSetValueEx(key,
-                           kBeaconState,
-                           0,
-                           REG_DWORD,
-                           reinterpret_cast<LPBYTE>(&blacklist_state),
-                           sizeof(blacklist_state));
+  result = SetDWValue(&key, kBeaconState, BLACKLIST_SETUP_RUNNING);
   ::RegCloseKey(key);
 
   return (result == ERROR_SUCCESS);
@@ -147,15 +178,28 @@ bool ResetBeacon() {
   if (result != ERROR_SUCCESS)
     return false;
 
-  DWORD blacklist_state = BLACKLIST_ENABLED;
-  result = ::RegSetValueEx(key,
-                           kBeaconState,
-                           0,
-                           REG_DWORD,
-                           reinterpret_cast<LPBYTE>(&blacklist_state),
-                           sizeof(blacklist_state));
-  ::RegCloseKey(key);
+  DWORD blacklist_state = BLACKLIST_STATE_MAX;
+  DWORD blacklist_state_size = sizeof(blacklist_state);
+  DWORD type = 0;
+  result = ::RegQueryValueEx(key,
+                             kBeaconState,
+                             0,
+                             &type,
+                             reinterpret_cast<LPBYTE>(&blacklist_state),
+                             &blacklist_state_size);
 
+  if (result != ERROR_SUCCESS || type != REG_DWORD) {
+    ::RegCloseKey(key);
+    return false;
+  }
+
+  // Reaching this point with the setup running state means the setup
+  // succeeded and so we reset to enabled.  Any other state indicates that setup
+  // was skipped; in that case we leave the state alone for later recording.
+  if (blacklist_state == BLACKLIST_SETUP_RUNNING)
+    result = SetDWValue(&key, kBeaconState, BLACKLIST_ENABLED);
+
+  ::RegCloseKey(key);
   return (result == ERROR_SUCCESS);
 }
 
@@ -253,7 +297,8 @@ bool Initialize(bool force) {
   if (IsNonBrowserProcess())
     return false;
 
-  // Check to see if a beacon is present, abort if so.
+  // Check to see if the blacklist beacon is still set to running (indicating a
+  // failure) or disabled, and abort if so.
   if (!force && !LeaveSetupBeacon())
     return false;
 
@@ -268,30 +313,6 @@ bool Initialize(bool force) {
   if (!thunk)
     return false;
 
-  // Record that we are starting the thunk setup code.
-  HKEY key = NULL;
-  DWORD disposition = 0;
-  LONG result = ::RegCreateKeyEx(HKEY_CURRENT_USER,
-                                 kRegistryBeaconPath,
-                                 0,
-                                 NULL,
-                                 REG_OPTION_NON_VOLATILE,
-                                 KEY_QUERY_VALUE | KEY_SET_VALUE,
-                                 NULL,
-                                 &key,
-                                 &disposition);
-  if (result == ERROR_SUCCESS) {
-    DWORD blacklist_state = BLACKLIST_THUNK_SETUP;
-    ::RegSetValueEx(key,
-                    kBeaconState,
-                    0,
-                    REG_DWORD,
-                    reinterpret_cast<LPBYTE>(&blacklist_state),
-                    sizeof(blacklist_state));
-  } else {
-    key = NULL;
-  }
-
   BYTE* thunk_storage = reinterpret_cast<BYTE*>(&g_thunk_storage);
 
   // Mark the thunk storage as readable and writeable, since we
@@ -301,7 +322,6 @@ bool Initialize(bool force) {
                       sizeof(g_thunk_storage),
                       PAGE_EXECUTE_READWRITE,
                       &old_protect)) {
-    RecordSuccessfulThunkSetup(&key);
     return false;
   }
 
@@ -353,8 +373,6 @@ bool Initialize(bool force) {
                                                       PAGE_EXECUTE_READ,
                                                       &old_protect);
 
-  RecordSuccessfulThunkSetup(&key);
-
   AddDllsFromRegistryToBlacklist();
 
   return NT_SUCCESS(ret) && page_executable;
diff --git a/chrome_elf/blacklist/blacklist.h b/chrome_elf/blacklist/blacklist.h
index 7b8d7a6..58cc16b 100644
--- a/chrome_elf/blacklist/blacklist.h
+++ b/chrome_elf/blacklist/blacklist.h
@@ -21,16 +21,16 @@ extern const wchar_t* g_troublesome_dlls[kTroublesomeDllsMaxCount];
 extern NtMapViewOfSectionFunction g_nt_map_view_of_section_func;
 #endif
 
-// Attempts to leave a beacon in the current user's registry hive.
-// If the blacklist beacon doesn't say it is enabled or there are any other
-// errors when creating the beacon, returns false. Otherwise returns true.
-// The intent of the beacon is to act as an extra failure mode protection
-// whereby if Chrome for some reason fails to start during blacklist setup,
-// it will skip blacklisting on the subsequent run.
+// Attempts to leave a beacon in the current user's registry hive. If the
+// blacklist beacon doesn't say it is enabled or there are any other errors when
+// creating the beacon, returns false. Otherwise returns true. The intent of the
+// beacon is to act as an extra failure mode protection whereby if Chrome
+// repeatedly fails to start during blacklist setup, it will skip blacklisting
+// on the subsequent run.
 bool LeaveSetupBeacon();
 
-// Looks for the beacon that LeaveSetupBeacon() creates and resets it to
-// to show the setup was successful.
+// Looks for the setup running beacon that LeaveSetupBeacon() creates and resets
+// it to to show the setup was successful.
 // Returns true if the beacon was successfully set to BLACKLIST_ENABLED.
 bool ResetBeacon();
 
diff --git a/chrome_elf/blacklist/test/blacklist_test.cc b/chrome_elf/blacklist/test/blacklist_test.cc
index 1f91956..2e3f9a4 100644
--- a/chrome_elf/blacklist/test/blacklist_test.cc
+++ b/chrome_elf/blacklist/test/blacklist_test.cc
@@ -42,10 +42,25 @@ __declspec(dllimport) bool TestDll_SuccessfullyBlocked(
     int* size);
 }
 
+namespace {
+
 class BlacklistTest : public testing::Test {
+ protected:
+  BlacklistTest() : override_manager_() {
+    override_manager_.OverrideRegistry(HKEY_CURRENT_USER, L"beacon_test");
+  }
+
+  scoped_ptr<base::win::RegKey> blacklist_registry_key_;
+  registry_util::RegistryOverrideManager override_manager_;
+
+ private:
   virtual void SetUp() {
     // Force an import from blacklist_test_main_dll.
     InitBlacklistTestDll();
+    blacklist_registry_key_.reset(
+        new base::win::RegKey(HKEY_CURRENT_USER,
+                              blacklist::kRegistryBeaconPath,
+                              KEY_QUERY_VALUE | KEY_SET_VALUE));
   }
 
   virtual void TearDown() {
@@ -55,28 +70,19 @@ class BlacklistTest : public testing::Test {
   }
 };
 
-namespace {
-
 struct TestData {
   const wchar_t* dll_name;
   const wchar_t* dll_beacon;
 } test_data[] = {{kTestDllName2, kDll2Beacon}, {kTestDllName3, kDll3Beacon}};
 
 TEST_F(BlacklistTest, Beacon) {
-  registry_util::RegistryOverrideManager override_manager;
-  override_manager.OverrideRegistry(HKEY_CURRENT_USER, L"beacon_test");
-
-  base::win::RegKey blacklist_registry_key(HKEY_CURRENT_USER,
-                                           blacklist::kRegistryBeaconPath,
-                                           KEY_QUERY_VALUE | KEY_SET_VALUE);
-
-  // Ensure that the beacon state starts off enabled for this version.
-  LONG result = blacklist_registry_key.WriteValue(blacklist::kBeaconState,
-                                                  blacklist::BLACKLIST_ENABLED);
+  // Ensure that the beacon state starts off 'running' for this version.
+  LONG result = blacklist_registry_key_->WriteValue(
+      blacklist::kBeaconState, blacklist::BLACKLIST_SETUP_RUNNING);
   EXPECT_EQ(ERROR_SUCCESS, result);
 
-  result = blacklist_registry_key.WriteValue(blacklist::kBeaconVersion,
-                                             TEXT(CHROME_VERSION_STRING));
+  result = blacklist_registry_key_->WriteValue(blacklist::kBeaconVersion,
+                                               TEXT(CHROME_VERSION_STRING));
   EXPECT_EQ(ERROR_SUCCESS, result);
 
   // First call should find the beacon and reset it.
@@ -84,12 +90,6 @@ TEST_F(BlacklistTest, Beacon) {
 
   // First call should succeed as the beacon is enabled.
   EXPECT_TRUE(blacklist::LeaveSetupBeacon());
-
-  // Second call should fail indicating the beacon wasn't set as enabled.
-  EXPECT_FALSE(blacklist::LeaveSetupBeacon());
-
-  // Resetting the beacon should work when setup beacon is present.
-  EXPECT_TRUE(blacklist::ResetBeacon());
 }
 
 TEST_F(BlacklistTest, AddAndRemoveModules) {
@@ -244,4 +244,83 @@ TEST_F(BlacklistTest, AddDllsFromRegistryToBlacklist) {
   CheckBlacklistedDllsNotLoaded();
 }
 
+void TestResetBeacon(scoped_ptr<base::win::RegKey>& key,
+                     DWORD input_state,
+                     DWORD expected_output_state) {
+  LONG result = key->WriteValue(blacklist::kBeaconState, input_state);
+  EXPECT_EQ(ERROR_SUCCESS, result);
+
+  EXPECT_TRUE(blacklist::ResetBeacon());
+  DWORD blacklist_state = blacklist::BLACKLIST_STATE_MAX;
+  result = key->ReadValueDW(blacklist::kBeaconState, &blacklist_state);
+  EXPECT_EQ(ERROR_SUCCESS, result);
+  EXPECT_EQ(expected_output_state, blacklist_state);
+}
+
+TEST_F(BlacklistTest, ResetBeacon) {
+  // Ensure that ResetBeacon resets properly on successful runs and not on
+  // failed or disabled runs.
+  TestResetBeacon(blacklist_registry_key_,
+                  blacklist::BLACKLIST_SETUP_RUNNING,
+                  blacklist::BLACKLIST_ENABLED);
+
+  TestResetBeacon(blacklist_registry_key_,
+                  blacklist::BLACKLIST_SETUP_FAILED,
+                  blacklist::BLACKLIST_SETUP_FAILED);
+
+  TestResetBeacon(blacklist_registry_key_,
+                  blacklist::BLACKLIST_DISABLED,
+                  blacklist::BLACKLIST_DISABLED);
+}
+
+TEST_F(BlacklistTest, SetupFailed) {
+  // Ensure that when the number of failed tries reaches the maximum allowed,
+  // the blacklist state is set to failed.
+  LONG result = blacklist_registry_key_->WriteValue(
+      blacklist::kBeaconState, blacklist::BLACKLIST_SETUP_RUNNING);
+  EXPECT_EQ(ERROR_SUCCESS, result);
+
+  // Set the attempt count so that on the next failure the blacklist is
+  // disabled.
+  result = blacklist_registry_key_->WriteValue(
+      blacklist::kBeaconAttemptCount, blacklist::kBeaconMaxAttempts - 1);
+  EXPECT_EQ(ERROR_SUCCESS, result);
+
+  EXPECT_FALSE(blacklist::LeaveSetupBeacon());
+
+  DWORD attempt_count = 0;
+  blacklist_registry_key_->ReadValueDW(blacklist::kBeaconAttemptCount,
+                                       &attempt_count);
+  EXPECT_EQ(attempt_count, blacklist::kBeaconMaxAttempts);
+
+  DWORD blacklist_state = blacklist::BLACKLIST_STATE_MAX;
+  result = blacklist_registry_key_->ReadValueDW(blacklist::kBeaconState,
+                                                &blacklist_state);
+  EXPECT_EQ(ERROR_SUCCESS, result);
+  EXPECT_EQ(blacklist_state, blacklist::BLACKLIST_SETUP_FAILED);
+}
+
+TEST_F(BlacklistTest, SetupSucceeded) {
+  // Starting with the enabled beacon should result in the setup running state
+  // and the attempt counter reset to zero.
+  LONG result = blacklist_registry_key_->WriteValue(
+      blacklist::kBeaconState, blacklist::BLACKLIST_ENABLED);
+  EXPECT_EQ(ERROR_SUCCESS, result);
+  result = blacklist_registry_key_->WriteValue(blacklist::kBeaconAttemptCount,
+                                               blacklist::kBeaconMaxAttempts);
+  EXPECT_EQ(ERROR_SUCCESS, result);
+
+  EXPECT_TRUE(blacklist::LeaveSetupBeacon());
+
+  DWORD blacklist_state = blacklist::BLACKLIST_STATE_MAX;
+  blacklist_registry_key_->ReadValueDW(blacklist::kBeaconState,
+                                       &blacklist_state);
+  EXPECT_EQ(blacklist_state, blacklist::BLACKLIST_SETUP_RUNNING);
+
+  DWORD attempt_count = blacklist::kBeaconMaxAttempts;
+  blacklist_registry_key_->ReadValueDW(blacklist::kBeaconAttemptCount,
+                                       &attempt_count);
+  EXPECT_EQ(static_cast<DWORD>(0), attempt_count);
+}
+
 }  // namespace
diff --git a/chrome_elf/chrome_elf_constants.cc b/chrome_elf/chrome_elf_constants.cc
index 1dbc749..a7ee1c9 100644
--- a/chrome_elf/chrome_elf_constants.cc
+++ b/chrome_elf/chrome_elf_constants.cc
@@ -21,5 +21,8 @@ const wchar_t kRegistryFinchListPath[] =
     L"SOFTWARE\\Google\\Chrome\\BLFinchList";
 const wchar_t kBeaconVersion[] = L"version";
 const wchar_t kBeaconState[] = L"state";
+const wchar_t kBeaconAttemptCount[] = L"failed_count";
+
+const DWORD kBeaconMaxAttempts = 2;
 
 }  // namespace blacklist
diff --git a/chrome_elf/chrome_elf_constants.h b/chrome_elf/chrome_elf_constants.h
index 4ff45b2..92e44ca 100644
--- a/chrome_elf/chrome_elf_constants.h
+++ b/chrome_elf/chrome_elf_constants.h
@@ -7,6 +7,8 @@
 #ifndef CHROME_ELF_CHROME_ELF_CONSTANTS_H_
 #define CHROME_ELF_CHROME_ELF_CONSTANTS_H_
 
+#include <windows.h>
+
 // directory names
 extern const wchar_t kAppDataDirName[];
 extern const wchar_t kCanaryAppDataDirName[];
@@ -25,20 +27,21 @@ extern const wchar_t kRegistryFinchListPath[];
 // The properties for the blacklist beacon.
 extern const wchar_t kBeaconVersion[];
 extern const wchar_t kBeaconState[];
+extern const wchar_t kBeaconAttemptCount[];
+
+// The number of failures that can occur on startup with the beacon enabled
+// before we give up and turn off the blacklist.
+extern const DWORD kBeaconMaxAttempts;
 
 // The states for the blacklist setup code.
 enum BlacklistState {
   BLACKLIST_DISABLED = 0,
   BLACKLIST_ENABLED,
-  // The blacklist setup code is running. If this is still set at startup,
-  // it means the last setup crashed.
+  // The blacklist setup code is running. If this is the state at startup, it
+  // means the last setup crashed.
   BLACKLIST_SETUP_RUNNING,
-  // The blacklist thunk setup code is running. If this is still set at startup,
-  // it means the last setup crashed during thunk setup.
-  BLACKLIST_THUNK_SETUP,
-  // The blacklist code is currently intercepting MapViewOfSection. If this is
-  // still set at startup, it means we crashed during interception.
-  BLACKLIST_INTERCEPTING,
+  // If the last setup crashed, we reassign the state to failed.
+  BLACKLIST_SETUP_FAILED,
   // Always keep this at the end.
   BLACKLIST_STATE_MAX,
 };