Introduce a SetRegValueWorkItem overload that accepts a callback instead (...)

chrome/installer/util/set_reg_value_work_item.cc

Index: chrome/installer/util/set_reg_value_work_item.cc
diff --git a/chrome/installer/util/set_reg_value_work_item.cc b/chrome/installer/util/set_reg_value_work_item.cc
index 596fe7c8895ecd81944ddf262f16f4386b63a92c..515a79f4fa4614d9b7579c792ce7b6ad80e146ef 100644
--- a/chrome/installer/util/set_reg_value_work_item.cc
+++ b/chrome/installer/util/set_reg_value_work_item.cc
@@ -9,6 +9,43 @@
 #include "base/win/registry.h"
 #include "chrome/installer/util/logging_installer.h"
 
+namespace {
+
+// Transforms |str_value| into the byte-by-byte representation of its underlying
+// string, stores the result in |binary_data|.
+void StringToBinaryData(const std::wstring& str_value,
+                        std::vector<uint8>* binary_data) {
+  DCHECK(binary_data);
+  const uint8* data = reinterpret_cast<const uint8*>(str_value.c_str());
+  binary_data->assign(data, data + (str_value.length() + 1) * sizeof(wchar_t));
+}
+
+// Transforms |binary_data| into its wstring representation (assuming
+// |binary_data| is a sequence of wchar_t's).
+void BinaryDataToString(const std::vector<uint8>& binary_data,
+                        std::wstring* str_value) {
+  DCHECK(str_value);
+  if (binary_data.size() < sizeof(wchar_t)) {
+    str_value->clear();
+    return;
+  }
+
+  // The length of the string contained in |binary_data|. This is at least one
+  // per the above condition and may contain a null character if |binary_data|
+  // is null-terminated.
+  const size_t str_len_maybe_with_null = binary_data.size() / sizeof(wchar_t);
+
+  str_value->assign(reinterpret_cast<const wchar_t*>(&binary_data[0]),
+                    str_len_maybe_with_null);
+
+  // If the last character in the assigned buffer was a null-character, remove
+  // it (C++ strings are not NULL terminated).

[grt (UTC plus 2), 2015/06/30 20:26:42]

this comment is a bit misleading: while the terminator is not counted in
std::basic_string<>::size(), std::basic_string<>.c_str() is null-terminated
(StringToBinaryData relies on this). how about just "  // Trim off a trailing
string terminator, if present."

[gab, 2015/07/01 02:55:19]

sgtm, done.

+  if ((*str_value)[str_len_maybe_with_null - 1] == L'\0')

[gab, 2015/06/29 20:11:01]

Think this is safe or should I check the resulting string's length?

[grt (UTC plus 2), 2015/06/30 20:26:42]

I think it's safe. str_len_maybe_with_null >= 1, so at the extreme it erases all
of str_value, right?

[gab, 2015/07/01 02:55:19]

Well I was thinking if somehow this ends up constructing a shorter string than
str_len_maybe_with_null (e.g., somehow this is bigger than allowed max string
size?!) than this can result in a buffer overflow?

[grt (UTC plus 2), 2015/07/01 15:37:16]

Ah, I see. You can do away with str_len_maybe_with_null with something like
this:
  if (!str_value->empty()) {
    auto iter = str_value->end();
    if (*--iter == L'\0')
      str_value->erase(iter);
  }

[gab, 2015/07/02 01:59:05]

I like that, done.

+    str_value->erase(str_len_maybe_with_null - 1);
+}
+
+}  // namespace
+
 SetRegValueWorkItem::~SetRegValueWorkItem() {
 }
 
@@ -29,8 +66,7 @@ SetRegValueWorkItem::SetRegValueWorkItem(HKEY predefined_root,
   DCHECK(wow64_access == 0 ||
          wow64_access == KEY_WOW64_32KEY ||
          wow64_access == KEY_WOW64_64KEY);
-  const uint8* data = reinterpret_cast<const uint8*>(value_data.c_str());
-  value_.assign(data, data + (value_data.length() + 1) * sizeof(wchar_t));
+  StringToBinaryData(value_data, &value_);
 }
 
 SetRegValueWorkItem::SetRegValueWorkItem(HKEY predefined_root,
@@ -75,6 +111,27 @@ SetRegValueWorkItem::SetRegValueWorkItem(HKEY predefined_root,
   value_.assign(data, data + sizeof(value_data));
 }
 
+SetRegValueWorkItem::SetRegValueWorkItem(
+    HKEY predefined_root,
+    const std::wstring& key_path,
+    REGSAM wow64_access,
+    const std::wstring& value_name,
+    const GetValueFromExistingCallback& get_value_callback)
+    : predefined_root_(predefined_root),
+      key_path_(key_path),
+      value_name_(value_name),
+      get_value_callback_(get_value_callback),
+      overwrite_(true),
+      wow64_access_(wow64_access),
+      status_(SET_VALUE),
+      type_(REG_SZ),
+      previous_type_(0) {
+  DCHECK(wow64_access == 0 ||
+         wow64_access == KEY_WOW64_32KEY ||
+         wow64_access == KEY_WOW64_64KEY);
+  // Nothing to do, |get_value_callback| will fill |value_| later.
+}
+
 bool SetRegValueWorkItem::Do() {
   LONG result = ERROR_SUCCESS;
   base::win::RegKey key;
@@ -118,6 +175,21 @@ bool SetRegValueWorkItem::Do() {
     }
   }
 
+  if (!get_value_callback_.is_null()) {
+    // Although this could be made more generic, for now this assumes the
+    // |type_| of |value_| is REG_SZ.
+    DCHECK(type_ == REG_SZ);
+
+    // Fill |previous_value_str| with the wstring representation of the binary
+    // data in |previous_value_| as long as it's of type REG_SZ (leave it empty
+    // otherwise).
+    std::wstring previous_value_str;
+    if (previous_type_ == REG_SZ)
+      BinaryDataToString(previous_value_, &previous_value_str);
+
+    StringToBinaryData(get_value_callback_.Run(previous_value_str), &value_);
+  }
+
   result = key.WriteValue(value_name_.c_str(), &value_[0],
                           static_cast<DWORD>(value_.size()), type_);
   if (result != ERROR_SUCCESS) {