[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.
+  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
 //------------------------------------------------------------------------------
 
 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]);

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

please use std::unique_ptr where possible to avoid the possibility of memory
leaks. it would be cleaner for this to be:

  std::unique_ptr<BYTE[]> buffer(new BYTE[size_needed]);
  KEY_VALUE_FULL_INFORMATION* value_info =
      reinterpret_cast<KEY_VALUE_FULL_INFORMATION*>(buffer.get());

and then get rid of delete[] on line 827. this also gets rid of the need for the
|success| local var, since early-exit becomes safe. my personal choice would be
to make the failure case return false so that the success case reaches the end
of the function.

(note the avoidance of base::MakeUnique for allocating the array, as it would
zero-initialize it.)

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

Done.

 
   // 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;
+
+    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;
   }
 
   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)) {

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

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