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 0928f753dcf671ba854088b2a2ab65b3336a72c5..fa968eb927e6f618bc4c89dbee35f136c710f32f 100644 |
--- a/chrome_elf/blacklist/test/blacklist_test.cc |
+++ b/chrome_elf/blacklist/test/blacklist_test.cc |
@@ -7,6 +7,7 @@ |
#include <memory> |
#include "base/environment.h" |
+#include "base/files/file.h" |
#include "base/files/file_path.h" |
#include "base/files/scoped_temp_dir.h" |
#include "base/i18n/case_conversion.h" |
@@ -21,8 +22,8 @@ |
#include "base/win/registry.h" |
#include "chrome/common/chrome_version.h" |
#include "chrome_elf/blacklist/blacklist.h" |
-#include "chrome_elf/blacklist/test/blacklist_test_main_dll.h" |
#include "chrome_elf/chrome_elf_constants.h" |
+#include "chrome_elf/chrome_elf_reg.h" |
#include "testing/gtest/include/gtest/gtest.h" |
const wchar_t kTestDllName1[] = L"blacklist_test_dll_1.dll"; |
@@ -34,24 +35,31 @@ const wchar_t kDll3Beacon[] = L"{9E056AEC-169E-400c-B2D0-5A07E3ACE2EB}"; |
extern const wchar_t* kEnvVars[]; |
-extern "C" { |
-// When modifying the blacklist in the test process, use the exported test dll |
-// functions on the test blacklist dll, not the ones linked into the test |
-// executable itself. |
-__declspec(dllimport) void TestDll_AddDllsFromRegistryToBlacklist(); |
-__declspec(dllimport) bool TestDll_AddDllToBlacklist(const wchar_t* dll_name); |
-__declspec(dllimport) int TestDll_BlacklistSize(); |
-__declspec(dllimport) void TestDll_BlockedDll(size_t blocked_index); |
-__declspec(dllimport) int TestDll_GetBlacklistIndex(const wchar_t* dll_name); |
-__declspec(dllimport) bool TestDll_IsBlacklistInitialized(); |
-__declspec(dllimport) bool TestDll_RemoveDllFromBlacklist( |
- const wchar_t* dll_name); |
-__declspec(dllimport) bool TestDll_SuccessfullyBlocked( |
+namespace { |
+ |
+// Functions we need from blacklist_test_main_dll.dll |
+typedef void (*TestDll_AddDllsFromRegistryToBlacklistFunction)(); |
+typedef bool (*TestDll_AddDllToBlacklistFunction)(const wchar_t* dll_name); |
+typedef int (*TestDll_BlacklistSizeFunction)(); |
+typedef void (*TestDll_BlockedDllFunction)(size_t blocked_index); |
+typedef int (*TestDll_GetBlacklistIndexFunction)(const wchar_t* dll_name); |
+typedef bool (*TestDll_IsBlacklistInitializedFunction)(); |
+typedef bool (*TestDll_RemoveDllFromBlacklistFunction)(const wchar_t* dll_name); |
+typedef bool (*TestDll_SuccessfullyBlockedFunction)( |
const wchar_t** blocked_dlls, |
int* size); |
-} |
- |
-namespace { |
+typedef void (*InitTestDllFunction)(); |
+ |
+TestDll_AddDllsFromRegistryToBlacklistFunction |
+ TestDll_AddDllsFromRegistryToBlacklist = nullptr; |
+TestDll_AddDllToBlacklistFunction TestDll_AddDllToBlacklist = nullptr; |
+TestDll_BlacklistSizeFunction TestDll_BlacklistSize = nullptr; |
+TestDll_BlockedDllFunction TestDll_BlockedDll = nullptr; |
+TestDll_GetBlacklistIndexFunction TestDll_GetBlacklistIndex = nullptr; |
+TestDll_IsBlacklistInitializedFunction TestDll_IsBlacklistInitialized = nullptr; |
+TestDll_RemoveDllFromBlacklistFunction TestDll_RemoveDllFromBlacklist = nullptr; |
+TestDll_SuccessfullyBlockedFunction TestDll_SuccessfullyBlocked = nullptr; |
+InitTestDllFunction InitTestDll = nullptr; |
struct TestData { |
const wchar_t* dll_name; |
@@ -64,7 +72,6 @@ struct TestData { |
class BlacklistTest : public testing::Test { |
protected: |
BlacklistTest() : override_manager_(), num_initially_blocked_(0) { |
- override_manager_.OverrideRegistry(HKEY_CURRENT_USER); |
} |
void CheckBlacklistedDllsNotLoaded() { |
@@ -126,9 +133,90 @@ class BlacklistTest : public testing::Test { |
int num_initially_blocked_; |
private: |
+ // This function puts registry-key redirection paths into |
+ // a process-specific temp file, for our test DLLs to access. |
+ void IpcOverrides() { |
+ wchar_t* file_name = L"chrome_elf_tests.txt"; |
robertshield
2016/04/20 05:16:09
static const
penny
2016/05/28 01:34:22
Done.
|
+ |
+ wchar_t path[MAX_PATH] = {}; |
+ ::wsprintf(path, L"chrome_elf_test_%u\\", ::GetCurrentProcessId()); |
robertshield
2016/04/20 15:45:47
Please also use process creation time when naming
penny
2016/05/28 01:34:22
See main comment in reply.
|
+ |
+ base::FilePath temp_dir; |
+ ASSERT_TRUE(base::PathService::Get(base::DIR_TEMP, &temp_dir)); |
+ temp_dir = temp_dir.Append(path); |
+ |
+ ASSERT_TRUE(reg_override_dir_.Set(temp_dir)); |
robertshield
2016/04/20 15:45:47
I think this will delete the user's %TEMP% folder,
penny
2016/05/28 01:34:22
It's all good actually - doesn't touch %TEMP%.
re
|
+ |
+ temp_dir = temp_dir.Append(file_name); |
+ // This file will be deleted when the ScopedTempDir is cleaned up. |
+ base::File file(temp_dir, |
+ base::File::FLAG_OPEN_ALWAYS | base::File::FLAG_WRITE); |
+ ASSERT_TRUE(file.IsValid()); |
+ |
+ std::wstring content; |
+ if (!nt::HKCU_override.empty()) { |
+ content.append(L"hkcu=\""); |
+ content.append(nt::HKCU_override.c_str()); |
+ content.append(L"\"\n"); |
+ } |
+ if (!nt::HKLM_override.empty()) { |
+ content.append(L"hklm=\""); |
+ content.append(nt::HKLM_override.c_str()); |
+ content.append(L"\"\n"); |
+ } |
+ ASSERT_NE(-1, file.Write(0, (char*)content.c_str(), |
+ (int)(content.length() * sizeof(wchar_t)))); |
+ } |
+ |
void SetUp() override { |
- // Force an import from blacklist_test_main_dll. |
- InitBlacklistTestDll(); |
+ override_manager_.OverrideRegistry(HKEY_CURRENT_USER, &nt::HKCU_override); |
+ |
+ // Make the override path available to our test DLL. |
+ IpcOverrides(); |
+ |
+ // Load the main test Dll now. |
+ // Note: this has to happen after we set up the registry overrides. |
+ HMODULE dll = nullptr; |
+ dll = ::LoadLibraryW(L"blacklist_test_main_dll.dll"); |
+ if (!dll) |
+ return; |
+ TestDll_AddDllsFromRegistryToBlacklist = |
+ reinterpret_cast<TestDll_AddDllsFromRegistryToBlacklistFunction>( |
+ ::GetProcAddress(dll, "TestDll_AddDllsFromRegistryToBlacklist")); |
+ TestDll_AddDllToBlacklist = |
+ reinterpret_cast<TestDll_AddDllToBlacklistFunction>( |
+ ::GetProcAddress(dll, "TestDll_AddDllToBlacklist")); |
+ TestDll_BlacklistSize = reinterpret_cast<TestDll_BlacklistSizeFunction>( |
+ ::GetProcAddress(dll, "TestDll_BlacklistSize")); |
+ TestDll_BlockedDll = reinterpret_cast<TestDll_BlockedDllFunction>( |
+ ::GetProcAddress(dll, "TestDll_BlockedDll")); |
+ TestDll_GetBlacklistIndex = |
+ reinterpret_cast<TestDll_GetBlacklistIndexFunction>( |
+ ::GetProcAddress(dll, "TestDll_GetBlacklistIndex")); |
+ TestDll_IsBlacklistInitialized = |
+ reinterpret_cast<TestDll_IsBlacklistInitializedFunction>( |
+ ::GetProcAddress(dll, "TestDll_IsBlacklistInitialized")); |
+ TestDll_RemoveDllFromBlacklist = |
+ reinterpret_cast<TestDll_RemoveDllFromBlacklistFunction>( |
+ ::GetProcAddress(dll, "TestDll_RemoveDllFromBlacklist")); |
+ TestDll_SuccessfullyBlocked = |
+ reinterpret_cast<TestDll_SuccessfullyBlockedFunction>( |
+ ::GetProcAddress(dll, "TestDll_SuccessfullyBlocked")); |
+ InitTestDll = reinterpret_cast<InitTestDllFunction>( |
+ ::GetProcAddress(dll, "InitTestDll")); |
+ if (!TestDll_AddDllsFromRegistryToBlacklist || !TestDll_AddDllToBlacklist || |
+ !TestDll_BlacklistSize || !TestDll_BlockedDll || |
+ !TestDll_GetBlacklistIndex || !TestDll_IsBlacklistInitialized || |
+ !TestDll_RemoveDllFromBlacklist || !TestDll_SuccessfullyBlocked || |
+ !InitTestDll) |
+ return; |
+ |
+ // We have to call this exported function every time this test setup runs. |
+ // If the tests are running in single process mode, the test DLL does not |
+ // get reloaded everytime - but we need to make sure it updates |
+ // appropriately. |
+ InitTestDll(); |
+ |
blacklist_registry_key_.reset( |
new base::win::RegKey(HKEY_CURRENT_USER, |
blacklist::kRegistryBeaconPath, |
@@ -143,6 +231,9 @@ class BlacklistTest : public testing::Test { |
TestDll_RemoveDllFromBlacklist(kTestDllName2); |
TestDll_RemoveDllFromBlacklist(kTestDllName3); |
} |
+ |
+ // A scoped temporary directory to be destroyed with this test. |
+ base::ScopedTempDir reg_override_dir_; |
}; |
TEST_F(BlacklistTest, Beacon) { |
@@ -255,15 +346,27 @@ TEST_F(BlacklistTest, AddDllsFromRegistryToBlacklist) { |
KEY_QUERY_VALUE | KEY_SET_VALUE); |
key.DeleteKey(L""); |
- // Add the test dlls to the registry (with their name as both key and value). |
+ // Add the test dlls to the registry. |
+ // (REG_MULTI_SZ: eos separated, double null terminated.) |
base::win::RegKey finch_blacklist_registry_key( |
HKEY_CURRENT_USER, |
blacklist::kRegistryFinchListPath, |
KEY_QUERY_VALUE | KEY_SET_VALUE); |
+ |
+ std::vector<wchar_t>(reg_buffer); |
for (size_t i = 0; i < arraysize(test_data); ++i) { |
- finch_blacklist_registry_key.WriteValue(test_data[i].dll_name, |
- test_data[i].dll_name); |
+ if (reg_buffer.size() > 0) |
+ reg_buffer.push_back(L'\0'); |
+ const wchar_t* dll = test_data[i].dll_name; |
+ // Append the name, not including terminator. |
+ reg_buffer.insert(reg_buffer.end(), dll, dll + ::wcslen(dll)); |
} |
+ reg_buffer.push_back(L'\0'); |
+ reg_buffer.push_back(L'\0'); |
+ |
+ finch_blacklist_registry_key.WriteValue( |
+ blacklist::kRegistryFinchListValueName, reg_buffer.data(), |
+ (DWORD)(reg_buffer.size() * sizeof(wchar_t)), REG_MULTI_SZ); |
TestDll_AddDllsFromRegistryToBlacklist(); |
CheckBlacklistedDllsNotLoaded(); |