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

chrome_elf/blacklist/test/blacklist_test.cc

 // Copyright 2013 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "base/environment.h"
 #include "base/files/file_path.h"
 #include "base/files/scoped_temp_dir.h"
 #include "base/i18n/case_conversion.h"
 #include "base/path_service.h"
 #include "base/scoped_native_library.h"
@@ skipping 23 matching lines @@
 // executable itself.
 __declspec(dllimport) bool TestDll_AddDllToBlacklist(const wchar_t* dll_name);
 __declspec(dllimport) bool TestDll_IsBlacklistInitialized();
 __declspec(dllimport) bool TestDll_RemoveDllFromBlacklist(
     const wchar_t* dll_name);
 __declspec(dllimport) bool TestDll_SuccessfullyBlocked(
     const wchar_t** blocked_dlls,
     int* size);
 }
 
+namespace {
+
 class BlacklistTest : public testing::Test {
+ protected:
+  BlacklistTest() : override_manager_() {
+    override_manager_.OverrideRegistry(HKEY_CURRENT_USER, L"beacon_test");
+  }
+
+  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_ =
+        new base::win::RegKey(HKEY_CURRENT_USER,
+                              blacklist::kRegistryBeaconPath,
+                              KEY_QUERY_VALUE | KEY_SET_VALUE);
   }
 
   virtual void TearDown() {
     TestDll_RemoveDllFromBlacklist(kTestDllName1);
     TestDll_RemoveDllFromBlacklist(kTestDllName2);
     TestDll_RemoveDllFromBlacklist(kTestDllName3);
+    delete blacklist_registry_key_;
   }
 };
 
 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.
   EXPECT_TRUE(blacklist::ResetBeacon());
 
   // 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) {
   EXPECT_TRUE(blacklist::AddDllToBlacklist(L"foo.dll"));
   // Adding the same item twice should be idempotent.
   EXPECT_TRUE(blacklist::AddDllToBlacklist(L"foo.dll"));
   EXPECT_TRUE(blacklist::RemoveDllFromBlacklist(L"foo.dll"));
   EXPECT_FALSE(blacklist::RemoveDllFromBlacklist(L"foo.dll"));
 
   // Increase the blacklist size by 1 to include the NULL pointer
@@ skipping 105 matching lines @@
     dll_blacklisted_different_case.Reset(NULL);
 
     EXPECT_TRUE(TestDll_RemoveDllFromBlacklist(uppercase_name.c_str()));
 
     // The blocked dll was removed, so we shouldn't get anything returned
     // here.
     TestDll_SuccessfullyBlocked(NULL, &num_blocked_dlls);
     EXPECT_EQ(0, num_blocked_dlls);
   }
 }
+
+void TestResetBeacon(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