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 1f91956ab72006bf18c72c6a8c59a5d0f532d8c7..2e3f9a4af66b14935ef9b438d481559a7c2003dd 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 |