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..fbdeb0a88b744591f64e6776363d1ca015845d94 100644
--- a/chrome_elf/blacklist/test/blacklist_test.cc
+++ b/chrome_elf/blacklist/test/blacklist_test.cc
@@ -41,34 +41,44 @@ __declspec(dllimport) bool TestDll_SuccessfullyBlocked(
     int* size);
 }
 
+namespace {
+
 class BlacklistTest : public testing::Test {
+ public:

[csharp, 2014/06/11 13:34:43]

protected is generally used here instead of public.

[krstnmnlsn, 2014/06/11 16:35:27]

Done.

+  base::win::RegKey* blacklist_registry_key;

[csharp, 2014/06/11 13:34:43]

Why not just have them as non-pointer members?

Also, member variables should have their names end with _

[krstnmnlsn, 2014/06/11 16:35:27]

I think the best I can do is making override_manager a non-pointer.
If blacklist_registry_key is a member either:
(a) It gets default constructed, and then we don't seem to get a public
assignment operator, so I can't fix it?
(b) I write a constructor with an initialization list, but
OverrideRegistry(HKEY_CURRENT_USER, L"beacon_test") needs to run before
blacklist_registry_key's constructor and I don't think that's possible?

Is it worth writing a constructor to make override_manager not a pointer?

[csharp, 2014/06/11 17:07:13]

I wouldn't worry about not being able to make blacklist_registry_key a pointer.

I would say adding the constructor is worth it for override_manager_ no longer
being a pointer

+  registry_util::RegistryOverrideManager* override_manager;
+
+ private:
   virtual void SetUp() {
     // Force an import from blacklist_test_main_dll.
     InitBlacklistTestDll();
+
+    // Setup registry overrides and registry key.
+    override_manager = new registry_util::RegistryOverrideManager();
+    override_manager->OverrideRegistry(HKEY_CURRENT_USER, L"beacon_test");
+    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 override_manager;
+    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.
@@ -76,12 +86,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 +211,84 @@ 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());
+  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