[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;
+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;
+}
+
+}  // namespace
+
+namespace nt {
+
+std::wstring HKLM_override;
+std::wstring HKCU_override;
+
+bool CreateRegKey(ROOT_KEY root,
+                  const wchar_t* key_path,
+                  ACCESS_MASK access,
+                  HANDLE* out_handle OPTIONAL) {
+  if (!g_initialized)
+    InitNativeRegApi();
+
+  NTSTATUS status = STATUS_UNSUCCESSFUL;
+  UNICODE_STRING key_path_uni = {};
+  OBJECT_ATTRIBUTES obj = {};
+  HANDLE key_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);
+  // Assuming CreateOptions to preserve across reboot, and currently
+  // don't care whether already existed or not.
+  status = g_nt_create_key(&key_handle, access, &obj, 0, nullptr,
+                           REG_OPTION_NON_VOLATILE, nullptr);
+
+  if (NT_SUCCESS(status)) {
+    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) {
+  wchar_t* string = const_cast<wchar_t*>(value.c_str());
+  if (std::numeric_limits<DWORD>::max() < ::wcslen(string))

[robertshield, 2016/06/17 13:40:03]

could use value.length() instead of all this

[penny, 2016/06/22 23:13:14]

Done.

+    return false;
+
+  DWORD length = static_cast<DWORD>(::wcslen(string));
+
+  return SetRegKeyValue(key, value_name, REG_SZ,
+                        reinterpret_cast<BYTE*>(string),
+                        (length + 1) * sizeof(wchar_t));
+}
+
+// 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