[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 7eb18f2ab2872f4c5a42b47f5b621f8af5413a1f..8eab4de037ba4f5ba8f0889ff5573a643e35a309 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"
@@ -20,8 +21,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/nt_registry/nt_registry.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 const wchar_t kTestDllName1[] = L"blacklist_test_dll_1.dll";
@@ -33,24 +34,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;
@@ -63,7 +71,6 @@ struct TestData {
 class BlacklistTest : public testing::Test {
  protected:
   BlacklistTest() : override_manager_(), num_initially_blocked_(0) {
-    override_manager_.OverrideRegistry(HKEY_CURRENT_USER);
   }
 
   void CheckBlacklistedDllsNotLoaded() {
@@ -125,9 +132,69 @@ class BlacklistTest : public testing::Test {
   int num_initially_blocked_;
 
  private:
+  // This function puts registry-key redirection paths into
+  // process-specific environment variables, for our test DLLs to access.
+  // This will only work as long as the IPC is within the same process.
+  void IpcOverrides() {
+    if (!nt::HKCU_override.empty()) {
+      ASSERT_TRUE(::SetEnvironmentVariableW(L"hkcu_override",
+                                            nt::HKCU_override.c_str()));
+    }
+    if (!nt::HKLM_override.empty()) {
+      ASSERT_TRUE(::SetEnvironmentVariableW(L"hklm_override",
+                                            nt::HKLM_override.c_str()));
+    }
+  }
+
   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,
@@ -142,6 +209,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) {
@@ -254,15 +324,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();