1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
|
// Copyright 2013 The Chromium Authors
// 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/blocklist_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(blocklist::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(
blocklist::kBeaconState, blocklist::BLOCKLIST_SETUP_RUNNING);
EXPECT_EQ(ERROR_SUCCESS, result);
result = beacon_registry_key_->WriteValue(blocklist::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(blocklist::kBeaconState, input_state);
EXPECT_EQ(ERROR_SUCCESS, result);
EXPECT_TRUE(ResetBeacon());
DWORD blocklist_state = blocklist::BLOCKLIST_STATE_MAX;
result = key->ReadValueDW(blocklist::kBeaconState, &blocklist_state);
EXPECT_EQ(ERROR_SUCCESS, result);
EXPECT_EQ(expected_output_state, blocklist_state);
}
TEST_F(BeaconTest, ResetBeacon) {
// Ensure that ResetBeacon resets properly on successful runs and not on
// failed or disabled runs.
TestResetBeacon(beacon_registry_key_, blocklist::BLOCKLIST_SETUP_RUNNING,
blocklist::BLOCKLIST_ENABLED);
TestResetBeacon(beacon_registry_key_, blocklist::BLOCKLIST_SETUP_FAILED,
blocklist::BLOCKLIST_SETUP_FAILED);
TestResetBeacon(beacon_registry_key_, blocklist::BLOCKLIST_DISABLED,
blocklist::BLOCKLIST_DISABLED);
}
TEST_F(BeaconTest, SetupFailed) {
// Ensure that when the number of failed tries reaches the maximum allowed,
// the blocklist state is set to failed.
LONG result = beacon_registry_key_->WriteValue(
blocklist::kBeaconState, blocklist::BLOCKLIST_SETUP_RUNNING);
EXPECT_EQ(ERROR_SUCCESS, result);
// Set the attempt count so that on the next failure the blocklist is
// disabled.
result = beacon_registry_key_->WriteValue(blocklist::kBeaconAttemptCount,
blocklist::kBeaconMaxAttempts - 1);
EXPECT_EQ(ERROR_SUCCESS, result);
EXPECT_FALSE(LeaveSetupBeacon());
DWORD attempt_count = 0;
beacon_registry_key_->ReadValueDW(blocklist::kBeaconAttemptCount,
&attempt_count);
EXPECT_EQ(attempt_count, blocklist::kBeaconMaxAttempts);
DWORD blocklist_state = blocklist::BLOCKLIST_STATE_MAX;
result = beacon_registry_key_->ReadValueDW(blocklist::kBeaconState,
&blocklist_state);
EXPECT_EQ(ERROR_SUCCESS, result);
EXPECT_EQ(blocklist_state,
static_cast<DWORD>(blocklist::BLOCKLIST_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(blocklist::kBeaconState,
blocklist::BLOCKLIST_ENABLED);
EXPECT_EQ(ERROR_SUCCESS, result);
result = beacon_registry_key_->WriteValue(blocklist::kBeaconAttemptCount,
blocklist::kBeaconMaxAttempts);
EXPECT_EQ(ERROR_SUCCESS, result);
EXPECT_TRUE(LeaveSetupBeacon());
DWORD blocklist_state = blocklist::BLOCKLIST_STATE_MAX;
beacon_registry_key_->ReadValueDW(blocklist::kBeaconState, &blocklist_state);
EXPECT_EQ(blocklist_state,
static_cast<DWORD>(blocklist::BLOCKLIST_SETUP_RUNNING));
DWORD attempt_count = blocklist::kBeaconMaxAttempts;
beacon_registry_key_->ReadValueDW(blocklist::kBeaconAttemptCount,
&attempt_count);
EXPECT_EQ(static_cast<DWORD>(0), attempt_count);
}
} // namespace third_party_dlls
|
