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