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 "chrome/installer/util/registry_key_backup.h"
6
7#include <algorithm>
8#include <map>
9#include <utility>
10#include <vector>
11
12#include "base/logging.h"
13#include "base/win/registry.h"
14
15using base::win::RegKey;
16
17namespace {
18
19const REGSAM kKeyReadNoNotify = (KEY_READ) & ~(KEY_NOTIFY);
20
21// A container for a registry value.
22class ValueData {
23 public:
24  ValueData();
25  ~ValueData();
26
27  // Initializes this object with a name (the first |name_size| characters in
28  // |name_buffer|, |type|, and data (the first |data_size| bytes in |data|).
29  void Initialize(const wchar_t* name_buffer, DWORD name_size,
30                  DWORD type, const uint8* data, DWORD data_size);
31
32  // The possibly empty name of this value.
33  const std::wstring& name_str() const { return name_; }
34
35  // The name of this value, or NULL for the default (unnamed) value.
36  const wchar_t* name() const { return name_.empty() ? NULL : name_.c_str(); }
37
38  // The type of this value.
39  DWORD type() const { return type_; }
40
41  // A pointer to a buffer of |data_len()| bytes containing the value's data,
42  // or NULL if the value has no data.
43  const uint8* data() const { return data_.empty() ? NULL : &data_[0]; }
44
45  // The size, in bytes, of the value's data.
46  DWORD data_len() const { return static_cast<DWORD>(data_.size()); }
47
48 private:
49  // This value's name, or the empty string if this is the default (unnamed)
50  // value.
51  std::wstring name_;
52  // This value's data.
53  std::vector<uint8> data_;
54  // This value's type (e.g., REG_DWORD, REG_SZ, REG_QWORD, etc).
55  DWORD type_;
56
57  // Copy constructible and assignable for use in STL containers.
58};
59
60}  // namespace
61
62// A container for a registry key, its values, and its subkeys.
63class RegistryKeyBackup::KeyData {
64 public:
65  KeyData();
66  ~KeyData();
67
68  // Initializes this object by reading the values and subkeys of |key|.
69  // Security descriptors are not backed up.  Returns true if the operation was
70  // successful; false otherwise, in which case the state of this object is not
71  // modified.
72  bool Initialize(const RegKey& key);
73
74  // Writes the contents of this object to |key|, which must have been opened
75  // with at least REG_SET_VALUE and KEY_CREATE_SUB_KEY access rights.  Returns
76  // true if the operation was successful; false otherwise, in which case the
77  // contents of |key| may have been modified.
78  bool WriteTo(RegKey* key) const;
79
80 private:
81  // The values of this key.
82  std::vector<ValueData> values_;
83  // Map of subkey names to the corresponding KeyData.
84  std::map<std::wstring, KeyData> subkeys_;
85
86  // Copy constructible and assignable for use in STL containers.
87};
88
89ValueData::ValueData() : type_(REG_NONE) {
90}
91
92ValueData::~ValueData() {
93}
94
95void ValueData::Initialize(
96    const wchar_t* name_buffer,
97    DWORD name_size,
98    DWORD type,
99    const uint8* data,
100    DWORD data_size) {
101  name_.assign(name_buffer, name_size);
102  type_ = type;
103  data_.assign(data, data + data_size);
104}
105
106RegistryKeyBackup::KeyData::KeyData() {
107}
108
109RegistryKeyBackup::KeyData::~KeyData() {
110}
111
112bool RegistryKeyBackup::KeyData::Initialize(const RegKey& key) {
113  std::vector<ValueData> values;
114  std::map<std::wstring, KeyData> subkeys;
115
116  DWORD num_subkeys = 0;
117  DWORD max_subkey_name_len = 0;
118  DWORD num_values = 0;
119  DWORD max_value_name_len = 0;
120  DWORD max_value_len = 0;
121  LONG result = RegQueryInfoKey(key.Handle(), NULL, NULL, NULL,
122                                &num_subkeys, &max_subkey_name_len, NULL,
123                                &num_values, &max_value_name_len,
124                                &max_value_len, NULL, NULL);
125  if (result != ERROR_SUCCESS) {
126    LOG(ERROR) << "Failed getting info of key to backup, result: " << result;
127    return false;
128  }
129  DWORD max_name_len = std::max(max_subkey_name_len, max_value_name_len) + 1;
130  std::vector<wchar_t> name_buffer(max_name_len);
131
132  // Backup the values.
133  if (num_values != 0) {
134    values.reserve(num_values);
135    std::vector<uint8> value_buffer(max_value_len != 0 ? max_value_len : 1);
136    DWORD name_size = 0;
137    DWORD value_type = REG_NONE;
138    DWORD value_size = 0;
139
140    for (DWORD i = 0; i < num_values; ) {
141      name_size = static_cast<DWORD>(name_buffer.size());
142      value_size = static_cast<DWORD>(value_buffer.size());
143      result = RegEnumValue(key.Handle(), i, &name_buffer[0], &name_size,
144                            NULL, &value_type, &value_buffer[0], &value_size);
145      switch (result) {
146        case ERROR_NO_MORE_ITEMS:
147          num_values = i;
148          break;
149        case ERROR_SUCCESS:
150          values.push_back(ValueData());
151          values.back().Initialize(&name_buffer[0], name_size, value_type,
152                                   &value_buffer[0], value_size);
153          ++i;
154          break;
155        case ERROR_MORE_DATA:
156          if (value_size > value_buffer.size())
157            value_buffer.resize(value_size);
158          // |name_size| does not include space for the terminating NULL.
159          if (name_size + 1 > name_buffer.size())
160            name_buffer.resize(name_size + 1);
161          break;
162        default:
163          LOG(ERROR) << "Failed backing up value " << i << ", result: "
164                     << result;
165          return false;
166      }
167    }
168    DLOG_IF(WARNING, RegEnumValue(key.Handle(), num_values, &name_buffer[0],
169                                  &name_size, NULL, &value_type, NULL,
170                                  NULL) != ERROR_NO_MORE_ITEMS)
171        << "Concurrent modifications to registry key during backup operation.";
172  }
173
174  // Backup the subkeys.
175  if (num_subkeys != 0) {
176    DWORD name_size = 0;
177
178    // Get the names of them.
179    for (DWORD i = 0; i < num_subkeys; ) {
180      name_size = static_cast<DWORD>(name_buffer.size());
181      result = RegEnumKeyEx(key.Handle(), i, &name_buffer[0], &name_size,
182                            NULL, NULL, NULL, NULL);
183      switch (result) {
184        case ERROR_NO_MORE_ITEMS:
185          num_subkeys = i;
186          break;
187        case ERROR_SUCCESS:
188          subkeys.insert(std::make_pair(&name_buffer[0], KeyData()));
189          ++i;
190          break;
191        case ERROR_MORE_DATA:
192          name_buffer.resize(name_size + 1);
193          break;
194        default:
195          LOG(ERROR) << "Failed getting name of subkey " << i
196                     << " for backup, result: " << result;
197          return false;
198      }
199    }
200    DLOG_IF(WARNING,
201            RegEnumKeyEx(key.Handle(), num_subkeys, NULL, &name_size, NULL,
202                         NULL, NULL, NULL) != ERROR_NO_MORE_ITEMS)
203        << "Concurrent modifications to registry key during backup operation.";
204
205    // Get their values.
206    RegKey subkey;
207    for (std::map<std::wstring, KeyData>::iterator it = subkeys.begin();
208         it != subkeys.end(); ++it) {
209      result = subkey.Open(key.Handle(), it->first.c_str(), kKeyReadNoNotify);
210      if (result != ERROR_SUCCESS) {
211        LOG(ERROR) << "Failed opening subkey \"" << it->first
212                   << "\" for backup, result: " << result;
213        return false;
214      }
215      if (!it->second.Initialize(subkey)) {
216        LOG(ERROR) << "Failed backing up subkey \"" << it->first << "\"";
217        return false;
218      }
219    }
220  }
221
222  values_.swap(values);
223  subkeys_.swap(subkeys);
224
225  return true;
226}
227
228bool RegistryKeyBackup::KeyData::WriteTo(RegKey* key) const {
229  DCHECK(key);
230
231  LONG result = ERROR_SUCCESS;
232
233  // Write the values.
234  for (std::vector<ValueData>::const_iterator it = values_.begin();
235       it != values_.end(); ++it) {
236    const ValueData& value = *it;
237    result = RegSetValueEx(key->Handle(), value.name(), 0, value.type(),
238                           value.data(), value.data_len());
239    if (result != ERROR_SUCCESS) {
240      LOG(ERROR) << "Failed writing value \"" << value.name_str()
241                 << "\", result: " << result;
242      return false;
243    }
244  }
245
246  // Write the subkeys.
247  RegKey subkey;
248  for (std::map<std::wstring, KeyData>::const_iterator it = subkeys_.begin();
249       it != subkeys_.end(); ++it) {
250    const std::wstring& name = it->first;
251
252    result = subkey.Create(key->Handle(), name.c_str(), KEY_WRITE);
253    if (result != ERROR_SUCCESS) {
254      LOG(ERROR) << "Failed creating subkey \"" << name << "\", result: "
255                 << result;
256      return false;
257    }
258    if (!it->second.WriteTo(&subkey)) {
259      LOG(ERROR) << "Failed writing subkey \"" << name << "\", result: "
260                 << result;
261      return false;
262    }
263  }
264
265  return true;
266}
267
268RegistryKeyBackup::RegistryKeyBackup() {
269}
270
271RegistryKeyBackup::~RegistryKeyBackup() {
272}
273
274bool RegistryKeyBackup::Initialize(HKEY root,
275                                   const wchar_t* key_path,
276                                   REGSAM wow64_access) {
277  DCHECK(key_path);
278  DCHECK(wow64_access == 0 ||
279         wow64_access == KEY_WOW64_32KEY ||
280         wow64_access == KEY_WOW64_64KEY);
281
282  RegKey key;
283  scoped_ptr<KeyData> key_data;
284
285  // Does the key exist?
286  LONG result = key.Open(root, key_path, kKeyReadNoNotify | wow64_access);
287  if (result == ERROR_SUCCESS) {
288    key_data.reset(new KeyData());
289    if (!key_data->Initialize(key)) {
290      LOG(ERROR) << "Failed to backup key at " << key_path;
291      return false;
292    }
293  } else if (result != ERROR_FILE_NOT_FOUND) {
294    LOG(ERROR) << "Failed to open key at " << key_path
295               << " to create backup, result: " << result;
296    return false;
297  }
298
299  key_data_.swap(key_data);
300  return true;
301}
302
303bool RegistryKeyBackup::WriteTo(HKEY root,
304                                const wchar_t* key_path,
305                                REGSAM wow64_access) const {
306  DCHECK(key_path);
307  DCHECK(wow64_access == 0 ||
308         wow64_access == KEY_WOW64_32KEY ||
309         wow64_access == KEY_WOW64_64KEY);
310
311  bool success = false;
312
313  if (key_data_.get() != NULL) {
314    RegKey dest_key;
315    LONG result = dest_key.Create(root, key_path, KEY_WRITE | wow64_access);
316    if (result != ERROR_SUCCESS) {
317      LOG(ERROR) << "Failed to create destination key at " << key_path
318                 << " to write backup, result: " << result;
319    } else {
320      success = key_data_->WriteTo(&dest_key);
321      LOG_IF(ERROR, !success) << "Failed to write key data.";
322    }
323  } else {
324    success = true;
325  }
326
327  return success;
328}
329