1 // Copyright 2011 The Chromium Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #include "base/test/test_reg_util_win.h"
6
7 #include <stdint.h>
8
9 #include "base/memory/ptr_util.h"
10 #include "base/strings/strcat.h"
11 #include "base/strings/string_number_conversions.h"
12 #include "base/strings/string_split.h"
13 #include "base/strings/string_util.h"
14 #include "base/strings/utf_string_conversions.h"
15 #include "base/time/time_override.h"
16 #include "base/uuid.h"
17 #include "testing/gtest/include/gtest/gtest.h"
18
19 #include <windows.h>
20
21 namespace registry_util {
22
23 namespace {
24
25 // Overriding HKLM is not permitted in some environments. This is controlled by
26 // this bool and disallowed by calling
27 // DisallowHKLMRegistryOverrideForIntegrationTests.
28 bool g_hklm_override_allowed = true;
29
30 constexpr char16_t kTimestampDelimiter[] = u"$";
31 constexpr wchar_t kTempTestKeyPath[] = L"Software\\Chromium\\TempTestKeys";
32
DeleteStaleTestKeys(const base::Time & now,const std::wstring & test_key_root)33 void DeleteStaleTestKeys(const base::Time& now,
34 const std::wstring& test_key_root) {
35 base::win::RegKey test_root_key;
36 if (test_root_key.Open(HKEY_CURRENT_USER,
37 test_key_root.c_str(),
38 KEY_ALL_ACCESS) != ERROR_SUCCESS) {
39 // This will occur on first-run, but is harmless.
40 return;
41 }
42
43 base::win::RegistryKeyIterator iterator_test_root_key(HKEY_CURRENT_USER,
44 test_key_root.c_str());
45 for (; iterator_test_root_key.Valid(); ++iterator_test_root_key) {
46 std::wstring key_name = iterator_test_root_key.Name();
47 std::vector<base::StringPiece16> tokens = base::SplitStringPiece(
48 base::AsStringPiece16(key_name), kTimestampDelimiter,
49 base::KEEP_WHITESPACE, base::SPLIT_WANT_NONEMPTY);
50 if (tokens.empty())
51 continue;
52 int64_t key_name_as_number = 0;
53
54 if (!base::StringToInt64(tokens[0], &key_name_as_number)) {
55 test_root_key.DeleteKey(key_name.c_str());
56 continue;
57 }
58
59 base::Time key_time = base::Time::FromInternalValue(key_name_as_number);
60 base::TimeDelta age = now - key_time;
61
62 if (age > base::Hours(24))
63 test_root_key.DeleteKey(key_name.c_str());
64 }
65 }
66
GenerateTempKeyPath(const std::wstring & test_key_root,const base::Time & timestamp)67 std::wstring GenerateTempKeyPath(const std::wstring& test_key_root,
68 const base::Time& timestamp) {
69 return base::AsWString(base::StrCat(
70 {base::AsStringPiece16(test_key_root), u"\\",
71 base::NumberToString16(timestamp.ToInternalValue()), kTimestampDelimiter,
72 base::ASCIIToUTF16(
73 base::Uuid::GenerateRandomV4().AsLowercaseString())}));
74 }
75
76 } // namespace
77
ScopedRegistryKeyOverride(HKEY override,const std::wstring & key_path)78 RegistryOverrideManager::ScopedRegistryKeyOverride::ScopedRegistryKeyOverride(
79 HKEY override,
80 const std::wstring& key_path)
81 : override_(override), key_path_(key_path) {}
82
83 RegistryOverrideManager::
~ScopedRegistryKeyOverride()84 ScopedRegistryKeyOverride::~ScopedRegistryKeyOverride() {
85 ::RegOverridePredefKey(override_, NULL);
86 base::win::RegKey(HKEY_CURRENT_USER, L"", KEY_QUERY_VALUE)
87 .DeleteKey(key_path_.c_str());
88 }
89
RegistryOverrideManager()90 RegistryOverrideManager::RegistryOverrideManager()
91 : timestamp_(base::subtle::TimeNowIgnoringOverride()),
92 test_key_root_(kTempTestKeyPath) {
93 // Use |base::subtle::TimeNowIgnoringOverride()| instead of
94 // |base::Time::Now()| can give us the real current time instead of the mock
95 // time in 1970 when MOCK_TIME is enabled. This can prevent test bugs where
96 // new instances of RegistryOverrideManager will clean up any redirected
97 // registry paths that have the timestamp from 1970, which then cause the
98 // currently running tests to fail since their expected reg keys were deleted
99 // by the other test.
100 DeleteStaleTestKeys(timestamp_, test_key_root_);
101 }
102
RegistryOverrideManager(const base::Time & timestamp,const std::wstring & test_key_root)103 RegistryOverrideManager::RegistryOverrideManager(
104 const base::Time& timestamp,
105 const std::wstring& test_key_root)
106 : timestamp_(timestamp), test_key_root_(test_key_root) {
107 DeleteStaleTestKeys(timestamp_, test_key_root_);
108 }
109
~RegistryOverrideManager()110 RegistryOverrideManager::~RegistryOverrideManager() {}
111
OverrideRegistry(HKEY override)112 void RegistryOverrideManager::OverrideRegistry(HKEY override) {
113 OverrideRegistry(override, nullptr);
114 }
115
OverrideRegistry(HKEY override,std::wstring * override_path)116 void RegistryOverrideManager::OverrideRegistry(HKEY override,
117 std::wstring* override_path) {
118 CHECK(override != HKEY_LOCAL_MACHINE || g_hklm_override_allowed)
119 << "Use of RegistryOverrideManager to override HKLM is not permitted in "
120 "this environment.";
121
122 std::wstring key_path = GenerateTempKeyPath(test_key_root_, timestamp_);
123
124 base::win::RegKey temp_key;
125 ASSERT_EQ(ERROR_SUCCESS, temp_key.Create(HKEY_CURRENT_USER, key_path.c_str(),
126 KEY_ALL_ACCESS));
127 ASSERT_EQ(ERROR_SUCCESS, ::RegOverridePredefKey(override, temp_key.Handle()));
128
129 overrides_.push_back(
130 std::make_unique<ScopedRegistryKeyOverride>(override, key_path));
131 if (override_path)
132 override_path->assign(key_path);
133 }
134
SetAllowHKLMRegistryOverrideForIntegrationTests(bool allow)135 void RegistryOverrideManager::SetAllowHKLMRegistryOverrideForIntegrationTests(
136 bool allow) {
137 g_hklm_override_allowed = allow;
138 }
139
GenerateTempKeyPath()140 std::wstring GenerateTempKeyPath() {
141 return GenerateTempKeyPath(kTempTestKeyPath,
142 base::subtle::TimeNowIgnoringOverride());
143 }
144
145 } // namespace registry_util
146