Home | History | Annotate | Download | only in util 
      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,
    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 
    303 bool 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