[Chrome ELF] New NT registry API.

chrome_elf/blacklist/test/blacklist_test.cc

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 the file.

[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, which isn't what is wanted.
Maybe you want to call CreateUniqueTempDir() instead, and have the file you
stash information in live in there instead?

[penny, 2016/05/28 01:34:22]

It's all good actually - doesn't touch %TEMP%.

reg_override_dir_ is a base::ScopedTempDir.

scoped_temp_dir.h:
  // Takes ownership of directory at |path|, creating it if necessary.
  // Don't call multiple times unless Take() has been called first.
  bool Set(const FilePath& path) WARN_UNUSED_RESULT;

The contents of the base::FilePath when set() is called is
DIR_TEMP\chrome_elf_test_PID\.  This directory is created by the call to Set(),
and is deleted during destruction of the BlacklistTest object (at end of scope).
 DIR_TEMP itself isn't touched.

Side note, I'm not using CreateUniqueTempDir, because I want to decide the name
of the temp directory (so that I can easily recreate that temp name in another
testing DLL, without having to communicate the temp dir location).

+
+    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();