[Chrome ELF] New NT registry API.

chrome_elf/chrome_elf_reg.cc

Index: chrome_elf/chrome_elf_reg.cc
diff --git a/chrome_elf/chrome_elf_reg.cc b/chrome_elf/chrome_elf_reg.cc
new file mode 100644
index 0000000000000000000000000000000000000000..6a5b7ae93490dfb3ba7f76a492725c40d9221381
--- /dev/null
+++ b/chrome_elf/chrome_elf_reg.cc
@@ -0,0 +1,539 @@
+// 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/chrome_elf_reg.h"
+
+#include "chrome_elf/chrome_elf_util.h"  // IsSystemInstall()
+
+namespace {
+
+// function pointers used for registry access.
+RtlInitUnicodeStringFunction rtl_init_unicode_string = nullptr;
+NtCreateKeyFunction nt_create_key = nullptr;
+NtDeleteKeyFunction nt_delete_key = nullptr;
+NtOpenKeyExFunction nt_open_key_ex = nullptr;
+NtCloseFunction nt_close = nullptr;
+NtQueryValueKeyFunction nt_query_value_key = nullptr;
+NtSetValueKeyFunction nt_set_value_key = nullptr;
+
+// lazy init.  No concern about concurrency in chrome_elf.
+bool initialized = false;
+bool system_install = false;
+const size_t kRegMaxPathLen = 255;
+wchar_t kRegPathHKLM[] = L"\\Registry\\Machine\\";
+wchar_t kRegPathHKCU[kRegMaxPathLen] = L"";
+base::string16 current_user_sid_string;
+base::string16 override_path;

[robertshield, 2016/04/20 16:55:53]

These are globals, please label them as g_override_path, etc. to make it
clearer.

[penny, 2016/05/28 01:34:23]

Done.

+
+bool InitNativeRegApi() {
+  HMODULE ntdll = ::GetModuleHandleW(L"ntdll.dll");
+
+  // Setup the global function pointers for registry access.
+  rtl_init_unicode_string = reinterpret_cast<RtlInitUnicodeStringFunction>(
+      ::GetProcAddress(ntdll, "RtlInitUnicodeString"));
+
+  nt_create_key = reinterpret_cast<NtCreateKeyFunction>(
+      ::GetProcAddress(ntdll, "NtCreateKey"));
+
+  nt_delete_key = reinterpret_cast<NtDeleteKeyFunction>(
+      ::GetProcAddress(ntdll, "NtDeleteKey"));
+
+  nt_open_key_ex = reinterpret_cast<NtOpenKeyExFunction>(
+      ::GetProcAddress(ntdll, "NtOpenKeyEx"));
+
+  nt_close =
+      reinterpret_cast<NtCloseFunction>(::GetProcAddress(ntdll, "NtClose"));
+
+  nt_query_value_key = reinterpret_cast<NtQueryValueKeyFunction>(
+      ::GetProcAddress(ntdll, "NtQueryValueKey"));
+
+  nt_set_value_key = reinterpret_cast<NtSetValueKeyFunction>(
+      ::GetProcAddress(ntdll, "NtSetValueKey"));
+
+  if (!rtl_init_unicode_string || !nt_create_key || !nt_open_key_ex ||
+      !nt_delete_key || !nt_close || !nt_query_value_key || !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(kRegPathHKCU, current_user_reg_path.Buffer, kRegMaxPathLen - 1);
+  ::wcsncat(kRegPathHKCU, L"\\", (kRegMaxPathLen - ::wcslen(kRegPathHKCU) - 1));
+  // Keep the sid string as well.
+  wchar_t* ptr = ::wcsrchr(current_user_reg_path.Buffer, L'\\');
+  ptr++;
+  current_user_sid_string.assign(ptr);
+  rtl_free_unicode_str(&current_user_reg_path);
+
+  // Figure out if we're a system or user install.
+  system_install = IsSystemInstall(GetCommandLineW());
+
+  initialized = true;
+  return true;
+}
+
+const wchar_t* ConvertRootKey(nt::ROOT_KEY root) {
+  nt::ROOT_KEY key = root;
+
+  if (!root) {
+    // AUTO
+    key = system_install ? nt::HKLM : nt::HKCU;
+  }
+
+  if ((key == nt::HKCU) && (!nt::HKCU_override.empty())) {
+    override_path.assign(kRegPathHKCU);
+    override_path.append(nt::HKCU_override.c_str());
+    override_path.append(L"\\");
+    return override_path.c_str();
+  } else if ((key == nt::HKLM) && (!nt::HKLM_override.empty())) {
+    override_path.assign(kRegPathHKCU);
+    override_path.append(nt::HKLM_override.c_str());
+    override_path.append(L"\\");
+    return override_path.c_str();
+  }
+
+  if (key == nt::HKCU)
+    return kRegPathHKCU;
+  else
+    return kRegPathHKLM;
+}
+
+}  // namespace
+
+namespace nt {
+
+base::string16 HKLM_override;
+base::string16 HKCU_override;
+
+bool CreateRegKey(ROOT_KEY root,
+                  const wchar_t* key_path,
+                  ACCESS_MASK access,
+                  HANDLE* out_handle OPTIONAL) {
+  if (!initialized)
+    InitNativeRegApi();
+
+  NTSTATUS status = STATUS_UNSUCCESSFUL;
+  UNICODE_STRING key_path_uni = {};
+  OBJECT_ATTRIBUTES obj = {};
+  HANDLE key_handle = INVALID_HANDLE_VALUE;
+
+  base::string16 full_path(ConvertRootKey(root));
+  full_path.append(key_path);
+
+  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 = 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 DeleteRegKey(HANDLE key) {
+  if (!initialized)
+    InitNativeRegApi();
+
+  NTSTATUS status = STATUS_UNSUCCESSFUL;
+
+  status = 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;
+}
+
+bool OpenRegKey(ROOT_KEY root,
+                const wchar_t* key_path,
+                ACCESS_MASK access,
+                HANDLE* out_handle,
+                NTSTATUS* error_code OPTIONAL) {
+  if (!initialized)
+    InitNativeRegApi();
+
+  NTSTATUS status = STATUS_UNSUCCESSFUL;
+  UNICODE_STRING key_path_uni = {};
+  OBJECT_ATTRIBUTES obj = {};
+  *out_handle = INVALID_HANDLE_VALUE;
+
+  base::string16 full_path(ConvertRootKey(root));
+  full_path.append(key_path);
+
+  rtl_init_unicode_string(&key_path_uni, full_path.c_str());
+  InitializeObjectAttributes(&obj, &key_path_uni, OBJ_CASE_INSENSITIVE, NULL,
+                             NULL);
+
+  status = 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;
+}
+
+void CloseRegKey(HANDLE key) {
+  if (!initialized)
+    InitNativeRegApi();
+  nt_close(key);
+}
+
+//------------------------------------------------------------------------------
+// Getter functions
+//------------------------------------------------------------------------------
+
+bool QueryRegKeyValue(HANDLE key,
+                      const wchar_t* value_name,
+                      ULONG* out_type,
+                      BYTE** out_buffer,
+                      DWORD* out_size) {
+  if (!initialized)
+    InitNativeRegApi();
+
+  NTSTATUS ntstatus = STATUS_UNSUCCESSFUL;
+  UNICODE_STRING value_uni = {};
+  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 = 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 = 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 GetRegValue_DWORD(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 GetRegValue_DWORD(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 (!GetRegValue_DWORD(key, value_name, out_dword)) {
+    CloseRegKey(key);
+    return false;
+  }
+
+  CloseRegKey(key);
+  return true;
+}
+
+// wrapper function
+bool GetRegValue_SZ(HANDLE key,
+                    const wchar_t* value_name,
+                    base::string16* 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 GetRegValue_SZ(ROOT_KEY root,
+                    const wchar_t* key_path,
+                    const wchar_t* value_name,
+                    base::string16* out_sz) {
+  HANDLE key = INVALID_HANDLE_VALUE;
+
+  if (!OpenRegKey(root, key_path, KEY_QUERY_VALUE | KEY_WOW64_32KEY, &key,
+                  NULL))
+    return false;
+
+  if (!GetRegValue_SZ(key, value_name, out_sz)) {
+    CloseRegKey(key);
+    return false;
+  }
+
+  CloseRegKey(key);
+  return true;
+}
+
+// wrapper function
+bool GetRegValue_MULTI_SZ(HANDLE key,
+                          const wchar_t* value_name,
+                          std::vector<base::string16>* 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);
+  base::string16 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 GetRegValue_MULTI_SZ(ROOT_KEY root,
+                          const wchar_t* key_path,
+                          const wchar_t* value_name,
+                          std::vector<base::string16>* out_multi_sz) {
+  HANDLE key = INVALID_HANDLE_VALUE;
+
+  if (!OpenRegKey(root, key_path, KEY_QUERY_VALUE | KEY_WOW64_32KEY, &key,
+                  NULL))
+    return false;
+
+  if (!GetRegValue_MULTI_SZ(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,
+                    BYTE* data,
+                    DWORD data_size) {
+  if (!initialized)
+    InitNativeRegApi();
+
+  NTSTATUS ntstatus = STATUS_UNSUCCESSFUL;
+  UNICODE_STRING value_uni = {};
+  rtl_init_unicode_string(&value_uni, value_name);
+
+  ntstatus = nt_set_value_key(key, &value_uni, 0, type, data, data_size);
+
+  if (NT_SUCCESS(ntstatus))
+    return true;
+
+  return false;
+}
+
+// wrapper function
+bool SetRegValue_DWORD(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 SetRegValue_DWORD(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 (!SetRegValue_DWORD(key, value_name, value)) {
+    CloseRegKey(key);
+    return false;
+  }
+
+  return true;
+}
+
+// wrapper function
+bool SetRegValue_SZ(HANDLE key,
+                    const wchar_t* value_name,
+                    base::string16* value) {
+  wchar_t* string = const_cast<wchar_t*>(value->c_str());
+  if (std::numeric_limits<DWORD>::max() < ::wcslen(string))
+    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 SetRegValue_SZ(ROOT_KEY root,
+                    const wchar_t* key_path,
+                    const wchar_t* value_name,
+                    base::string16* value) {
+  HANDLE key = INVALID_HANDLE_VALUE;
+
+  if (!OpenRegKey(root, key_path, KEY_SET_VALUE | KEY_WOW64_32KEY, &key, NULL))
+    return false;
+
+  if (!SetRegValue_SZ(key, value_name, value)) {
+    CloseRegKey(key);
+    return false;
+  }
+
+  return true;
+}
+
+// wrapper function
+bool SetRegValue_MULTI_SZ(HANDLE key,
+                          const wchar_t* value_name,
+                          std::vector<base::string16>* 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 (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 SetRegValue_MULTI_SZ(ROOT_KEY root,
+                          const wchar_t* key_path,
+                          const wchar_t* value_name,
+                          std::vector<base::string16>* values) {
+  HANDLE key = INVALID_HANDLE_VALUE;
+
+  if (!OpenRegKey(root, key_path, KEY_SET_VALUE | KEY_WOW64_32KEY, &key, NULL))
+    return false;
+
+  if (!SetRegValue_MULTI_SZ(key, value_name, values)) {
+    CloseRegKey(key);
+    return false;
+  }
+
+  return true;
+}
+
+//------------------------------------------------------------------------------
+// Utils
+//------------------------------------------------------------------------------
+
+base::string16 GetCurrentUserSidString() {
+  if (!initialized)
+    InitNativeRegApi();
+
+  return current_user_sid_string;
+}
+
+};  // namespace nt