[NtRegistry] Ensure REG_SZ and REG_MULTI_SZ are null terminated.

chrome_elf/nt_registry/nt_registry.cc

 // Copyright 2016 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_elf/nt_registry/nt_registry.h"
 
 #include <assert.h>
 #include <stdlib.h>
 
 namespace {
@@ skipping 524 matching lines @@
       out_base->push_back(L'\\');
     }
     subkeys->erase(subkeys->begin(), subkeys->begin() + num_base_tokens);
   } else {
     out_base->assign(converted_root);
   }
 
   return true;
 }
 
+// String safety.
+// - NOTE: only working with wchar_t here.
+// - Also ensures the content of |value_bytes| is at least a terminator.
+// - Pass "true" for |multi| for MULTISZ.
+void EnsureTerminatedSZ(std::vector<BYTE>* value_bytes, bool multi) {
+  DWORD terminator_size = sizeof(wchar_t);
+
+  if (multi)
+    terminator_size = 2 * sizeof(wchar_t);
+
+  // Ensure content is at least the size of a terminator.
+  for (size_t i = value_bytes->size(); i < terminator_size; i++)

[grt (UTC plus 2), 2017/07/20 05:22:53]

  if (value_bytes->size() < terminator_size) {
    value_bytes->insert(value_bytes->end(), terminator_size -
value_bytes->size(), 0);
  }

[penny, 2017/07/20 18:38:37]

Done.  <sigh>  Thank you for your 1337 c++ skillz Greg.  Forgot all about
insert() with an iterator.  That's much cleaner - all through this CL.

+    value_bytes->push_back(0);
+
+  // Sanity check content size based on character size.
+  DWORD modulo = value_bytes->size() % sizeof(wchar_t);
+  while (modulo) {

[grt (UTC plus 2), 2017/07/20 05:22:53]

  value_bytes->insert(value_bytes->end(), modulo, 0);

[penny, 2017/07/20 18:38:37]

Done.

+    value_bytes->push_back(0);
+    modulo--;
+  }
+
+  // Now finally check for trailing terminator.
+  bool terminated = true;
+  size_t last_element = value_bytes->size() - 1;
+  for (size_t i = 0; i < terminator_size; i++) {
+    if ((*value_bytes)[last_element - i] != 0) {
+      terminated = false;
+      break;
+    }
+  }
+
+  if (terminated)
+    return;
+
+  // Append a full terminator to be safe.
+  for (size_t i = 0; i < terminator_size; i++)

[grt (UTC plus 2), 2017/07/20 05:22:53]

  value_bytes->insert(value_bytes->end(), terminator_size, 0);

[penny, 2017/07/20 18:38:37]

Done.

+    value_bytes->push_back(0);
+
+  return;
+}
+
 //------------------------------------------------------------------------------
 // Misc wrapper functions - LOCAL
 //------------------------------------------------------------------------------
 
 NTSTATUS CreateKeyWrapper(const std::wstring& key_path,
                           ACCESS_MASK access,
                           HANDLE* out_handle,
                           ULONG* create_or_open OPTIONAL) {
   UNICODE_STRING key_path_uni = {};
   g_rtl_init_unicode_string(&key_path_uni, key_path.c_str());
@@ skipping 187 matching lines @@
   g_nt_close(key);
 }
 
 //------------------------------------------------------------------------------
 // Getter functions
 //------------------------------------------------------------------------------
 
 bool QueryRegKeyValue(HANDLE key,
                       const wchar_t* value_name,
                       ULONG* out_type,
-                      BYTE** out_buffer,
-                      DWORD* out_size) {
+                      std::vector<BYTE>* out_buffer) {
   if (!g_initialized)
     InitNativeRegApi();
 
   NTSTATUS ntstatus = STATUS_UNSUCCESSFUL;
   UNICODE_STRING value_uni = {};
   g_rtl_init_unicode_string(&value_uni, value_name);
   DWORD size_needed = 0;
   bool success = false;
 
   // First call to find out how much room we need for the value!
   ntstatus = g_nt_query_value_key(key, &value_uni, KeyValueFullInformation,
                                   nullptr, 0, &size_needed);
   if (ntstatus != STATUS_BUFFER_TOO_SMALL)
     return false;
 
   KEY_VALUE_FULL_INFORMATION* value_info =
       reinterpret_cast<KEY_VALUE_FULL_INFORMATION*>(new BYTE[size_needed]);
 
   // Second call to get the value.
   ntstatus = g_nt_query_value_key(key, &value_uni, KeyValueFullInformation,
                                   value_info, size_needed, &size_needed);
   if (NT_SUCCESS(ntstatus)) {
     *out_type = value_info->Type;
-    *out_size = value_info->DataLength;
-    *out_buffer = new BYTE[*out_size];
-    ::memcpy(*out_buffer,
-             (reinterpret_cast<BYTE*>(value_info) + value_info->DataOffset),
-             *out_size);
+    DWORD data_size = value_info->DataLength;
+    out_buffer->clear();

[grt (UTC plus 2), 2017/07/20 05:22:53]

nit: by using assign() below, clear() is only needed in the !data_size case.

[penny, 2017/07/20 18:38:36]

Done.

+    if (data_size) {
+      // Move the data into |out_buffer| vector.
+      BYTE* data = reinterpret_cast<BYTE*>(value_info) + value_info->DataOffset;
+      for (size_t i = 0; i < data_size; i++)

[grt (UTC plus 2), 2017/07/20 05:22:53]

    out_buffer->assign(data, data + data_size);

[penny, 2017/07/20 18:38:37]

assign.  thank you.  

+        out_buffer->push_back(data[i]);
+    }
     success = true;
   }
 
   delete[] value_info;
   return success;
 }
 
 // wrapper function
 bool QueryRegValueDWORD(HANDLE key,
                         const wchar_t* value_name,
                         DWORD* out_dword) {
   ULONG type = REG_NONE;
-  BYTE* value_bytes = nullptr;
-  DWORD ret_size = 0;
+  std::vector<BYTE> value_bytes;
 
-  if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) ||
-      type != REG_DWORD)
+  if (!QueryRegKeyValue(key, value_name, &type, &value_bytes) ||
+      type != REG_DWORD) {
+    return false;
+  }
+
+  if (value_bytes.size() < sizeof(*out_dword))
     return false;
 
-  *out_dword = *(reinterpret_cast<DWORD*>(value_bytes));
+  *out_dword = *(reinterpret_cast<DWORD*>(value_bytes.data()));
 
-  delete[] value_bytes;
   return true;
 }
 
 // wrapper function
 bool QueryRegValueDWORD(ROOT_KEY root,
                         WOW64_OVERRIDE wow64_override,
                         const wchar_t* key_path,
                         const wchar_t* value_name,
                         DWORD* out_dword) {
   HANDLE key = INVALID_HANDLE_VALUE;
 
   if (!OpenRegKey(root, key_path, KEY_QUERY_VALUE | wow64_override, &key, NULL))
     return false;
 
   if (!QueryRegValueDWORD(key, value_name, out_dword)) {
     CloseRegKey(key);
     return false;
   }
 
   CloseRegKey(key);
   return true;
 }
 
 // wrapper function
 bool QueryRegValueSZ(HANDLE key,
                      const wchar_t* value_name,
                      std::wstring* out_sz) {
-  BYTE* value_bytes = nullptr;
-  DWORD ret_size = 0;
+  std::vector<BYTE> value_bytes;
   ULONG type = REG_NONE;
 
-  if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) ||
-      type != REG_SZ)
+  if (!QueryRegKeyValue(key, value_name, &type, &value_bytes) ||
+      (type != REG_SZ && type != REG_EXPAND_SZ)) {
     return false;
+  }
 
-  *out_sz = reinterpret_cast<wchar_t*>(value_bytes);
+  EnsureTerminatedSZ(&value_bytes, false);
 
-  delete[] value_bytes;
+  *out_sz = reinterpret_cast<wchar_t*>(value_bytes.data());

[grt (UTC plus 2), 2017/07/20 05:22:53]

should this preserve embedded nulls? if yes, use out_sz->assign(value, value +
len_less_terminator).

[penny, 2017/07/20 18:38:37]

I never expected the result of the QueryRegValueSZ wrapper function to preserve
embedded nulls.  That's a very abnormal use case for an SZ type.  That's what
multisz is really for.  If there's a special case, caller can use the raw
QueryRegKeyValue function.

[grt (UTC plus 2), 2017/07/21 08:09:00]

Sounds reasonable to me. Please document this behavior in the doc comment for
both QueryRegValueSZ functions in the .h.

[penny, 2017/07/25 18:30:47]

Done.

+
   return true;
 }
 
 // wrapper function
 bool QueryRegValueSZ(ROOT_KEY root,
                      WOW64_OVERRIDE wow64_override,
                      const wchar_t* key_path,
                      const wchar_t* value_name,
                      std::wstring* out_sz) {
   HANDLE key = INVALID_HANDLE_VALUE;
 
   if (!OpenRegKey(root, key_path, KEY_QUERY_VALUE | wow64_override, &key, NULL))
     return false;
 
   if (!QueryRegValueSZ(key, value_name, out_sz)) {
     CloseRegKey(key);
     return false;
   }
 
   CloseRegKey(key);
   return true;
 }
 
 // wrapper function
 bool QueryRegValueMULTISZ(HANDLE key,
                           const wchar_t* value_name,
                           std::vector<std::wstring>* out_multi_sz) {
-  BYTE* value_bytes = nullptr;
-  DWORD ret_size = 0;
+  std::vector<BYTE> value_bytes;
   ULONG type = REG_NONE;
 
-  if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) ||
-      type != REG_MULTI_SZ)
+  if (!QueryRegKeyValue(key, value_name, &type, &value_bytes) ||
+      type != REG_MULTI_SZ) {
     return false;
+  }
 
-  // Make sure the vector is empty to start.
-  (*out_multi_sz).resize(0);
+  EnsureTerminatedSZ(&value_bytes, true);
 
-  wchar_t* pointer = reinterpret_cast<wchar_t*>(value_bytes);
+  // Make sure the out vector is empty to start.
+  (*out_multi_sz).clear();
+
+  wchar_t* pointer = reinterpret_cast<wchar_t*>(value_bytes.data());
   std::wstring temp = pointer;

[robertshield, 2017/07/20 17:14:04]

Won't this truncate any characters after the first null char in value_bytes?

[penny, 2017/07/20 18:38:37]

No.  It merely creates a new wstring object, and copies in the string starting
at pointer - up until the first \0.  std::wstring initialization doesn't mangle
the source.

The C pointer |pointer| merely moves along the content of |value_bytes|, to the
start of each string (which are separated by \0).  The loop instantiates |temp|
each time starting from |pointer|.

   // Loop.  Each string is separated by '\0'.  Another '\0' at very end (so 2 in
   // a row).
   while (temp.length() != 0) {

[grt (UTC plus 2), 2017/07/20 05:22:53]

nit: !temp.empty()

[penny, 2017/07/20 18:38:37]

I wasn't entirely sure whether an empty string ("" or \0) would be considered
empty().  Am a bit confused about how std::wstring handles EoS characters. 
Wanted to be safe with length().
e.g.
std::wstring blah = L"";
vs
std::wstring blah2();

Internally, is |blah| an "empty" object, or is that size() == 1 and length() ==
0?  Assuming this class even stores the null character as an element.

Based on your advice, I at least assume empty() returns true for a string of
length 0, size 1?  <sigh>  Gimme a C array any day of the week.

[grt (UTC plus 2), 2017/07/21 08:09:00]

The spec more-or-less says that a std::basic_string does not contain a
terminator -- there is no guarantee that the memory returned by data() is
terminated, only that the first length() characters in it are the string
contents.

So then what about its c_str() method? It returns a pointer to null-terminated
char array that contains the length() characters of data() followed by a string
terminator. It just so happens that all sane implementations put a terminator at
the end of the data() buffer so that c_str() and data() return the same pointer
value, but that's really an implementation detail.

Back to length() == 0 vs empty() -- the latter basically boils down to this
inline method:

  bool empty() const { return length() == 0; }

In Chromium, we prefer to use empty() when testing for this since it's more
explicit.

In your examples, blah and blah2 are both empty and have length() == 0.

Oh, and as for size() and length(), they're synonyms! I imagine the two exist
because:
- length() -- folks are used to thinking of how "long" a string is thanks to
strlen
- size() -- std::basic_string is a container (like a vector), so give it a
size() member for use in other templates that take a container type

Hope this helps clear up some of the ambiguity!

[penny, 2017/07/25 18:30:47]

Done.  Clear as mud!

     (*out_multi_sz).push_back(temp);

[grt (UTC plus 2), 2017/07/20 05:22:53]

nit: avoid making another copy of |temp| with:
    pointer += temp.length() + 1;
    (*out_multi_sz).push_back(std::move(temp));
    temp = pointer;

[penny, 2017/07/20 18:38:36]

Done.  I didn't even think push_back could take the result of a std::move - but
TIL.  Makes sense.  Thanks!

[grt (UTC plus 2), 2017/07/21 08:09:00]

I didn't know either -- I just figured something like that existed and looked it
up here: http://en.cppreference.com/w/cpp/container/vector/push_back

 
     pointer += temp.length() + 1;
     temp = pointer;
   }
 
-  // Handle the case of "empty multi_sz".
-  if (out_multi_sz->size() == 0)
-    out_multi_sz->push_back(L"");
-
-  delete[] value_bytes;
   return true;
 }
 
 // wrapper function
 bool QueryRegValueMULTISZ(ROOT_KEY root,
                           WOW64_OVERRIDE wow64_override,
                           const wchar_t* key_path,
                           const wchar_t* value_name,
                           std::vector<std::wstring>* out_multi_sz) {
   HANDLE key = INVALID_HANDLE_VALUE;
@@ skipping 183 matching lines @@
   if (!g_initialized)
     InitNativeRegApi();
 
   if (root == HKCU || (root == AUTO && !g_system_install))
     return g_HKCU_override;
 
   return g_HKLM_override;
 }
 
 };  // namespace nt