1 // Copyright (c) 2011 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 #include "base/test/test_reg_util_win.h"
6
7 #include <stdint.h>
8
9 #include "base/guid.h"
10 #include "base/memory/ptr_util.h"
11 #include "base/strings/strcat.h"
12 #include "base/strings/string_number_conversions.h"
13 #include "base/strings/string_split.h"
14 #include "base/strings/string_util.h"
15 #include "base/strings/utf_string_conversions.h"
16 #include "testing/gtest/include/gtest/gtest.h"
17
18 #include <windows.h>
19
20 namespace registry_util {
21
22 namespace {
23
24 constexpr base::char16 kTimestampDelimiter[] = STRING16_LITERAL("$");
25 constexpr wchar_t kTempTestKeyPath[] = L"Software\\Chromium\\TempTestKeys";
26
DeleteStaleTestKeys(const base::Time & now,const std::wstring & test_key_root)27 void DeleteStaleTestKeys(const base::Time& now,
28 const std::wstring& test_key_root) {
29 base::win::RegKey test_root_key;
30 if (test_root_key.Open(HKEY_CURRENT_USER,
31 test_key_root.c_str(),
32 KEY_ALL_ACCESS) != ERROR_SUCCESS) {
33 // This will occur on first-run, but is harmless.
34 return;
35 }
36
37 base::win::RegistryKeyIterator iterator_test_root_key(HKEY_CURRENT_USER,
38 test_key_root.c_str());
39 for (; iterator_test_root_key.Valid(); ++iterator_test_root_key) {
40 std::wstring key_name = iterator_test_root_key.Name();
41 std::vector<base::StringPiece16> tokens = base::SplitStringPiece(
42 base::AsStringPiece16(key_name), kTimestampDelimiter,
43 base::KEEP_WHITESPACE, base::SPLIT_WANT_NONEMPTY);
44 if (tokens.empty())
45 continue;
46 int64_t key_name_as_number = 0;
47
48 if (!base::StringToInt64(tokens[0], &key_name_as_number)) {
49 test_root_key.DeleteKey(key_name.c_str());
50 continue;
51 }
52
53 base::Time key_time = base::Time::FromInternalValue(key_name_as_number);
54 base::TimeDelta age = now - key_time;
55
56 if (age > base::TimeDelta::FromHours(24))
57 test_root_key.DeleteKey(key_name.c_str());
58 }
59 }
60
GenerateTempKeyPath(const std::wstring & test_key_root,const base::Time & timestamp)61 std::wstring GenerateTempKeyPath(const std::wstring& test_key_root,
62 const base::Time& timestamp) {
63 return base::AsWString(base::StrCat(
64 {base::AsStringPiece16(test_key_root), STRING16_LITERAL("\\"),
65 base::NumberToString16(timestamp.ToInternalValue()), kTimestampDelimiter,
66 base::ASCIIToUTF16(base::GenerateGUID())}));
67 }
68
69 } // namespace
70
ScopedRegistryKeyOverride(HKEY override,const std::wstring & key_path)71 RegistryOverrideManager::ScopedRegistryKeyOverride::ScopedRegistryKeyOverride(
72 HKEY override,
73 const std::wstring& key_path)
74 : override_(override), key_path_(key_path) {}
75
76 RegistryOverrideManager::
~ScopedRegistryKeyOverride()77 ScopedRegistryKeyOverride::~ScopedRegistryKeyOverride() {
78 ::RegOverridePredefKey(override_, NULL);
79 base::win::RegKey(HKEY_CURRENT_USER, L"", KEY_QUERY_VALUE)
80 .DeleteKey(key_path_.c_str());
81 }
82
RegistryOverrideManager()83 RegistryOverrideManager::RegistryOverrideManager()
84 : timestamp_(base::Time::Now()), test_key_root_(kTempTestKeyPath) {
85 DeleteStaleTestKeys(timestamp_, test_key_root_);
86 }
87
RegistryOverrideManager(const base::Time & timestamp,const std::wstring & test_key_root)88 RegistryOverrideManager::RegistryOverrideManager(
89 const base::Time& timestamp,
90 const std::wstring& test_key_root)
91 : timestamp_(timestamp), test_key_root_(test_key_root) {
92 DeleteStaleTestKeys(timestamp_, test_key_root_);
93 }
94
~RegistryOverrideManager()95 RegistryOverrideManager::~RegistryOverrideManager() {}
96
OverrideRegistry(HKEY override)97 void RegistryOverrideManager::OverrideRegistry(HKEY override) {
98 OverrideRegistry(override, nullptr);
99 }
100
OverrideRegistry(HKEY override,std::wstring * override_path)101 void RegistryOverrideManager::OverrideRegistry(HKEY override,
102 std::wstring* override_path) {
103 std::wstring key_path = GenerateTempKeyPath(test_key_root_, timestamp_);
104
105 base::win::RegKey temp_key;
106 ASSERT_EQ(ERROR_SUCCESS, temp_key.Create(HKEY_CURRENT_USER, key_path.c_str(),
107 KEY_ALL_ACCESS));
108 ASSERT_EQ(ERROR_SUCCESS, ::RegOverridePredefKey(override, temp_key.Handle()));
109
110 overrides_.push_back(
111 std::make_unique<ScopedRegistryKeyOverride>(override, key_path));
112 if (override_path)
113 override_path->assign(key_path);
114 }
115
GenerateTempKeyPath()116 std::wstring GenerateTempKeyPath() {
117 return GenerateTempKeyPath(kTempTestKeyPath, base::Time::Now());
118 }
119
120 } // namespace registry_util
121