[Chrome ELF] New NT registry API.

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"
+
+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;

[robertshield, 2016/06/26 02:11:04]

I believe that static non-POD types are disallowed by the style guide due to
non-deterministic creation / destruction ordering:
https://google.github.io/styleguide/cppguide.html#Static_and_Global_Variables

If possible, please make these wchar_t arrays as above.

[penny, 2016/07/11 19:59:04]

Thanks for that.  I was unaware of this restriction!  I figured using an object
would be more space efficient since I don't know how much room will be needed. 
I've made them POD arrays anyhow.

+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) {
+  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);
+}
+
+}  // namespace
+
+namespace nt {
+
+std::wstring HKLM_override;

[penny, 2016/07/11 19:59:04]

Robert, I adjusted these to be POD arrays as well.

+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 with a path seperator.
+  std::wstring sub_path(key_path);
+  if (sub_path.front() == L'\\')
+    sub_path.erase(0, 1);
+  // Make sure |key_path| DOES end with a path seperator.
+  if (sub_path.back() != L'\\')
+    sub_path.push_back(L'\\');
+
+  std::wstring current_path(root_path);
+  std::vector<HANDLE> rollback;
+  bool success = true;
+
+  // Recursively create the rest of the sub keys.
+  for (std::wstring::iterator it = sub_path.begin(); it != sub_path.end();
+       ++it) {
+    if (*it == L'\\') {
+      // Process the latest subkey.
+      ULONG created = 0;
+      HANDLE key_handle = INVALID_HANDLE_VALUE;
+      status =
+          CreateKeyWrapper(current_path.c_str(), access, &key_handle, &created);
+      if (!NT_SUCCESS(status)) {
+        success = false;
+        break;
+      }
+      // Save any handle of subkey created, in case of rollback.

[robertshield, 2016/06/26 02:11:04]

micro-nit: s/handle of subkey/subkey handle/

[penny, 2016/07/11 19:59:04]

Done.

+      if (created == REG_CREATED_NEW_KEY)
+        rollback.push_back(&key_handle);
+      else
+        g_nt_close(key_handle);
+      current_path.push_back(L'\\');
+    } else {
+      current_path.push_back(*it);
+    }
+  }
+
+  if (!success) {
+    // Delete any subkeys created.
+    for (auto handle : rollback) {
+      g_nt_delete_key(handle);

[robertshield, 2016/06/26 02:11:03]

This doesn't delete |root_handle| in case of failure, is that intended?

[penny, 2016/07/11 19:59:04]

I've adjusted this a bit.  Now I only open the base hive first. 
"\Registry\Machine" or "\Registry\User\SID".  This will always exist already.

+    }
+  }
+
+  for (auto handle : rollback) {
+    g_nt_close(handle);
+  }
+  g_nt_close(root_handle);

[robertshield, 2016/06/26 02:11:04]

If we do want to delete |root_handle| above in case of failure, we could
probably get rid of it as a separate variable and make it the first element in
the |rollback| list, which would simplify some of the clean up code.

[penny, 2016/07/11 19:59:04]

Acknowledged.

+
+  if (!success)
+    return false;
+
+  // See if caller wants the handle open.
+  if (out_handle) {

[robertshield, 2016/06/26 02:11:04]

slight optimization: could you check |out_handle| before closing the handles
above, pluck the last one off and close the rest, e.g.:


if (!success) {
  // Delete any subkeys created.
  for (auto handle : rollback) {
    g_nt_delete_key(handle);
  }
} else if (out_handle && !rollback.empty()) {
  *out_handle = rollback.back();
  rollback.pop_back();
}

for (auto handle : rollback) {
  g_nt_close(handle);
}
g_nt_close(root_handle);

return success;

[penny, 2016/07/11 19:59:04]

Unfortunately, not every handle is on the vector - only the ones that had to be
created.

But, I've changed this now, and know when I'm dealing with the last handle.  So
no need to "reopen" the final handle if the user wants it (which is most often
the case).

+    // Reopen the final handle.
+    status =
+        CreateKeyWrapper(current_path.c_str(), access, out_handle, nullptr);
+    if (!NT_SUCCESS(status))
+      return false;
+  }
+
+  return true;
+}
+
+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