[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..a826a660ef68b8056aab3c0db927122bff3e802e 100644
--- a/chrome_elf/nt_registry/nt_registry.cc
+++ b/chrome_elf/nt_registry/nt_registry.cc
@@ -542,6 +542,45 @@ 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.
+  if (value_bytes->size() < terminator_size) {
+    value_bytes->insert(value_bytes->end(),
+                        (terminator_size - value_bytes->size()), 0);

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

nit: i would omit these extra parens around the arithmetic -- they don't
disambiguate anything imo.

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

Done.

+  }
+
+  // Sanity check content size based on character size.
+  DWORD modulo = value_bytes->size() % sizeof(wchar_t);
+  value_bytes->insert(value_bytes->end(), modulo, 0);
+
+  // 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.
+  value_bytes->insert(value_bytes->end(), terminator_size, 0);
+
+  return;
+}
+
 //------------------------------------------------------------------------------
 // Misc wrapper functions - LOCAL
 //------------------------------------------------------------------------------
@@ -749,8 +788,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 +812,15 @@ 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;
+
+    if (data_size) {
+      // Move the data into |out_buffer| vector.
+      BYTE* data = reinterpret_cast<BYTE*>(value_info) + value_info->DataOffset;
+      out_buffer->assign(data, data + data_size);
+    } else {
+      out_buffer->clear();
+    }
     success = true;
   }
 
@@ -791,16 +833,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 +872,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)) {

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

maybe document in .h that REG_EXPAND_SZ values are handled by this, but that no
expansion takes place?

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

The .h function documentation declares support for REG_EXPAND_SZ.
Registry APIs don't do expansion of this type.  It's just a notification label
if an application wants to label a reg value as needing to be expanded later. 
Registry APIs treat SZ and EXPAND_SZ the same.

     return false;
+  }
+
+  EnsureTerminatedSZ(&value_bytes, false);
 
-  *out_sz = reinterpret_cast<wchar_t*>(value_bytes);
+  *out_sz = reinterpret_cast<wchar_t*>(value_bytes.data());
 
-  delete[] value_bytes;
   return true;
 }
 
@@ -866,33 +911,29 @@ 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();

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

  out_multi_sz->clear();

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

Done.

+
+  wchar_t* pointer = reinterpret_cast<wchar_t*>(value_bytes.data());
   std::wstring temp = pointer;
   // Loop.  Each string is separated by '\0'.  Another '\0' at very end (so 2 in
   // a row).
   while (temp.length() != 0) {
-    (*out_multi_sz).push_back(temp);
-
     pointer += temp.length() + 1;
+    (*out_multi_sz).push_back(std::move(temp));

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

  out_multi_sz->push_back(std::move(temp));

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

Done.

     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;
 }