[NtRegistry] Ensure REG_SZ and REG_MULTI_SZ are null terminated.

chrome_elf/nt_registry/nt_registry.cc

Index: chrome_elf/nt_registry/nt_registry.cc
diff --git a/chrome_elf/nt_registry/nt_registry.cc b/chrome_elf/nt_registry/nt_registry.cc
index 4c65af8daff7c76ea5a9e4a0255da0225aeb58a4..6e1039c707cbc53144ce0fd6d67042659e54300e 100644
--- a/chrome_elf/nt_registry/nt_registry.cc
+++ b/chrome_elf/nt_registry/nt_registry.cc
@@ -542,6 +542,47 @@ bool ParseFullRegPath(const std::wstring& 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
 //------------------------------------------------------------------------------
@@ -749,8 +790,7 @@ void CloseRegKey(HANDLE key) {
 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();
 
@@ -774,11 +814,14 @@ bool QueryRegKeyValue(HANDLE key,
                                   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;
   }
 
@@ -791,16 +834,18 @@ 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) ||
+      type != REG_DWORD) {
+    return false;
+  }
 
-  if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) ||
-      type != REG_DWORD)
+  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;
 }
 
@@ -828,17 +873,18 @@ bool QueryRegValueDWORD(ROOT_KEY root,
 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;
+  }
+
+  EnsureTerminatedSZ(&value_bytes, false);
 
-  *out_sz = reinterpret_cast<wchar_t*>(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.

 
-  delete[] value_bytes;
   return true;
 }
 
@@ -866,18 +912,20 @@ bool QueryRegValueSZ(ROOT_KEY root,
 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).
@@ -888,11 +936,6 @@ bool QueryRegValueMULTISZ(HANDLE key,
     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;
 }