Added CopyRegKeyWorkItem in support of IE low rights policy fixes.

chrome/installer/util/registry_key_backup.cc

Index: chrome/installer/util/registry_key_backup.cc
diff --git a/chrome/installer/util/delete_reg_key_work_item.cc b/chrome/installer/util/registry_key_backup.cc
similarity index 57%
copy from chrome/installer/util/delete_reg_key_work_item.cc
copy to chrome/installer/util/registry_key_backup.cc
index b044f447c47d015beebd4d00aeaa29549bfd0d59..93734b32501395f85e56ce162b4f56d98817afff 100644
--- a/chrome/installer/util/delete_reg_key_work_item.cc
+++ b/chrome/installer/util/registry_key_backup.cc
@@ -1,97 +1,97 @@
-// Copyright (c) 2010 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-#include "chrome/installer/util/delete_reg_key_work_item.h"
+#include "chrome/installer/util/registry_key_backup.h"
 
-#include <shlwapi.h>
 #include <algorithm>
 #include <limits>
 #include <vector>
 
 #include "base/logging.h"
-#include "base/rand_util.h"
-#include "base/stringprintf.h"
 #include "base/win/registry.h"
-#include "chrome/installer/util/logging_installer.h"
 
 using base::win::RegKey;
 
 namespace {
 const REGSAM kKeyReadNoNotify = (KEY_READ) & ~(KEY_NOTIFY);
-}
+}  // namespace
 
-// A container for a registry key, its values, and its subkeys.  We don't use
-// more obvious methods for various reasons:
-// - RegCopyTree isn't supported pre-Vista, so we'd have to do something
-//   different for XP anyway.
-// - SHCopyKey can't copy subkeys into a volatile destination, so we'd have to
-//   worry about polluting the registry.
-// We don't persist security attributes since we only delete keys that we own,
-// and we don't set custom attributes on them anyway.
-class DeleteRegKeyWorkItem::RegKeyBackup {
+// A container for a registry key, its values, and its subkeys.
+class RegistryKeyBackup::KeyData {
  public:
-  RegKeyBackup();
+  KeyData();
+  ~KeyData();
   bool Initialize(const RegKey& key);
   bool WriteTo(RegKey* key) const;
 
  private:
-  // A container for a registry value.
-  class RegValueBackup {
-   public:
-    RegValueBackup();
-    void Initialize(const wchar_t* name_buffer, DWORD name_size,
-                    DWORD type, const uint8* data, DWORD data_size);
-    const std::wstring& name_str() const { return name_; }
-    const wchar_t* name() const { return name_.empty() ? NULL : name_.c_str(); }
-    DWORD type() const { return type_; }
-    const uint8* data() const { return data_.empty() ? NULL : &data_[0]; }
-    DWORD data_len() const { return static_cast<DWORD>(data_.size()); }
-
-   private:
-    std::wstring name_;
-    std::vector<uint8> data_;
-    DWORD type_;
-
-    DISALLOW_COPY_AND_ASSIGN(RegValueBackup);
-  };
-
-  scoped_array<RegValueBackup> values_;
+  class ValueData;
+
+  scoped_array<ValueData> values_;
   scoped_array<std::wstring> subkey_names_;
-  scoped_array<RegKeyBackup> subkeys_;
-  ptrdiff_t num_values_;
-  ptrdiff_t num_subkeys_;
+  scoped_array<KeyData> subkeys_;
+  DWORD num_values_;
+  DWORD num_subkeys_;
+
+  DISALLOW_COPY_AND_ASSIGN(KeyData);
+};
+
+// A container for a registry value.
+class RegistryKeyBackup::KeyData::ValueData {
+ public:
+  ValueData();
+  ~ValueData();
+  void Initialize(const wchar_t* name_buffer, DWORD name_size,
+                  DWORD type, const uint8* data, DWORD data_size);
+  const std::wstring& name_str() const { return name_; }
+  const wchar_t* name() const { return name_.empty() ? NULL : name_.c_str(); }
+  DWORD type() const { return type_; }
+  const uint8* data() const { return data_.empty() ? NULL : &data_[0]; }
+  DWORD data_len() const { return static_cast<DWORD>(data_.size()); }
+
+ private:
+  std::wstring name_;
+  std::vector<uint8> data_;
+  DWORD type_;
 
-  DISALLOW_COPY_AND_ASSIGN(RegKeyBackup);
+  DISALLOW_COPY_AND_ASSIGN(ValueData);
 };
 
-DeleteRegKeyWorkItem::RegKeyBackup::RegValueBackup::RegValueBackup()
+RegistryKeyBackup::KeyData::ValueData::ValueData()
     : type_(REG_NONE) {
 }
 
-void DeleteRegKeyWorkItem::RegKeyBackup::RegValueBackup::Initialize(
+RegistryKeyBackup::KeyData::ValueData::~ValueData()
+{
+}
+
+void RegistryKeyBackup::KeyData::ValueData::Initialize(
     const wchar_t* name_buffer,
     DWORD name_size,
-    DWORD type, const uint8* data,
+    DWORD type,
+    const uint8* data,
     DWORD data_size) {
   name_.assign(name_buffer, name_size);
   type_ = type;
   data_.assign(data, data + data_size);
 }
 
-DeleteRegKeyWorkItem::RegKeyBackup::RegKeyBackup()
+RegistryKeyBackup::KeyData::KeyData()
     : num_values_(0),
       num_subkeys_(0) {
 }
 
+RegistryKeyBackup::KeyData::~KeyData()
+{
+}
+
 // Initializes this object by reading the values and subkeys of |key|.
 // Security descriptors are not backed up.
-bool DeleteRegKeyWorkItem::RegKeyBackup::Initialize(const RegKey& key) {
-  DCHECK(key.Valid());
-
-  scoped_array<RegValueBackup> values;
+bool RegistryKeyBackup::KeyData::Initialize(const RegKey& key) {
+  scoped_array<ValueData> values;
   scoped_array<std::wstring> subkey_names;
-  scoped_array<RegKeyBackup> subkeys;
+  scoped_array<KeyData> subkeys;
 
   DWORD num_subkeys = 0;
   DWORD max_subkey_name_len = 0;
@@ -117,7 +117,7 @@ bool DeleteRegKeyWorkItem::RegKeyBackup::Initialize(const RegKey& key) {
 
   // Backup the values.
   if (num_values != 0) {
-    values.reset(new RegValueBackup[num_values]);
+    values.reset(new ValueData[num_values]);
     scoped_array<uint8> value_buffer(new uint8[max_value_len]);
     DWORD name_size = 0;
     DWORD value_type = REG_NONE;
@@ -140,14 +140,16 @@ bool DeleteRegKeyWorkItem::RegKeyBackup::Initialize(const RegKey& key) {
         case ERROR_MORE_DATA:
           if (value_size > max_value_len) {
             max_value_len = value_size;
+            value_buffer.reset();  // Release to heap before new allocation.
             value_buffer.reset(new uint8[max_value_len]);
           } else {
-            DCHECK(max_name_len - 1 < name_size);
+            DCHECK_LT(max_name_len - 1, name_size);
             if (name_size >= std::numeric_limits<DWORD>::max() - 1) {
               LOG(ERROR) << "Failed backing up key; value name out of range.";
               return false;
             }
             max_name_len = name_size + 1;
+            name_buffer.reset();  // Release to heap before new allocation.
             name_buffer.reset(new wchar_t[max_name_len]);
           }
           break;
@@ -166,7 +168,7 @@ bool DeleteRegKeyWorkItem::RegKeyBackup::Initialize(const RegKey& key) {
   // Backup the subkeys.
   if (num_subkeys != 0) {
     subkey_names.reset(new std::wstring[num_subkeys]);
-    subkeys.reset(new RegKeyBackup[num_subkeys]);
+    subkeys.reset(new KeyData[num_subkeys]);
     DWORD name_size = 0;
 
     // Get the names of them.
@@ -228,12 +230,14 @@ bool DeleteRegKeyWorkItem::RegKeyBackup::Initialize(const RegKey& key) {
 }
 
 // Writes the values and subkeys of this object into |key|.
-bool DeleteRegKeyWorkItem::RegKeyBackup::WriteTo(RegKey* key) const {
+bool RegistryKeyBackup::KeyData::WriteTo(RegKey* key) const {
+  DCHECK(key);
+
   LONG result = ERROR_SUCCESS;
 
   // Write the values.
-  for (int i = 0; i < num_values_; ++i) {
-    const RegValueBackup& value = values_[i];
+  for (DWORD i = 0; i < num_values_; ++i) {
+    const ValueData& value = values_[i];
     result = RegSetValueEx(key->Handle(), value.name(), 0, value.type(),
                            value.data(), value.data_len());
     if (result != ERROR_SUCCESS) {
@@ -245,7 +249,7 @@ bool DeleteRegKeyWorkItem::RegKeyBackup::WriteTo(RegKey* key) const {
 
   // Write the subkeys.
   RegKey subkey;
-  for (int i = 0; i < num_subkeys_; ++i) {
+  for (DWORD i = 0; i < num_subkeys_; ++i) {
     const std::wstring& name = subkey_names_[i];
 
     result = subkey.Create(key->Handle(), name.c_str(), KEY_WRITE);
@@ -264,75 +268,54 @@ bool DeleteRegKeyWorkItem::RegKeyBackup::WriteTo(RegKey* key) const {
   return true;
 }
 
-DeleteRegKeyWorkItem::~DeleteRegKeyWorkItem() {
+RegistryKeyBackup::RegistryKeyBackup() {
 }
 
-DeleteRegKeyWorkItem::DeleteRegKeyWorkItem(HKEY predefined_root,
-                                           const std::wstring& path)
-    : predefined_root_(predefined_root),
-      path_(path) {
-  // It's a safe bet that we don't want to delete one of the root trees.
-  DCHECK(!path.empty());
+RegistryKeyBackup::~RegistryKeyBackup() {
 }
 
-bool DeleteRegKeyWorkItem::Do() {
-  scoped_ptr<RegKeyBackup> backup;
-
-  // Only try to make a backup if we're not configured to ignore failures.
-  if (!ignore_failure_) {
-    RegKey original_key;
-
-    // Does the key exist?
-    LONG result = original_key.Open(predefined_root_, path_.c_str(),
-                                    kKeyReadNoNotify);
-    if (result == ERROR_SUCCESS) {
-      backup.reset(new RegKeyBackup());
-      if (!backup->Initialize(original_key)) {
-        LOG(ERROR) << "Failed to backup key at " << path_;
-        return ignore_failure_;
-      }
-    } else if (result != ERROR_FILE_NOT_FOUND) {
-      LOG(ERROR) << "Failed to open key at " << path_
-                 << " to create backup, result: " << result;
-      return ignore_failure_;
-    }
-  }
+bool RegistryKeyBackup::Initialize(HKEY root, const wchar_t* key_path) {
+  DCHECK(key_path);
 
-  // Delete the key.
-  LONG result = SHDeleteKey(predefined_root_, path_.c_str());
-  if (result != ERROR_SUCCESS && result != ERROR_FILE_NOT_FOUND) {
-    LOG(ERROR) << "Failed to delete key at " << path_ << ", result: "
-               << result;
-    return ignore_failure_;
-  }
+  RegKey key;
+  scoped_ptr<KeyData> key_data;
 
-  // We've succeeded, so remember any backup we may have made.
-  backup_.swap(backup);
+  // Does the key exist?
+  LONG result = key.Open(root, key_path, kKeyReadNoNotify);
+  if (result == ERROR_SUCCESS) {
+    key_data.reset(new KeyData());
+    if (!key_data->Initialize(key)) {
+      LOG(ERROR) << "Failed to backup key at " << key_path;
+      return false;
+    }
+  } else if (result != ERROR_FILE_NOT_FOUND) {
+    LOG(ERROR) << "Failed to open key at " << key_path
+               << " to create backup, result: " << result;
+    return false;
+  }
 
+  key_data_.swap(key_data);
   return true;
 }
 
-void DeleteRegKeyWorkItem::Rollback() {
-  if (ignore_failure_ || backup_.get() == NULL)
-    return;
+bool RegistryKeyBackup::WriteTo(HKEY root, const wchar_t* key_path) const {
+  DCHECK(key_path);
 
-  // Delete anything in the key before restoring the backup in case someone else
-  // put new data in the key after Do().
-  LONG result = SHDeleteKey(predefined_root_, path_.c_str());
-  if (result != ERROR_SUCCESS && result != ERROR_FILE_NOT_FOUND) {
-    LOG(ERROR) << "Failed to delete key at " << path_ << " in rollback, "
-                  "result: " << result;
-  }
+  bool success = false;
 
-  // Restore the old contents.  The restoration takes on its default security
-  // attributes; any custom attributes are lost.
-  RegKey original_key;
-  result = original_key.Create(predefined_root_, path_.c_str(), KEY_WRITE);
-  if (result != ERROR_SUCCESS) {
-    LOG(ERROR) << "Failed to create original key at " << path_
-               << " in rollback, result: " << result;
+  if (key_data_.get() != NULL) {
+    RegKey dest_key;
+    LONG result = dest_key.Create(root, key_path, KEY_WRITE);
+    if (result != ERROR_SUCCESS) {
+      LOG(ERROR) << "Failed to create destination key at " << key_path
+                 << " to write backup, result: " << result;
+    } else {
+      success = key_data_->WriteTo(&dest_key);
+      LOG_IF(ERROR, !success) << "Failed to write key data.";
+    }
   } else {
-    if (!backup_->WriteTo(&original_key))
-      LOG(ERROR) << "Failed to restore key in rollback, result: " << result;
+    success = true;
   }
+
+  return success;
 }