1 // Copyright (c) 2013 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 "chrome/common/importer/importer_test_registry_overrider_win.h"
6 
7 #include <windows.h>
8 
9 #include <memory>
10 #include <string>
11 
12 #include "base/environment.h"
13 #include "base/guid.h"
14 #include "base/strings/utf_string_conversions.h"
15 #include "base/win/registry.h"
16 
17 namespace {
18 
19 // The key to which a random subkey will be appended. This key itself will never
20 // be deleted.
21 const wchar_t kTestHKCUOverrideKeyPrefix[] = L"SOFTWARE\\Chromium Unit Tests\\";
22 const char kTestHKCUOverrideEnvironmentVariable[] =
23     "IE_IMPORTER_TEST_OVERRIDE_HKCU";
24 
25 // Reads the environment variable set by a previous call to
26 // SetTestRegistryOverride() into |key| if it exists and |key| is not NULL.
27 // Returns true if the variable was successfully read.
GetTestKeyFromEnvironment(base::string16 * key)28 bool GetTestKeyFromEnvironment(base::string16* key) {
29   std::unique_ptr<base::Environment> env(base::Environment::Create());
30   std::string value;
31   bool result = env->GetVar(kTestHKCUOverrideEnvironmentVariable, &value);
32   if (result)
33     *key = base::UTF8ToUTF16(value);
34   return result;
35 }
36 
37 }  // namespace
38 
39 ////////////////////////////////////////////////////////////////////////////////
40 // ImporterTestRegistryOverrider, public:
41 
ImporterTestRegistryOverrider()42 ImporterTestRegistryOverrider::ImporterTestRegistryOverrider()
43     : temporary_key_(kTestHKCUOverrideKeyPrefix +
44                      base::UTF8ToUTF16(base::GenerateGUID())) {
45   DCHECK(!GetTestKeyFromEnvironment(NULL));
46 
47   std::unique_ptr<base::Environment> env(base::Environment::Create());
48   bool success = env->SetVar(kTestHKCUOverrideEnvironmentVariable,
49                              base::UTF16ToUTF8(temporary_key_));
50   DCHECK(success);
51 }
52 
~ImporterTestRegistryOverrider()53 ImporterTestRegistryOverrider::~ImporterTestRegistryOverrider() {
54   base::win::RegKey reg_key(HKEY_CURRENT_USER, temporary_key_.c_str(),
55                             KEY_ALL_ACCESS);
56   DCHECK(reg_key.Valid());
57   reg_key.DeleteKey(L"");
58 
59   std::unique_ptr<base::Environment> env(base::Environment::Create());
60   bool success = env->UnSetVar(kTestHKCUOverrideEnvironmentVariable);
61   DCHECK(success);
62 }
63 
64 // static
GetTestRegistryOverride()65 base::string16 ImporterTestRegistryOverrider::GetTestRegistryOverride() {
66   base::string16 key;
67   if (!GetTestKeyFromEnvironment(&key))
68     return base::string16();
69   return key;
70 }
71