rlz/test/rlz_test_helpers.cc

   1 // Copyright (c) 2012 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 // Main entry point for all unit tests.
   6
   7 #include "rlz_test_helpers.h"
   8
   9 #include <map>
  10 #include <vector>
  11
  12 #include "base/strings/string16.h"
  13 #include "rlz/lib/rlz_lib.h"
  14 #include "testing/gtest/include/gtest/gtest.h"
  15
  16 #if defined(OS_WIN)
  17 #include <shlwapi.h>
  18 #include "base/win/registry.h"
  19 #include "base/win/windows_version.h"
  20 #elif defined(OS_POSIX)
  21 #include "base/files/file_path.h"
  22 #include "rlz/lib/rlz_value_store.h"
  23 #endif
  24
  25 #if defined(OS_WIN)
  26
  27 namespace {
  28
  29 // Path to recursively copy into the replacemment hives.  These are needed
  30 // to make sure certain win32 APIs continue to run correctly once the real
  31 // hives are replaced.
  32 const wchar_t kHKLMAccessProviders[] =
  33     L"System\\CurrentControlSet\\Control\\Lsa\\AccessProviders";
  34
  35 struct RegistryValue {
  36   base::string16 name;
  37   DWORD type;
  38   std::vector<uint8> data;
  39 };
  40
  41 struct RegistryKeyData {
  42   std::vector<RegistryValue> values;
  43   std::map<base::string16, RegistryKeyData> keys;
  44 };
  45
  46 void ReadRegistryTree(const base::win::RegKey& src, RegistryKeyData* data) {
  47   // First read values.
  48   {
  49     base::win::RegistryValueIterator i(src.Handle(), L"");
  50     data->values.clear();
  51     data->values.reserve(i.ValueCount());
  52     for (; i.Valid(); ++i) {
  53       RegistryValue& value = *data->values.insert(data->values.end(),
  54                                                   RegistryValue());
  55       const uint8* data = reinterpret_cast<const uint8*>(i.Value());
  56       value.name.assign(i.Name());
  57       value.type = i.Type();
  58       value.data.assign(data, data + i.ValueSize());
  59     }
  60   }
  61
  62   // Next read subkeys recursively.
  63   for (base::win::RegistryKeyIterator i(src.Handle(), L"");
  64        i.Valid(); ++i) {
  65     ReadRegistryTree(base::win::RegKey(src.Handle(), i.Name(), KEY_READ),
  66                      &data->keys[base::string16(i.Name())]);
  67   }
  68 }
  69
  70 void WriteRegistryTree(const RegistryKeyData& data, base::win::RegKey* dest) {
  71   // First write values.
  72   for (size_t i = 0; i < data.values.size(); ++i) {
  73     const RegistryValue& value = data.values[i];
  74     dest->WriteValue(value.name.c_str(),
  75                      value.data.size() ? &value.data[0] : NULL,
  76                      static_cast<DWORD>(value.data.size()),
  77                      value.type);
  78   }
  79
  80   // Next write values recursively.
  81   for (std::map<base::string16, RegistryKeyData>::const_iterator iter =
  82            data.keys.begin();
  83        iter != data.keys.end(); ++iter) {
  84     base::win::RegKey key(dest->Handle(), iter->first.c_str(), KEY_ALL_ACCESS);
  85     WriteRegistryTree(iter->second, &key);
  86   }
  87 }
  88
  89 // Initialize temporary HKLM/HKCU registry hives used for testing.
  90 // Testing RLZ requires reading and writing to the Windows registry.  To keep
  91 // the tests isolated from the machine's state, as well as to prevent the tests
  92 // from causing side effects in the registry, HKCU and HKLM are overridden for
  93 // the duration of the tests. RLZ tests don't expect the HKCU and KHLM hives to
  94 // be empty though, and this function initializes the minimum value needed so
  95 // that the test will run successfully.
  96 void InitializeRegistryOverridesForTesting(
  97     registry_util::RegistryOverrideManager* override_manager) {
  98   // For the moment, the HKCU hive requires no initialization.
  99   const bool do_copy = (base::win::GetVersion() >= base::win::VERSION_WIN7);
 100   RegistryKeyData data;
 101
 102   if (do_copy) {
 103     // Copy the following HKLM subtrees to the temporary location so that the
 104     // win32 APIs used by the tests continue to work:
 105     //
 106     //    HKLM\System\CurrentControlSet\Control\Lsa\AccessProviders
 107     //
 108     // This seems to be required since Win7.
 109     ReadRegistryTree(base::win::RegKey(HKEY_LOCAL_MACHINE,
 110                                        kHKLMAccessProviders,
 111                                        KEY_READ), &data);
 112   }
 113
 114   override_manager->OverrideRegistry(HKEY_LOCAL_MACHINE);
 115   override_manager->OverrideRegistry(HKEY_CURRENT_USER);
 116
 117   if (do_copy) {
 118     base::win::RegKey key(
 119         HKEY_LOCAL_MACHINE, kHKLMAccessProviders, KEY_ALL_ACCESS);
 120     WriteRegistryTree(data, &key);
 121   }
 122 }
 123
 124 }  // namespace
 125
 126 #endif  // defined(OS_WIN)
 127
 128 void RlzLibTestNoMachineState::SetUp() {
 129 #if defined(OS_WIN)
 130   InitializeRegistryOverridesForTesting(&override_manager_);
 131 #elif defined(OS_MACOSX)
 132   base::mac::ScopedNSAutoreleasePool pool;
 133 #endif  // defined(OS_WIN)
 134 #if defined(OS_POSIX)
 135   ASSERT_TRUE(temp_dir_.CreateUniqueTempDir());
 136   rlz_lib::testing::SetRlzStoreDirectory(temp_dir_.path());
 137 #endif  // defined(OS_POSIX)
 138 }
 139
 140 void RlzLibTestNoMachineState::TearDown() {
 141 #if defined(OS_POSIX)
 142   rlz_lib::testing::SetRlzStoreDirectory(base::FilePath());
 143 #endif  // defined(OS_POSIX)
 144 }
 145
 146 void RlzLibTestBase::SetUp() {
 147   RlzLibTestNoMachineState::SetUp();
 148 #if defined(OS_WIN)
 149   rlz_lib::CreateMachineState();
 150 #endif  // defined(OS_WIN)
 151 }