1// Copyright 2013 The Chromium Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#include "base/environment.h"
6#include "base/files/file_path.h"
7#include "base/files/scoped_temp_dir.h"
8#include "base/i18n/case_conversion.h"
9#include "base/path_service.h"
10#include "base/scoped_native_library.h"
11#include "base/strings/string16.h"
12#include "base/strings/string_number_conversions.h"
13#include "base/strings/utf_string_conversions.h"
14#include "base/test/test_reg_util_win.h"
15#include "base/win/registry.h"
16#include "chrome_elf/blacklist/blacklist.h"
17#include "chrome_elf/blacklist/test/blacklist_test_main_dll.h"
18#include "chrome_elf/chrome_elf_constants.h"
19#include "testing/gtest/include/gtest/gtest.h"
20#include "version.h"  // NOLINT
21
22const wchar_t kTestDllName1[] = L"blacklist_test_dll_1.dll";
23const wchar_t kTestDllName2[] = L"blacklist_test_dll_2.dll";
24const wchar_t kTestDllName3[] = L"blacklist_test_dll_3.dll";
25
26const wchar_t kDll2Beacon[] = L"{F70A0100-2889-4629-9B44-610FE5C73231}";
27const wchar_t kDll3Beacon[] = L"{9E056AEC-169E-400c-B2D0-5A07E3ACE2EB}";
28
29extern const wchar_t* kEnvVars[];
30
31extern "C" {
32// When modifying the blacklist in the test process, use the exported test dll
33// functions on the test blacklist dll, not the ones linked into the test
34// executable itself.
35__declspec(dllimport) void TestDll_AddDllsFromRegistryToBlacklist();
36__declspec(dllimport) bool TestDll_AddDllToBlacklist(const wchar_t* dll_name);
37__declspec(dllimport) bool TestDll_IsBlacklistInitialized();
38__declspec(dllimport) bool TestDll_RemoveDllFromBlacklist(
39    const wchar_t* dll_name);
40__declspec(dllimport) bool TestDll_SuccessfullyBlocked(
41    const wchar_t** blocked_dlls,
42    int* size);
43}
44
45namespace {
46
47class BlacklistTest : public testing::Test {
48 protected:
49  BlacklistTest() : override_manager_() {
50    override_manager_.OverrideRegistry(HKEY_CURRENT_USER);
51  }
52
53  scoped_ptr<base::win::RegKey> blacklist_registry_key_;
54  registry_util::RegistryOverrideManager override_manager_;
55
56 private:
57  virtual void SetUp() {
58    // Force an import from blacklist_test_main_dll.
59    InitBlacklistTestDll();
60    blacklist_registry_key_.reset(
61        new base::win::RegKey(HKEY_CURRENT_USER,
62                              blacklist::kRegistryBeaconPath,
63                              KEY_QUERY_VALUE | KEY_SET_VALUE));
64  }
65
66  virtual void TearDown() {
67    TestDll_RemoveDllFromBlacklist(kTestDllName1);
68    TestDll_RemoveDllFromBlacklist(kTestDllName2);
69    TestDll_RemoveDllFromBlacklist(kTestDllName3);
70  }
71};
72
73struct TestData {
74  const wchar_t* dll_name;
75  const wchar_t* dll_beacon;
76} test_data[] = {
77    {kTestDllName2, kDll2Beacon},
78#if !defined(_WIN64)
79    // The third test dll is special in that it does not contain an export
80    // table. This prevents SafeGetImageInfo from extracting the name from there
81    // AND for some reason NtQueryVirtualMemory with MemorySectionName returns
82    // STATUS_ACCESS_VIOLATION in 64 bit builds for reasons still unknown.
83    // http://crbug.com/397137
84    {kTestDllName3, kDll3Beacon}
85#endif
86};
87
88TEST_F(BlacklistTest, Beacon) {
89  // Ensure that the beacon state starts off 'running' for this version.
90  LONG result = blacklist_registry_key_->WriteValue(
91      blacklist::kBeaconState, blacklist::BLACKLIST_SETUP_RUNNING);
92  EXPECT_EQ(ERROR_SUCCESS, result);
93
94  result = blacklist_registry_key_->WriteValue(blacklist::kBeaconVersion,
95                                               TEXT(CHROME_VERSION_STRING));
96  EXPECT_EQ(ERROR_SUCCESS, result);
97
98  // First call should find the beacon and reset it.
99  EXPECT_TRUE(blacklist::ResetBeacon());
100
101  // First call should succeed as the beacon is enabled.
102  EXPECT_TRUE(blacklist::LeaveSetupBeacon());
103}
104
105TEST_F(BlacklistTest, AddAndRemoveModules) {
106  EXPECT_TRUE(blacklist::AddDllToBlacklist(L"foo.dll"));
107  // Adding the same item twice should be idempotent.
108  EXPECT_TRUE(blacklist::AddDllToBlacklist(L"foo.dll"));
109  EXPECT_TRUE(blacklist::RemoveDllFromBlacklist(L"foo.dll"));
110  EXPECT_FALSE(blacklist::RemoveDllFromBlacklist(L"foo.dll"));
111
112  // Increase the blacklist size by 1 to include the NULL pointer
113  // that marks the end.
114  int empty_spaces = blacklist::kTroublesomeDllsMaxCount - (
115      blacklist::BlacklistSize() + 1);
116  std::vector<base::string16> added_dlls;
117  added_dlls.reserve(empty_spaces);
118  for (int i = 0; i < empty_spaces; ++i) {
119    added_dlls.push_back(base::IntToString16(i) + L".dll");
120    EXPECT_TRUE(blacklist::AddDllToBlacklist(added_dlls[i].c_str())) << i;
121  }
122  EXPECT_FALSE(blacklist::AddDllToBlacklist(L"overflow.dll"));
123  for (int i = 0; i < empty_spaces; ++i) {
124    EXPECT_TRUE(blacklist::RemoveDllFromBlacklist(added_dlls[i].c_str())) << i;
125  }
126  EXPECT_FALSE(blacklist::RemoveDllFromBlacklist(added_dlls[0].c_str()));
127  EXPECT_FALSE(blacklist::RemoveDllFromBlacklist(
128    added_dlls[empty_spaces - 1].c_str()));
129}
130
131TEST_F(BlacklistTest, SuccessfullyBlocked) {
132  // Ensure that we have at least 5 dlls to blacklist.
133  int blacklist_size = blacklist::BlacklistSize();
134  const int kDesiredBlacklistSize = 5;
135  for (int i = blacklist_size; i < kDesiredBlacklistSize; ++i) {
136    base::string16 new_dll_name(base::IntToString16(i) + L".dll");
137    EXPECT_TRUE(blacklist::AddDllToBlacklist(new_dll_name.c_str()));
138  }
139
140  // Block 5 dlls, one at a time, starting from the end of the list, and
141  // ensuring SuccesfullyBlocked correctly passes the list of blocked dlls.
142  for (int i = 0; i < kDesiredBlacklistSize; ++i) {
143    blacklist::BlockedDll(i);
144
145    int size = 0;
146    blacklist::SuccessfullyBlocked(NULL, &size);
147    EXPECT_EQ(i + 1, size);
148
149    std::vector<const wchar_t*> blocked_dlls(size);
150    blacklist::SuccessfullyBlocked(&(blocked_dlls[0]), &size);
151    EXPECT_EQ(i + 1, size);
152
153    for (size_t j = 0; j < blocked_dlls.size(); ++j) {
154      EXPECT_EQ(blocked_dlls[j], blacklist::g_troublesome_dlls[j]);
155    }
156  }
157}
158
159void CheckBlacklistedDllsNotLoaded() {
160  base::FilePath current_dir;
161  ASSERT_TRUE(PathService::Get(base::DIR_EXE, &current_dir));
162
163  for (int i = 0; i < arraysize(test_data); ++i) {
164    // Ensure that the dll has not been loaded both by inspecting the handle
165    // returned by LoadLibrary and by looking for an environment variable that
166    // is set when the DLL's entry point is called.
167    base::ScopedNativeLibrary dll_blacklisted(
168        current_dir.Append(test_data[i].dll_name));
169    EXPECT_FALSE(dll_blacklisted.is_valid());
170    EXPECT_EQ(0u, ::GetEnvironmentVariable(test_data[i].dll_beacon, NULL, 0));
171    dll_blacklisted.Reset(NULL);
172
173    // Ensure that the dll is recorded as blocked.
174    int array_size = 1;
175    const wchar_t* blocked_dll = NULL;
176    TestDll_SuccessfullyBlocked(&blocked_dll, &array_size);
177    EXPECT_EQ(1, array_size);
178    EXPECT_EQ(test_data[i].dll_name, base::string16(blocked_dll));
179
180    // Remove the DLL from the blacklist. Ensure that it loads and that its
181    // entry point was called.
182    EXPECT_TRUE(TestDll_RemoveDllFromBlacklist(test_data[i].dll_name));
183    base::ScopedNativeLibrary dll(current_dir.Append(test_data[i].dll_name));
184    EXPECT_TRUE(dll.is_valid());
185    EXPECT_NE(0u, ::GetEnvironmentVariable(test_data[i].dll_beacon, NULL, 0));
186    dll.Reset(NULL);
187
188    ::SetEnvironmentVariable(test_data[i].dll_beacon, NULL);
189
190    // Ensure that the dll won't load even if the name has different
191    // capitalization.
192    base::string16 uppercase_name = base::i18n::ToUpper(test_data[i].dll_name);
193    EXPECT_TRUE(TestDll_AddDllToBlacklist(uppercase_name.c_str()));
194    base::ScopedNativeLibrary dll_blacklisted_different_case(
195        current_dir.Append(test_data[i].dll_name));
196    EXPECT_FALSE(dll_blacklisted_different_case.is_valid());
197    EXPECT_EQ(0u, ::GetEnvironmentVariable(test_data[i].dll_beacon, NULL, 0));
198    dll_blacklisted_different_case.Reset(NULL);
199
200    EXPECT_TRUE(TestDll_RemoveDllFromBlacklist(uppercase_name.c_str()));
201
202    // The blocked dll was removed, so we shouldn't get anything returned
203    // here.
204    int num_blocked_dlls = 0;
205    TestDll_SuccessfullyBlocked(NULL, &num_blocked_dlls);
206    EXPECT_EQ(0, num_blocked_dlls);
207  }
208}
209
210TEST_F(BlacklistTest, LoadBlacklistedLibrary) {
211  base::FilePath current_dir;
212  ASSERT_TRUE(PathService::Get(base::DIR_EXE, &current_dir));
213
214  // Ensure that the blacklist is loaded.
215  ASSERT_TRUE(TestDll_IsBlacklistInitialized());
216
217  // Test that an un-blacklisted DLL can load correctly.
218  base::ScopedNativeLibrary dll1(current_dir.Append(kTestDllName1));
219  EXPECT_TRUE(dll1.is_valid());
220  dll1.Reset(NULL);
221
222  int num_blocked_dlls = 0;
223  TestDll_SuccessfullyBlocked(NULL, &num_blocked_dlls);
224  EXPECT_EQ(0, num_blocked_dlls);
225
226  // Add all DLLs to the blacklist then check they are blocked.
227  for (int i = 0; i < arraysize(test_data); ++i) {
228    EXPECT_TRUE(TestDll_AddDllToBlacklist(test_data[i].dll_name));
229  }
230  CheckBlacklistedDllsNotLoaded();
231}
232
233TEST_F(BlacklistTest, AddDllsFromRegistryToBlacklist) {
234  // Ensure that the blacklist is loaded.
235  ASSERT_TRUE(TestDll_IsBlacklistInitialized());
236
237  // Delete the finch registry key to clear its values.
238  base::win::RegKey key(HKEY_CURRENT_USER,
239                        blacklist::kRegistryFinchListPath,
240                        KEY_QUERY_VALUE | KEY_SET_VALUE);
241  key.DeleteKey(L"");
242
243  // Add the test dlls to the registry (with their name as both key and value).
244  base::win::RegKey finch_blacklist_registry_key(
245      HKEY_CURRENT_USER,
246      blacklist::kRegistryFinchListPath,
247      KEY_QUERY_VALUE | KEY_SET_VALUE);
248  for (int i = 0; i < arraysize(test_data); ++i) {
249    finch_blacklist_registry_key.WriteValue(test_data[i].dll_name,
250                                            test_data[i].dll_name);
251  }
252
253  TestDll_AddDllsFromRegistryToBlacklist();
254  CheckBlacklistedDllsNotLoaded();
255}
256
257void TestResetBeacon(scoped_ptr<base::win::RegKey>& key,
258                     DWORD input_state,
259                     DWORD expected_output_state) {
260  LONG result = key->WriteValue(blacklist::kBeaconState, input_state);
261  EXPECT_EQ(ERROR_SUCCESS, result);
262
263  EXPECT_TRUE(blacklist::ResetBeacon());
264  DWORD blacklist_state = blacklist::BLACKLIST_STATE_MAX;
265  result = key->ReadValueDW(blacklist::kBeaconState, &blacklist_state);
266  EXPECT_EQ(ERROR_SUCCESS, result);
267  EXPECT_EQ(expected_output_state, blacklist_state);
268}
269
270TEST_F(BlacklistTest, ResetBeacon) {
271  // Ensure that ResetBeacon resets properly on successful runs and not on
272  // failed or disabled runs.
273  TestResetBeacon(blacklist_registry_key_,
274                  blacklist::BLACKLIST_SETUP_RUNNING,
275                  blacklist::BLACKLIST_ENABLED);
276
277  TestResetBeacon(blacklist_registry_key_,
278                  blacklist::BLACKLIST_SETUP_FAILED,
279                  blacklist::BLACKLIST_SETUP_FAILED);
280
281  TestResetBeacon(blacklist_registry_key_,
282                  blacklist::BLACKLIST_DISABLED,
283                  blacklist::BLACKLIST_DISABLED);
284}
285
286TEST_F(BlacklistTest, SetupFailed) {
287  // Ensure that when the number of failed tries reaches the maximum allowed,
288  // the blacklist state is set to failed.
289  LONG result = blacklist_registry_key_->WriteValue(
290      blacklist::kBeaconState, blacklist::BLACKLIST_SETUP_RUNNING);
291  EXPECT_EQ(ERROR_SUCCESS, result);
292
293  // Set the attempt count so that on the next failure the blacklist is
294  // disabled.
295  result = blacklist_registry_key_->WriteValue(
296      blacklist::kBeaconAttemptCount, blacklist::kBeaconMaxAttempts - 1);
297  EXPECT_EQ(ERROR_SUCCESS, result);
298
299  EXPECT_FALSE(blacklist::LeaveSetupBeacon());
300
301  DWORD attempt_count = 0;
302  blacklist_registry_key_->ReadValueDW(blacklist::kBeaconAttemptCount,
303                                       &attempt_count);
304  EXPECT_EQ(attempt_count, blacklist::kBeaconMaxAttempts);
305
306  DWORD blacklist_state = blacklist::BLACKLIST_STATE_MAX;
307  result = blacklist_registry_key_->ReadValueDW(blacklist::kBeaconState,
308                                                &blacklist_state);
309  EXPECT_EQ(ERROR_SUCCESS, result);
310  EXPECT_EQ(blacklist_state, blacklist::BLACKLIST_SETUP_FAILED);
311}
312
313TEST_F(BlacklistTest, SetupSucceeded) {
314  // Starting with the enabled beacon should result in the setup running state
315  // and the attempt counter reset to zero.
316  LONG result = blacklist_registry_key_->WriteValue(
317      blacklist::kBeaconState, blacklist::BLACKLIST_ENABLED);
318  EXPECT_EQ(ERROR_SUCCESS, result);
319  result = blacklist_registry_key_->WriteValue(blacklist::kBeaconAttemptCount,
320                                               blacklist::kBeaconMaxAttempts);
321  EXPECT_EQ(ERROR_SUCCESS, result);
322
323  EXPECT_TRUE(blacklist::LeaveSetupBeacon());
324
325  DWORD blacklist_state = blacklist::BLACKLIST_STATE_MAX;
326  blacklist_registry_key_->ReadValueDW(blacklist::kBeaconState,
327                                       &blacklist_state);
328  EXPECT_EQ(blacklist_state, blacklist::BLACKLIST_SETUP_RUNNING);
329
330  DWORD attempt_count = blacklist::kBeaconMaxAttempts;
331  blacklist_registry_key_->ReadValueDW(blacklist::kBeaconAttemptCount,
332                                       &attempt_count);
333  EXPECT_EQ(static_cast<DWORD>(0), attempt_count);
334}
335
336}  // namespace
337