src/chrome/installer/util/registry_key_backup.cc

   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
  15 using base::win::RegKey;
  16
  17 namespace {
  18
  19 const REGSAM kKeyReadNoNotify = (KEY_READ) & ~(KEY_NOTIFY);
  20
  21 // A container for a registry value.
  22 class 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.
  63 class 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
  89 ValueData::ValueData() : type_(REG_NONE) {
  90 }
  91
  92 ValueData::~ValueData() {
  93 }
  94
  95 void 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
 106 RegistryKeyBackup::KeyData::KeyData() {
 107 }
 108
 109 RegistryKeyBackup::KeyData::~KeyData() {
 110 }
 111
 112 bool 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
 228 bool 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
 268 RegistryKeyBackup::RegistryKeyBackup() {
 269 }
 270
 271 RegistryKeyBackup::~RegistryKeyBackup() {
 272 }
 273
 274 bool RegistryKeyBackup::Initialize(HKEY root, const wchar_t* key_path) {
 275   DCHECK(key_path);
 276
 277   RegKey key;
 278   scoped_ptr<KeyData> key_data;
 279
 280   // Does the key exist?
 281   LONG result = key.Open(root, key_path, kKeyReadNoNotify);
 282   if (result == ERROR_SUCCESS) {
 283     key_data.reset(new KeyData());
 284     if (!key_data->Initialize(key)) {
 285       LOG(ERROR) << "Failed to backup key at " << key_path;
 286       return false;
 287     }
 288   } else if (result != ERROR_FILE_NOT_FOUND) {
 289     LOG(ERROR) << "Failed to open key at " << key_path
 290                << " to create backup, result: " << result;
 291     return false;
 292   }
 293
 294   key_data_.swap(key_data);
 295   return true;
 296 }
 297
 298 bool RegistryKeyBackup::WriteTo(HKEY root, const wchar_t* key_path) const {
 299   DCHECK(key_path);
 300
 301   bool success = false;
 302
 303   if (key_data_.get() != NULL) {
 304     RegKey dest_key;
 305     LONG result = dest_key.Create(root, key_path, KEY_WRITE);
 306     if (result != ERROR_SUCCESS) {
 307       LOG(ERROR) << "Failed to create destination key at " << key_path
 308                  << " to write backup, result: " << result;
 309     } else {
 310       success = key_data_->WriteTo(&dest_key);
 311       LOG_IF(ERROR, !success) << "Failed to write key data.";
 312     }
 313   } else {
 314     success = true;
 315   }
 316
 317   return success;
 318 }