Add rollback support to DeleteRegKeyWorkItem....

chrome/installer/util/delete_reg_key_work_item.cc

 // Copyright (c) 2010 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 <shlwapi.h>
-
-#include "chrome/installer/util/delete_reg_key_work_item.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;
+
+// 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.  Also, we've observed crashes in other
+//   shell lightweight API funcs, so perhaps it's best to avoid them.

[robertshield, 2011/03/01 16:04:43]

I feel we have insufficient evidence for shlwapi problems, please remove that
part of the comment.

[grt (UTC plus 2), 2011/03/01 20:54:02]

Done.

+// 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 {
+ public:
+  RegKeyBackup();
+  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_;
+  scoped_array<std::wstring> subkey_names_;
+  scoped_array<RegKeyBackup> subkeys_;
+  ptrdiff_t num_values_;
+  ptrdiff_t num_subkeys_;
+
+  DISALLOW_COPY_AND_ASSIGN(RegKeyBackup);
+};
+
+DeleteRegKeyWorkItem::RegKeyBackup::RegValueBackup::RegValueBackup()
+    : type_(REG_NONE) {
+}
+
+void DeleteRegKeyWorkItem::RegKeyBackup::RegValueBackup::Initialize(
+    const wchar_t* name_buffer,
+    DWORD name_size,
+    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()
+    : num_values_(0),
+      num_subkeys_(0) {
+}
+
+// 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;
+  scoped_array<std::wstring> subkey_names;
+  scoped_array<RegKeyBackup> subkeys;
+
+  DWORD num_subkeys = 0;
+  DWORD max_subkey_name_len = 0;
+  DWORD num_values = 0;
+  DWORD max_value_name_len = 0;
+  DWORD max_value_len = 0;
+  LONG result = RegQueryInfoKey(key.Handle(), NULL, NULL, NULL,
+                                &num_subkeys, &max_subkey_name_len, NULL,
+                                &num_values, &max_value_name_len,
+                                &max_value_len, NULL, NULL);
+  if (result != ERROR_SUCCESS) {
+    LOG(ERROR) << "Failed getting info of key to backup, result: " << result;
+    return false;
+  }
+  if (max_subkey_name_len >= std::numeric_limits<DWORD>::max() - 1 ||
+      max_value_name_len >= std::numeric_limits<DWORD>::max() - 1) {
+    LOG(ERROR)
+        << "Failed backing up key; subkeys and/or names are out of range.";
+    return false;
+  }
+  DWORD max_name_len = std::max(max_subkey_name_len, max_value_name_len) + 1;
+  scoped_array<wchar_t> name_buffer(new wchar_t[max_name_len]);
+
+  // Backup the values.
+  if (num_values != 0) {
+    values.reset(new RegValueBackup[num_values]);
+    scoped_array<uint8> value_buffer(new uint8[max_value_len]);
+    DWORD name_size = 0;
+    DWORD value_type = REG_NONE;
+    DWORD value_size = 0;
+
+    for (DWORD i = 0; i < num_values; ++i) {
+      name_size = max_name_len;
+      value_size = max_value_len;
+      result = RegEnumValue(key.Handle(), i, name_buffer.get(), &name_size,
+                            NULL, &value_type, value_buffer.get(), &value_size);
+      if (result != ERROR_SUCCESS) {

[robertshield, 2011/03/01 16:04:43]

do you want to handle ERROR_MORE_DATA here (which would happen if an existing
value was lengthened after teh call to RegQueryInfoKey)?

It feels like we should try to still work in this case (we return false from the
WorkItem's Do() if we return false here). 

[grt (UTC plus 2), 2011/03/01 20:54:02]

Done.  Also handles the case where values and/or keys are removed between
RegQueryInfoKey and RegEnum*.

+        LOG(ERROR) << "Failed backing up value " << i << ", result: " << result;
+        return false;
+      }
+      values[i].Initialize(name_buffer.get(), name_size, value_type,
+                           value_buffer.get(), value_size);
+    }
+    DLOG_IF(WARNING, RegEnumValue(key.Handle(), num_values, NULL, NULL, NULL,
+                                  &value_type, NULL, NULL) == ERROR_SUCCESS)
+        << "Concurrent modifications to registry key during backup operation.";
+  }
+
+  // Backup the subkeys.
+  if (num_subkeys != 0) {
+    subkey_names.reset(new std::wstring[num_subkeys]);
+    subkeys.reset(new RegKeyBackup[num_subkeys]);
+    DWORD name_size = 0;
+
+    // Get the names of them.
+    for (DWORD i = 0; i < num_subkeys; ++i) {
+      name_size = max_name_len;
+      result = RegEnumKeyEx(key.Handle(), i, name_buffer.get(), &name_size,
+                            NULL, NULL, NULL, NULL);

[robertshield, 2011/03/01 16:04:43]

ditto

[grt (UTC plus 2), 2011/03/01 20:54:02]

Done.

+      if (result != ERROR_SUCCESS) {
+        LOG(ERROR) << "Failed getting name of up subkey " << i
+                   << " for backup, result: " << result;
+        return false;
+      }
+      subkey_names[i].assign(name_buffer.get(), name_size);
+    }
+    DLOG_IF(WARNING, RegEnumKeyEx(key.Handle(), num_subkeys, NULL, &name_size,
+                                  NULL, NULL, NULL, NULL) == ERROR_SUCCESS)
+        << "Concurrent modifications to registry key during backup operation.";
+
+    // Get their values.
+    RegKey subkey;
+    for (DWORD i = 0; i < num_subkeys; ++i) {
+      result = subkey.Open(key.Handle(), subkey_names[i].c_str(), KEY_READ);
+      if (result != ERROR_SUCCESS) {
+        LOG(ERROR) << "Failed opening subkey \"" << subkey_names[i]
+                   << "\" for backup, result: " << result;
+        return false;
+      }
+      if (!subkeys[i].Initialize(subkey)) {
+        LOG(ERROR) << "Failed backing up subkey \"" << subkey_names[i] << "\"";
+        return false;
+      }
+    }
+  }
+
+  values_.swap(values);
+  subkey_names_.swap(subkey_names);
+  subkeys_.swap(subkeys);
+  num_values_ = num_values;
+  num_subkeys_ = num_subkeys;
+
+  return true;
+}
+
+// Writes the values and subkeys of this object into |key|.
+bool DeleteRegKeyWorkItem::RegKeyBackup::WriteTo(RegKey* key) const {
+  LONG result = ERROR_SUCCESS;
+
+  // Write the values.
+  for (int i = 0; i < num_values_; ++i) {
+    const RegValueBackup& value = values_[i];
+    result = RegSetValueEx(key->Handle(), value.name(), 0, value.type(),
+                           value.data(), value.data_len());
+    if (result != ERROR_SUCCESS) {
+      LOG(ERROR) << "Failed writing value \"" << value.name_str()
+                 << "\", result: " << result;
+      return false;
+    }
+  }
+
+  // Write the subkeys.
+  RegKey subkey;
+  for (int i = 0; i < num_subkeys_; ++i) {
+    const std::wstring& name = subkey_names_[i];
+
+    result = subkey.Create(key->Handle(), name.c_str(), KEY_WRITE);
+    if (result != ERROR_SUCCESS) {
+      LOG(ERROR) << "Failed creating subkey \"" << name << "\", result: "
+                 << result;
+      return false;
+    }
+    if (!subkeys_[i].WriteTo(&subkey)) {
+      LOG(ERROR) << "Failed writing subkey \"" << name << "\", result: "
+                 << result;
+      return false;
+    }
+  }
+
+  return true;
+}
 
 DeleteRegKeyWorkItem::~DeleteRegKeyWorkItem() {
 }
 
 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());
 }
 
 bool DeleteRegKeyWorkItem::Do() {
-  // TODO(robertshield): Implement rollback for this work item. Consider using
-  // RegSaveKey or RegCopyKey.
+  scoped_ptr<RegKeyBackup>  backup;

[robertshield, 2011/03/01 16:04:43]

nit: 1 space

[grt (UTC plus 2), 2011/03/01 20:54:02]

Nice catch.  Done.

+
+  // Only try to make a backup if we're not configured to ignore failures.
   if (!ignore_failure_) {
-    NOTREACHED();
-    LOG(ERROR) << "Unexpected configuration in DeleteRegKeyWorkItem.";
-    return false;
+    RegKey original_key;
+
+    // Does the key exist?
+    LONG result = original_key.Open(predefined_root_, path_.c_str(), KEY_READ);
+    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_;
+    }
   }
 
-  LSTATUS result = SHDeleteKey(predefined_root_, path_.c_str());
-  if (result != ERROR_SUCCESS) {
-    LOG(ERROR) << "Failed to delete key at " << path_ << ", result: " << result;
+  // 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_;
   }
 
+  // We've succeeded, so rememeber any backup we may have made.
+  backup_.swap(backup);
+
   return true;
 }
 
 void DeleteRegKeyWorkItem::Rollback() {
-  if (ignore_failure_)
+  if (ignore_failure_ || backup_.get() == NULL)
     return;
-  NOTREACHED();
+
+  // 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;
+  }
+
+  // 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;
+  } else {
+    if (!backup_->WriteTo(&original_key))
+      LOG(ERROR) << "Failed to restore key in rollback, result: " << result;
+  }
 }
-