[Chrome ELF] New NT registry API.

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
new file mode 100644
index 0000000000000000000000000000000000000000..588886b7e36d6a0dee79c1d68a0ead938d33b68b
--- /dev/null
+++ b/chrome_elf/nt_registry/nt_registry.cc
@@ -0,0 +1,633 @@
+// 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"
+
+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;
+
+// Lazy init.  No concern about concurrency in chrome_elf.
+bool g_initialized = false;
+bool g_system_install = false;
+const size_t g_kRegMaxPathLen = 255;
+wchar_t g_kRegPathHKLM[] = L"\\Registry\\Machine\\";
+wchar_t g_kRegPathHKCU[g_kRegMaxPathLen] = L"";
+std::wstring g_current_user_sid_string;
+std::wstring g_override_path;
+
+// Not using install_util, to prevent circular dependency.
+bool IsThisProcSystem() {
+  wchar_t program_dir[MAX_PATH] = {};
+  wchar_t* cmd_line = GetCommandLineW();
+  // If our command line starts with the "Program Files" or
+  // "Program Files (x86)" path, we're system.
+  DWORD ret = ::GetEnvironmentVariable(L"PROGRAMFILES", program_dir, MAX_PATH);
+  if (ret && ret < MAX_PATH && !::wcsncmp(cmd_line, program_dir, ret))
+    return true;
+
+  ret = ::GetEnvironmentVariable(L"PROGRAMFILES(X86)", program_dir, MAX_PATH);
+  if (ret && ret < MAX_PATH && !::wcsncmp(cmd_line, program_dir, ret))
+    return true;
+
+  return false;
+}
+
+bool InitNativeRegApi() {
+  HMODULE ntdll = ::GetModuleHandleW(L"ntdll.dll");
+
+  // Setup the global function pointers for registry access.
+  g_rtl_init_unicode_string = reinterpret_cast<RtlInitUnicodeStringFunction>(
+      ::GetProcAddress(ntdll, "RtlInitUnicodeString"));
+
+  g_nt_create_key = reinterpret_cast<NtCreateKeyFunction>(
+      ::GetProcAddress(ntdll, "NtCreateKey"));
+
+  g_nt_delete_key = reinterpret_cast<NtDeleteKeyFunction>(
+      ::GetProcAddress(ntdll, "NtDeleteKey"));
+
+  g_nt_open_key_ex = reinterpret_cast<NtOpenKeyExFunction>(
+      ::GetProcAddress(ntdll, "NtOpenKeyEx"));
+
+  g_nt_close =
+      reinterpret_cast<NtCloseFunction>(::GetProcAddress(ntdll, "NtClose"));
+
+  g_nt_query_value_key = reinterpret_cast<NtQueryValueKeyFunction>(
+      ::GetProcAddress(ntdll, "NtQueryValueKey"));
+
+  g_nt_set_value_key = reinterpret_cast<NtSetValueKeyFunction>(
+      ::GetProcAddress(ntdll, "NtSetValueKey"));
+
+  if (!g_rtl_init_unicode_string || !g_nt_create_key || !g_nt_open_key_ex ||
+      !g_nt_delete_key || !g_nt_close || !g_nt_query_value_key ||
+      !g_nt_set_value_key)
+    return false;
+
+  // We need to set HKCU based on the sid of the current user account.
+  RtlFormatCurrentUserKeyPathFunction rtl_current_user_string =
+      reinterpret_cast<RtlFormatCurrentUserKeyPathFunction>(
+          ::GetProcAddress(ntdll, "RtlFormatCurrentUserKeyPath"));
+
+  RtlFreeUnicodeStringFunction rtl_free_unicode_str =
+      reinterpret_cast<RtlFreeUnicodeStringFunction>(
+          ::GetProcAddress(ntdll, "RtlFreeUnicodeString"));
+
+  if (!rtl_current_user_string || !rtl_free_unicode_str)
+    return false;
+
+  UNICODE_STRING current_user_reg_path;
+  if (!NT_SUCCESS(rtl_current_user_string(&current_user_reg_path)))
+    return false;
+
+  // Finish setting up global HKCU path.
+  ::wcsncat(g_kRegPathHKCU, current_user_reg_path.Buffer, g_kRegMaxPathLen - 1);
+  ::wcsncat(g_kRegPathHKCU, L"\\",
+            (g_kRegMaxPathLen - ::wcslen(g_kRegPathHKCU) - 1));
+  // Keep the sid string as well.
+  wchar_t* ptr = ::wcsrchr(current_user_reg_path.Buffer, L'\\');
+  ptr++;
+  g_current_user_sid_string.assign(ptr);
+  rtl_free_unicode_str(&current_user_reg_path);
+
+  // Figure out if we're a system or user install.
+  g_system_install = IsThisProcSystem();
+
+  g_initialized = true;
+  return true;
+}
+
+const wchar_t* ConvertRootKey(nt::ROOT_KEY root) {
+  nt::ROOT_KEY key = root;
+
+  if (!root) {
+    // AUTO
+    key = g_system_install ? nt::HKLM : nt::HKCU;
+  }
+
+  if ((key == nt::HKCU) && (!nt::HKCU_override.empty())) {
+    g_override_path.assign(g_kRegPathHKCU);
+    g_override_path.append(nt::HKCU_override.c_str());
+    g_override_path.append(L"\\");
+    return g_override_path.c_str();
+  } else if ((key == nt::HKLM) && (!nt::HKLM_override.empty())) {
+    g_override_path.assign(g_kRegPathHKCU);
+    g_override_path.append(nt::HKLM_override.c_str());
+    g_override_path.append(L"\\");
+    return g_override_path.c_str();
+  }
+
+  if (key == nt::HKCU)
+    return g_kRegPathHKCU;
+  else
+    return g_kRegPathHKLM;
+}
+
+NTSTATUS CreateKeyWrapper(const std::wstring& key_path,
+                          ACCESS_MASK access,
+                          HANDLE* out_handle,
+                          ULONG* create_or_open OPTIONAL) {

[robertshield, 2016/06/23 02:29:36]

I'm haven't seen the OPTIONAL annotation used before in Chromium code. Do you
need it here?

[penny, 2016/06/25 21:13:43]

So, I started a little discussion with my fellow chrome-windows folks.  I think
these tags make function definitions nicely readable when there is a truly
optional argument.  This annotation has not been used in chromium, but it's also
not mentioned or denied in chromium or google coding standards.  It's definitely
a Windows annotation.

As my use of this is contained in the new nt_registry library (also used a
couple times in CreateRegKey and OpenRegKey), and this library/code is very much
Windows-specific, I'm feeling ok with leaving these annotations in this isolated
library.

It's really easy to remove in the future if anyone gets emotional about it.

+  UNICODE_STRING key_path_uni = {};
+  g_rtl_init_unicode_string(&key_path_uni, key_path.c_str());
+
+  OBJECT_ATTRIBUTES obj = {};
+  InitializeObjectAttributes(&obj, &key_path_uni, OBJ_CASE_INSENSITIVE, NULL,
+                             nullptr);
+
+  return g_nt_create_key(out_handle, access, &obj, 0, nullptr,
+                         REG_OPTION_NON_VOLATILE, create_or_open);
+}
+
+// |root_path| should already exist.
+bool CreateRegKeyRecursive(const std::wstring& root_path,
+                           const std::wstring& sub_key_path,
+                           ACCESS_MASK access,
+                           HANDLE* out_handle) {

[robertshield, 2016/06/23 02:29:36]

fwiw, you can do this more simply (and more efficiently) without recursion by
looping over |sub_key_path| keeping iterators to the beginning and end of each
sub key name. This works too though so only change it if you feel like it :-)

[penny, 2016/06/25 21:13:43]

Thanks Robert.  I was thinking it would be easier to support "unwinding" changes
on failure with recursion.  However, I've changed it now to use an iterator and
still manage to handle cleanup on failure.  A little more complexity, but no
function stack frames at least.  Let me know what you think of this version.

+  std::wstring new_sub_key_path;
+  std::wstring new_root_path(root_path);
+  std::wstring next_key(sub_key_path);
+
+  // See if this is the last sub-key.
+  size_t index = sub_key_path.find_first_of(L"\\", 0);

[robertshield, 2016/06/23 02:29:36]

you're searching for one character, so you can use the find() overload that
takes a wchar_t. 

[penny, 2016/06/25 21:13:43]

Acknowledged.

+  if (index != std::wstring::npos) {
+    // Save the remaining key path.
+    new_sub_key_path = sub_key_path.substr(index + 1, std::wstring::npos);
+    // Adjust next_key to just be the first sub key.
+    next_key.resize(index);
+  }
+  new_root_path.push_back(L'\\');
+  new_root_path.append(next_key);
+
+  // Create the next key.  It might already exist.
+  HANDLE temp_handle = INVALID_HANDLE_VALUE;
+  ULONG create_or_open = 0;
+  NTSTATUS status =
+      CreateKeyWrapper(new_root_path, access, &temp_handle, &create_or_open);
+  if (!NT_SUCCESS(status)) {
+    *out_handle = INVALID_HANDLE_VALUE;
+    return false;
+  }
+
+  // See if done.
+  if (new_sub_key_path.empty()) {
+    // Don't close temp_handle in this case.  Pass it out in |out_handle|.
+    *out_handle = temp_handle;
+    return true;
+  }
+
+  // Recursively call this function on |new_sub_key_path|.
+  if (!CreateRegKeyRecursive(new_root_path, new_sub_key_path, access,
+                             out_handle)) {
+    // On failure, "clean up" newly created subkeys only.  If it already
+    // existed, leave it.
+    if (create_or_open == REG_CREATED_NEW_KEY)
+      g_nt_delete_key(temp_handle);
+    g_nt_close(temp_handle);
+    return false;
+  }
+
+  g_nt_close(temp_handle);
+  return true;
+}
+
+}  // namespace
+
+namespace nt {
+
+std::wstring HKLM_override;
+std::wstring HKCU_override;
+
+//------------------------------------------------------------------------------
+// Create, open, delete, close functions
+//------------------------------------------------------------------------------
+
+bool CreateRegKey(ROOT_KEY root,
+                  const wchar_t* key_path,
+                  ACCESS_MASK access,
+                  HANDLE* out_handle OPTIONAL) {
+  if (!g_initialized)
+    InitNativeRegApi();
+
+  // Open the root first.  Using create instead of open, because
+  // there is small chance a full redirection root does not exist yet.
+  std::wstring root_path(ConvertRootKey(root));
+  HANDLE root_handle = INVALID_HANDLE_VALUE;
+  NTSTATUS status = CreateKeyWrapper(root_path, access, &root_handle, nullptr);
+  if (!NT_SUCCESS(status))
+    return false;
+
+  // Make sure |key_path| does not start or end with a path seperator.
+  std::wstring sub_path(key_path);
+  if (sub_path.front() == L'\\')
+    sub_path.erase(0, 1);
+  if (sub_path.back() == L'\\')
+    sub_path.pop_back();
+
+  // Recursively create the rest of the sub keys.
+  HANDLE key_handle = INVALID_HANDLE_VALUE;
+  bool success =
+      CreateRegKeyRecursive(root_path, sub_path, access, &key_handle);
+  CloseRegKey(root_handle);
+
+  if (success) {
+    if (out_handle)
+      *out_handle = key_handle;
+    else
+      CloseRegKey(key_handle);
+    return true;
+  }
+
+  return false;
+}
+
+bool OpenRegKey(ROOT_KEY root,
+                const wchar_t* key_path,
+                ACCESS_MASK access,
+                HANDLE* out_handle,
+                NTSTATUS* error_code OPTIONAL) {
+  if (!g_initialized)
+    InitNativeRegApi();
+
+  NTSTATUS status = STATUS_UNSUCCESSFUL;
+  UNICODE_STRING key_path_uni = {};
+  OBJECT_ATTRIBUTES obj = {};
+  *out_handle = INVALID_HANDLE_VALUE;
+
+  std::wstring full_path(ConvertRootKey(root));
+  full_path.append(key_path);
+
+  g_rtl_init_unicode_string(&key_path_uni, full_path.c_str());
+  InitializeObjectAttributes(&obj, &key_path_uni, OBJ_CASE_INSENSITIVE, NULL,
+                             NULL);
+
+  status = g_nt_open_key_ex(out_handle, access, &obj, 0);
+  // See if caller wants the NTSTATUS.
+  if (error_code)
+    *error_code = status;
+
+  if (NT_SUCCESS(status))
+    return true;
+
+  return false;
+}
+
+bool DeleteRegKey(HANDLE key) {
+  if (!g_initialized)
+    InitNativeRegApi();
+
+  NTSTATUS status = STATUS_UNSUCCESSFUL;
+
+  status = g_nt_delete_key(key);
+
+  if (NT_SUCCESS(status))
+    return true;
+
+  return false;
+}
+
+// wrapper function
+bool DeleteRegKey(ROOT_KEY root, const wchar_t* key_path) {
+  HANDLE key = INVALID_HANDLE_VALUE;
+
+  if (!OpenRegKey(root, key_path, DELETE, &key, nullptr))
+    return false;
+
+  if (!DeleteRegKey(key)) {
+    CloseRegKey(key);
+    return false;
+  }
+
+  CloseRegKey(key);
+  return true;
+}
+
+void CloseRegKey(HANDLE key) {
+  if (!g_initialized)
+    InitNativeRegApi();
+  g_nt_close(key);
+}
+
+//------------------------------------------------------------------------------
+// Getter functions
+//------------------------------------------------------------------------------
+
+bool QueryRegKeyValue(HANDLE key,
+                      const wchar_t* value_name,
+                      ULONG* out_type,
+                      BYTE** out_buffer,
+                      DWORD* out_size) {
+  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);
+    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;
+
+  if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) ||
+      type != REG_DWORD)
+    return false;
+
+  *out_dword = *(reinterpret_cast<DWORD*>(value_bytes));
+
+  delete[] value_bytes;
+  return true;
+}
+
+// wrapper function
+bool QueryRegValueDWORD(ROOT_KEY root,
+                        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 | KEY_WOW64_32KEY, &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;
+  ULONG type = REG_NONE;
+
+  if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) ||
+      type != REG_SZ)
+    return false;
+
+  *out_sz = reinterpret_cast<wchar_t*>(value_bytes);
+
+  delete[] value_bytes;
+  return true;
+}
+
+// wrapper function
+bool QueryRegValueSZ(ROOT_KEY root,
+                     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 | KEY_WOW64_32KEY, &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;
+  ULONG type = REG_NONE;
+
+  if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) ||
+      type != REG_MULTI_SZ)
+    return false;
+
+  // Make sure the vector is empty to start.
+  (*out_multi_sz).resize(0);
+
+  wchar_t* pointer = reinterpret_cast<wchar_t*>(value_bytes);
+  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;
+    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,
+                          const wchar_t* key_path,
+                          const wchar_t* value_name,
+                          std::vector<std::wstring>* out_multi_sz) {
+  HANDLE key = INVALID_HANDLE_VALUE;
+
+  if (!OpenRegKey(root, key_path, KEY_QUERY_VALUE | KEY_WOW64_32KEY, &key,
+                  NULL))
+    return false;
+
+  if (!QueryRegValueMULTISZ(key, value_name, out_multi_sz)) {
+    CloseRegKey(key);
+    return false;
+  }
+
+  CloseRegKey(key);
+  return true;
+}
+
+//------------------------------------------------------------------------------
+// Setter functions
+//------------------------------------------------------------------------------
+
+bool SetRegKeyValue(HANDLE key,
+                    const wchar_t* value_name,
+                    ULONG type,
+                    const BYTE* data,
+                    DWORD data_size) {
+  if (!g_initialized)
+    InitNativeRegApi();
+
+  NTSTATUS ntstatus = STATUS_UNSUCCESSFUL;
+  UNICODE_STRING value_uni = {};
+  g_rtl_init_unicode_string(&value_uni, value_name);
+
+  BYTE* non_const_data = const_cast<BYTE*>(data);
+  ntstatus =
+      g_nt_set_value_key(key, &value_uni, 0, type, non_const_data, data_size);
+
+  if (NT_SUCCESS(ntstatus))
+    return true;
+
+  return false;
+}
+
+// wrapper function
+bool SetRegValueDWORD(HANDLE key, const wchar_t* value_name, DWORD value) {
+  return SetRegKeyValue(key, value_name, REG_DWORD,
+                        reinterpret_cast<BYTE*>(&value), sizeof(value));
+}
+
+// wrapper function
+bool SetRegValueDWORD(ROOT_KEY root,
+                      const wchar_t* key_path,
+                      const wchar_t* value_name,
+                      DWORD value) {
+  HANDLE key = INVALID_HANDLE_VALUE;
+
+  if (!OpenRegKey(root, key_path, KEY_SET_VALUE | KEY_WOW64_32KEY, &key, NULL))
+    return false;
+
+  if (!SetRegValueDWORD(key, value_name, value)) {
+    CloseRegKey(key);
+    return false;
+  }
+
+  return true;
+}
+
+// wrapper function
+bool SetRegValueSZ(HANDLE key,
+                   const wchar_t* value_name,
+                   const std::wstring& value) {
+  // Make sure the number of bytes in |value|, including EoS, fits in a DWORD.
+  if (std::numeric_limits<DWORD>::max() <
+      ((value.length() + 1) * sizeof(wchar_t)))
+    return false;
+
+  DWORD size = (static_cast<DWORD>((value.length() + 1) * sizeof(wchar_t)));
+  return SetRegKeyValue(key, value_name, REG_SZ,
+                        reinterpret_cast<const BYTE*>(value.c_str()), size);
+}
+
+// wrapper function
+bool SetRegValueSZ(ROOT_KEY root,
+                   const wchar_t* key_path,
+                   const wchar_t* value_name,
+                   const std::wstring& value) {
+  HANDLE key = INVALID_HANDLE_VALUE;
+
+  if (!OpenRegKey(root, key_path, KEY_SET_VALUE | KEY_WOW64_32KEY, &key, NULL))
+    return false;
+
+  if (!SetRegValueSZ(key, value_name, value)) {
+    CloseRegKey(key);
+    return false;
+  }
+
+  return true;
+}
+
+// wrapper function
+bool SetRegValueMULTISZ(HANDLE key,
+                        const wchar_t* value_name,
+                        const std::vector<std::wstring>& values) {
+  std::vector<wchar_t> builder;
+
+  for (auto& string : values) {
+    // Just in case someone is passing in an illegal empty string
+    // (not allowed in REG_MULTI_SZ), ignore it.
+    if (!string.empty()) {
+      for (const wchar_t& w : string) {
+        builder.push_back(w);
+      }
+      builder.push_back(L'\0');
+    }
+  }
+  // Add second null terminator to end REG_MULTI_SZ.
+  builder.push_back(L'\0');
+  // Handle rare case where the vector passed in was empty,
+  // or only had an empty string.
+  if (builder.size() == 1)
+    builder.push_back(L'\0');
+
+  if (std::numeric_limits<DWORD>::max() < builder.size())
+    return false;
+
+  return SetRegKeyValue(
+      key, value_name, REG_MULTI_SZ, reinterpret_cast<BYTE*>(builder.data()),
+      (static_cast<DWORD>(builder.size()) + 1) * sizeof(wchar_t));
+}
+
+// wrapper function
+bool SetRegValueMULTISZ(ROOT_KEY root,
+                        const wchar_t* key_path,
+                        const wchar_t* value_name,
+                        const std::vector<std::wstring>& values) {
+  HANDLE key = INVALID_HANDLE_VALUE;
+
+  if (!OpenRegKey(root, key_path, KEY_SET_VALUE | KEY_WOW64_32KEY, &key, NULL))
+    return false;
+
+  if (!SetRegValueMULTISZ(key, value_name, values)) {
+    CloseRegKey(key);
+    return false;
+  }
+
+  return true;
+}
+
+//------------------------------------------------------------------------------
+// Utils
+//------------------------------------------------------------------------------
+
+std::wstring GetCurrentUserSidString() {
+  if (!g_initialized)
+    InitNativeRegApi();
+
+  return g_current_user_sid_string;
+}
+
+};  // namespace nt