[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>
 
+#include <memory>
+
 namespace {
 
 // Function pointers used for registry access.
 RtlInitUnicodeStringFunction g_rtl_init_unicode_string = nullptr;
 NtCreateKeyFunction g_nt_create_key = nullptr;
 NtDeleteKeyFunction g_nt_delete_key = nullptr;
 NtOpenKeyExFunction g_nt_open_key_ex = nullptr;
 NtCloseFunction g_nt_close = nullptr;
 NtQueryValueKeyFunction g_nt_query_value_key = nullptr;
 NtSetValueKeyFunction g_nt_set_value_key = nullptr;
@@ skipping 515 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);
+  }
+
+  // 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;
 
+  std::unique_ptr<BYTE[]> buffer(new BYTE[size_needed]);
   KEY_VALUE_FULL_INFORMATION* value_info =
-      reinterpret_cast<KEY_VALUE_FULL_INFORMATION*>(new BYTE[size_needed]);
+      reinterpret_cast<KEY_VALUE_FULL_INFORMATION*>(buffer.get());
 
   // 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);
-    success = true;
+  if (!NT_SUCCESS(ntstatus))
+    return false;
+
+  *out_type = value_info->Type;
+  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();
   }
 
-  delete[] value_info;
-  return success;
+  return true;
 }
 
 // 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());
+
   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;
   // 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);
-
+  while (!temp.empty()) {
     pointer += temp.length() + 1;
+    out_multi_sz->push_back(std::move(temp));
     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