blob: ee73e0d2465a41197f87fd71d95d28d2069e97ba [file] [log] [blame]
// 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 <stddef.h>
#include <memory>
#include <string>
#include "chrome/chrome_elf/third_party_dlls/beacon.h"
#include "base/test/test_reg_util_win.h"
#include "base/win/registry.h"
#include "chrome/chrome_elf/chrome_elf_constants.h"
#include "chrome/chrome_elf/nt_registry/nt_registry.h"
#include "chrome/install_static/install_util.h"
#include "testing/gtest/include/gtest/gtest.h"
namespace third_party_dlls {
namespace {
class BeaconTest : public testing::Test {
protected:
BeaconTest() : override_manager_() {}
std::unique_ptr<base::win::RegKey> beacon_registry_key_;
registry_util::RegistryOverrideManager override_manager_;
private:
void SetUp() override {
std::wstring temp;
ASSERT_NO_FATAL_FAILURE(
override_manager_.OverrideRegistry(HKEY_CURRENT_USER, &temp));
ASSERT_TRUE(nt::SetTestingOverride(nt::HKCU, temp));
beacon_registry_key_ = std::make_unique<base::win::RegKey>(
HKEY_CURRENT_USER,
install_static::GetRegistryPath()
.append(blacklist::kRegistryBeaconKeyName)
.c_str(),
KEY_QUERY_VALUE | KEY_SET_VALUE);
}
void TearDown() override {
ASSERT_TRUE(nt::SetTestingOverride(nt::HKCU, std::wstring()));
}
};
} // namespace
//------------------------------------------------------------------------------
// Beacon tests
//------------------------------------------------------------------------------
// Ensure that the beacon state starts off 'running' if a version is specified.
TEST_F(BeaconTest, Beacon) {
LONG result = beacon_registry_key_->WriteValue(
blacklist::kBeaconState, blacklist::BLACKLIST_SETUP_RUNNING);
EXPECT_EQ(ERROR_SUCCESS, result);
result = beacon_registry_key_->WriteValue(blacklist::kBeaconVersion,
L"beacon_version");
EXPECT_EQ(ERROR_SUCCESS, result);
// First call should find the beacon and reset it.
EXPECT_TRUE(ResetBeacon());
// First call should succeed as the beacon is enabled.
EXPECT_TRUE(LeaveSetupBeacon());
}
void TestResetBeacon(std::unique_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(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(BeaconTest, ResetBeacon) {
// Ensure that ResetBeacon resets properly on successful runs and not on
// failed or disabled runs.
TestResetBeacon(beacon_registry_key_, blacklist::BLACKLIST_SETUP_RUNNING,
blacklist::BLACKLIST_ENABLED);
TestResetBeacon(beacon_registry_key_, blacklist::BLACKLIST_SETUP_FAILED,
blacklist::BLACKLIST_SETUP_FAILED);
TestResetBeacon(beacon_registry_key_, blacklist::BLACKLIST_DISABLED,
blacklist::BLACKLIST_DISABLED);
}
TEST_F(BeaconTest, SetupFailed) {
// Ensure that when the number of failed tries reaches the maximum allowed,
// the blacklist state is set to failed.
LONG result = beacon_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 = beacon_registry_key_->WriteValue(blacklist::kBeaconAttemptCount,
blacklist::kBeaconMaxAttempts - 1);
EXPECT_EQ(ERROR_SUCCESS, result);
EXPECT_FALSE(LeaveSetupBeacon());
DWORD attempt_count = 0;
beacon_registry_key_->ReadValueDW(blacklist::kBeaconAttemptCount,
&attempt_count);
EXPECT_EQ(attempt_count, blacklist::kBeaconMaxAttempts);
DWORD blacklist_state = blacklist::BLACKLIST_STATE_MAX;
result = beacon_registry_key_->ReadValueDW(blacklist::kBeaconState,
&blacklist_state);
EXPECT_EQ(ERROR_SUCCESS, result);
EXPECT_EQ(blacklist_state,
static_cast<DWORD>(blacklist::BLACKLIST_SETUP_FAILED));
}
TEST_F(BeaconTest, SetupSucceeded) {
// Starting with the enabled beacon should result in the setup running state
// and the attempt counter reset to zero.
LONG result = beacon_registry_key_->WriteValue(blacklist::kBeaconState,
blacklist::BLACKLIST_ENABLED);
EXPECT_EQ(ERROR_SUCCESS, result);
result = beacon_registry_key_->WriteValue(blacklist::kBeaconAttemptCount,
blacklist::kBeaconMaxAttempts);
EXPECT_EQ(ERROR_SUCCESS, result);
EXPECT_TRUE(LeaveSetupBeacon());
DWORD blacklist_state = blacklist::BLACKLIST_STATE_MAX;
beacon_registry_key_->ReadValueDW(blacklist::kBeaconState, &blacklist_state);
EXPECT_EQ(blacklist_state,
static_cast<DWORD>(blacklist::BLACKLIST_SETUP_RUNNING));
DWORD attempt_count = blacklist::kBeaconMaxAttempts;
beacon_registry_key_->ReadValueDW(blacklist::kBeaconAttemptCount,
&attempt_count);
EXPECT_EQ(static_cast<DWORD>(0), attempt_count);
}
} // namespace third_party_dlls