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 {
   virtual void SetUp() {
     // Force an import from blacklist_test_main_dll.
     InitBlacklistTestDll();
   }
 
   virtual void TearDown() {
     TestDll_RemoveDllFromBlacklist(kTestDllName1);
     TestDll_RemoveDllFromBlacklist(kTestDllName2);
     TestDll_RemoveDllFromBlacklist(kTestDllName3);
   }
 };
 
+void ResetRegistry(base::win::RegKey* blacklist_registry_key) {
+  blacklist_registry_key->WriteValue(blacklist::kBeaconState,
+                                     blacklist::BLACKLIST_ENABLED);
+  blacklist_registry_key->WriteValue(blacklist::kBeaconAttemptCount,
+                                     static_cast<DWORD>(0));
+}
+
 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);
   EXPECT_EQ(ERROR_SUCCESS, result);
 
   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());

[csharp, 2014/06/10 14:07:02]

Nit: Please add a blank line above

[krstnmnlsn, 2014/06/10 22:03:27]

Done.

+  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.
+  base::win::RegKey blacklist_registry_key(HKEY_CURRENT_USER,

[csharp, 2014/06/10 14:07:02]

A lot of tests seem to be declaring this, is it worth making it a class variable
and just declaring it once in the constructor?

[krstnmnlsn, 2014/06/10 22:03:27]

Done.

+                                           blacklist::kRegistryBeaconPath,
+                                           KEY_QUERY_VALUE | KEY_SET_VALUE);
+
+  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);
+
+  ResetRegistry(&blacklist_registry_key);
+}
+
+TEST_F(BlacklistTest, SetupFailed) {
+  // Ensure that when the number of failed tries reaches the maximum allowed the

[csharp, 2014/06/10 14:07:02]

nit: allowed -> allowed,

[krstnmnlsn, 2014/06/10 22:03:27]

Done.

+  // blacklist state is set to failed.
+  base::win::RegKey blacklist_registry_key(HKEY_CURRENT_USER,
+                                           blacklist::kRegistryBeaconPath,
+                                           KEY_QUERY_VALUE | KEY_SET_VALUE);
+
+  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);
+
+  ResetRegistry(&blacklist_registry_key);
+}
+
+TEST_F(BlacklistTest, SetupSucceeded) {
+  // Starting with the enabled beacon should result in the setup running state
+  // and the attempt counter reset to zero.
+  base::win::RegKey blacklist_registry_key(HKEY_CURRENT_USER,
+                                           blacklist::kRegistryBeaconPath,
+                                           KEY_QUERY_VALUE | KEY_SET_VALUE);
+
+  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);
+
+  ResetRegistry(&blacklist_registry_key);
+}
+
+}  // namespace

[csharp, 2014/06/10 14:07:02]

Please add a newline to the end of the file (lint is complaining about it)

[krstnmnlsn, 2014/06/10 22:03:27]

Checked my previous commit (which also complained about this) and it seems to be
just the CL.