test_reg_util_win.cc 4.0 KB

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