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

chrome_elf/blacklist/test/blacklist_test.cc

Index: chrome_elf/blacklist/test/blacklist_test.cc
diff --git a/chrome_elf/blacklist/test/blacklist_test.cc b/chrome_elf/blacklist/test/blacklist_test.cc
index 05aeb14f1b123611451b66dd3111cc8665f8ef85..e038d4f22e3d1c5666d5edc7f95517c54a73005d 100644
--- a/chrome_elf/blacklist/test/blacklist_test.cc
+++ b/chrome_elf/blacklist/test/blacklist_test.cc
@@ -41,6 +41,8 @@ __declspec(dllimport) bool TestDll_SuccessfullyBlocked(
     int* size);
 }
 
+namespace {
+
 class BlacklistTest : public testing::Test {
   virtual void SetUp() {
     // Force an import from blacklist_test_main_dll.
@@ -54,6 +56,13 @@ class BlacklistTest : public testing::Test {
   }
 };
 
+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");
@@ -76,12 +85,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) {
@@ -207,3 +210,100 @@ TEST_F(BlacklistTest, LoadBlacklistedLibrary) {
     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.