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

chrome/installer/util/set_reg_value_work_item.cc

 // 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/set_reg_value_work_item.h"
 
 #include "base/logging.h"
 #include "base/strings/string_util.h"
 #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() {
 }
 
 SetRegValueWorkItem::SetRegValueWorkItem(HKEY predefined_root,
                                          const std::wstring& key_path,
                                          REGSAM wow64_access,
                                          const std::wstring& value_name,
                                          const std::wstring& value_data,
                                          bool overwrite)
     : predefined_root_(predefined_root),
       key_path_(key_path),
       value_name_(value_name),
       overwrite_(overwrite),
       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);
-  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,
                                          const std::wstring& key_path,
                                          REGSAM wow64_access,
                                          const std::wstring& value_name,
                                          DWORD value_data,
                                          bool overwrite)
     : predefined_root_(predefined_root),
       key_path_(key_path),
@@ skipping 24 matching lines @@
       status_(SET_VALUE),
       type_(REG_QWORD),
       previous_type_(0) {
   DCHECK(wow64_access == 0 ||
          wow64_access == KEY_WOW64_32KEY ||
          wow64_access == KEY_WOW64_64KEY);
   const uint8* data = reinterpret_cast<const uint8*>(&value_data);
   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;
   if (status_ != SET_VALUE) {
     // we already did something.
     VLOG(1) << "multiple calls to Do()";
     result = ERROR_CANTWRITE;
     return ignore_failure_;
   }
 
@@ skipping 23 matching lines @@
       previous_value_.resize(size);
       result = key.ReadValue(value_name_.c_str(), &previous_value_[0], &size,
                              &previous_type_);
       if (result != ERROR_SUCCESS) {
         previous_value_.clear();
         VLOG(1) << "Failed to save original value. Error: " << result;
       }
     }
   }
 
+  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) {
     VLOG(1) << "Failed to write value " << key_path_ << " error: " << result;
     return ignore_failure_;
   }
 
   status_ = previous_type_ ? VALUE_OVERWRITTEN : NEW_VALUE_CREATED;
   return true;
 }
@@ skipping 28 matching lines @@
     result = key.WriteValue(value_name_.c_str(), previous_value,
                             static_cast<DWORD>(previous_value_.size()),
                             previous_type_);
     VLOG(1) << "rollback: restoring " << value_name_ << " error: " << result;
   } else {
     NOTREACHED();
   }
 
   status_ = VALUE_ROLL_BACK;
 }