[Chrome ELF] New NT registry API.

chrome_elf/blacklist/test/blacklist_test.cc

 // Copyright 2013 The Chromium Authors. All rights reserved.
 // 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 "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"
 #include "base/macros.h"
 #include "base/path_service.h"
 #include "base/scoped_native_library.h"
 #include "base/strings/string16.h"
 #include "base/strings/string_number_conversions.h"
 #include "base/strings/utf_string_conversions.h"
 #include "base/test/test_reg_util_win.h"
 #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";
 const wchar_t kTestDllName2[] = L"blacklist_test_dll_2.dll";
 const wchar_t kTestDllName3[] = L"blacklist_test_dll_3.dll";
 
 const wchar_t kDll2Beacon[] = L"{F70A0100-2889-4629-9B44-610FE5C73231}";
 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);
-}
+typedef void (*InitTestDllFunction)();
 
-namespace {
+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;
   const wchar_t* dll_beacon;
 } test_data[] = {
     {kTestDllName2, kDll2Beacon},
     {kTestDllName3, kDll3Beacon}
 };
 
 class BlacklistTest : public testing::Test {
  protected:
   BlacklistTest() : override_manager_(), num_initially_blocked_(0) {
-    override_manager_.OverrideRegistry(HKEY_CURRENT_USER);
   }
 
   void CheckBlacklistedDllsNotLoaded() {
     base::FilePath current_dir;
     ASSERT_TRUE(PathService::Get(base::DIR_EXE, &current_dir));
 
     for (size_t i = 0; i < arraysize(test_data); ++i) {
       // Ensure that the dll has not been loaded both by inspecting the handle
       // returned by LoadLibrary and by looking for an environment variable that
       // is set when the DLL's entry point is called.
@@ skipping 41 matching lines @@
     }
   }
 
   std::unique_ptr<base::win::RegKey> blacklist_registry_key_;
   registry_util::RegistryOverrideManager override_manager_;
 
   // The number of dlls initially blocked by the blacklist.
   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,
                               KEY_QUERY_VALUE | KEY_SET_VALUE));
 
     // Find out how many dlls were blocked before the test starts.
     TestDll_SuccessfullyBlocked(NULL, &num_initially_blocked_);
   }
 
   void TearDown() override {
     TestDll_RemoveDllFromBlacklist(kTestDllName1);
     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) {
   // Ensure that the beacon state starts off 'running' for this version.
   LONG result = blacklist_registry_key_->WriteValue(
       blacklist::kBeaconState, blacklist::BLACKLIST_SETUP_RUNNING);
   EXPECT_EQ(ERROR_SUCCESS, result);
 
   result = blacklist_registry_key_->WriteValue(blacklist::kBeaconVersion,
                                                TEXT(CHROME_VERSION_STRING));
@@ skipping 92 matching lines @@
 TEST_F(BlacklistTest, AddDllsFromRegistryToBlacklist) {
   // Ensure that the blacklist is loaded.
   ASSERT_TRUE(TestDll_IsBlacklistInitialized());
 
   // Delete the finch registry key to clear its values.
   base::win::RegKey key(HKEY_CURRENT_USER,
                         blacklist::kRegistryFinchListPath,
                         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();
 }
 
 void TestResetBeacon(std::unique_ptr<base::win::RegKey>& key,
                      DWORD input_state,
                      DWORD expected_output_state) {
   LONG result = key->WriteValue(blacklist::kBeaconState, input_state);
   EXPECT_EQ(ERROR_SUCCESS, result);
@@ skipping 65 matching lines @@
                                        &blacklist_state);
   EXPECT_EQ(blacklist_state, blacklist::BLACKLIST_SETUP_RUNNING);
 
   DWORD attempt_count = blacklist::kBeaconMaxAttempts;
   blacklist_registry_key_->ReadValueDW(blacklist::kBeaconAttemptCount,
                                        &attempt_count);
   EXPECT_EQ(static_cast<DWORD>(0), attempt_count);
 }
 
 }  // namespace