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
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(¤t_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(¤t_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 |