Commit 191c5780 authored by Penny MacNeil's avatar Penny MacNeil Committed by Commit Bot

[NtRegistry] APIs to help with enumeration added.

QueryRegEnumerationInfo() and QueryRegSubkey().
Also added test suite.

BUG=769590
TEST=chrome_elf_unittests.exe, NtRegistryTest.*

Cq-Include-Trybots: master.tryserver.chromium.win:win10_chromium_x64_rel_ng
Change-Id: I6fb1c08c361c6d9fff001cbe67f9b43da4a988d8
Reviewed-on: https://chromium-review.googlesource.com/673830
Commit-Queue: Penny MacNeil <pennymac@chromium.org>
Reviewed-by: 's avatarGreg Thompson <grt@chromium.org>
Cr-Commit-Position: refs/heads/master@{#506086}
parent 9c0e4bd2
......@@ -17,6 +17,8 @@ NtCreateKeyFunction g_nt_create_key = nullptr;
NtDeleteKeyFunction g_nt_delete_key = nullptr;
NtOpenKeyExFunction g_nt_open_key_ex = nullptr;
NtCloseFunction g_nt_close = nullptr;
NtQueryKeyFunction g_nt_query_key = nullptr;
NtEnumerateKeyFunction g_nt_enumerate_key = nullptr;
NtQueryValueKeyFunction g_nt_query_value_key = nullptr;
NtSetValueKeyFunction g_nt_set_value_key = nullptr;
......@@ -28,6 +30,10 @@ wchar_t g_kRegPathHKLM[] = L"\\Registry\\Machine\\";
wchar_t g_kRegPathHKCU[nt::g_kRegMaxPathLen + 1] = L"";
wchar_t g_current_user_sid_string[nt::g_kRegMaxPathLen + 1] = L"";
// Max number of tries for system API calls when STATUS_BUFFER_TOO_SMALL can be
// returned.
enum { kMaxTries = 5 };
// For testing only.
wchar_t g_HKLM_override[nt::g_kRegMaxPathLen + 1] = L"";
wchar_t g_HKCU_override[nt::g_kRegMaxPathLen + 1] = L"";
......@@ -88,6 +94,12 @@ bool InitNativeRegApi() {
g_nt_close =
reinterpret_cast<NtCloseFunction>(::GetProcAddress(ntdll, "NtClose"));
g_nt_query_key = reinterpret_cast<NtQueryKeyFunction>(
::GetProcAddress(ntdll, "NtQueryKey"));
g_nt_enumerate_key = reinterpret_cast<NtEnumerateKeyFunction>(
::GetProcAddress(ntdll, "NtEnumerateKey"));
g_nt_query_value_key = reinterpret_cast<NtQueryValueKeyFunction>(
::GetProcAddress(ntdll, "NtQueryValueKey"));
......@@ -95,8 +107,8 @@ bool InitNativeRegApi() {
::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)
!g_nt_delete_key || !g_nt_close || !g_nt_query_key ||
!g_nt_enumerate_key || !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.
......@@ -620,8 +632,8 @@ bool CreateRegKey(ROOT_KEY root,
::wcsnlen(key_path, g_kRegMaxPathLen + 1) == g_kRegMaxPathLen + 1)
return false;
if (!g_initialized)
InitNativeRegApi();
if (!g_initialized && !InitNativeRegApi())
return false;
if (root == nt::AUTO)
root = g_system_install ? nt::HKLM : nt::HKCU;
......@@ -711,8 +723,8 @@ bool OpenRegKey(ROOT_KEY root,
::wcsnlen(key_path, g_kRegMaxPathLen + 1) == g_kRegMaxPathLen + 1)
return false;
if (!g_initialized)
InitNativeRegApi();
if (!g_initialized && !InitNativeRegApi())
return false;
NTSTATUS status = STATUS_UNSUCCESSFUL;
UNICODE_STRING key_path_uni = {};
......@@ -746,17 +758,12 @@ bool OpenRegKey(ROOT_KEY root,
}
bool DeleteRegKey(HANDLE key) {
if (!g_initialized)
InitNativeRegApi();
NTSTATUS status = STATUS_UNSUCCESSFUL;
status = g_nt_delete_key(key);
if (!g_initialized && !InitNativeRegApi())
return false;
if (NT_SUCCESS(status))
return true;
NTSTATUS status = g_nt_delete_key(key);
return false;
return NT_SUCCESS(status);
}
// wrapper function
......@@ -791,27 +798,27 @@ bool QueryRegKeyValue(HANDLE key,
const wchar_t* value_name,
ULONG* out_type,
std::vector<BYTE>* out_buffer) {
if (!g_initialized)
InitNativeRegApi();
if (!g_initialized && !InitNativeRegApi())
return false;
NTSTATUS ntstatus = STATUS_UNSUCCESSFUL;
UNICODE_STRING value_uni = {};
g_rtl_init_unicode_string(&value_uni, value_name);
DWORD size_needed = 0;
// 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;
std::unique_ptr<BYTE[]> buffer(new BYTE[size_needed]);
KEY_VALUE_FULL_INFORMATION* value_info =
reinterpret_cast<KEY_VALUE_FULL_INFORMATION*>(buffer.get());
// Use a loop here, to be a little more tolerant of concurrent registry
// changes.
NTSTATUS ntstatus = STATUS_UNSUCCESSFUL;
int tries = 0;
DWORD size_needed = 1;
std::vector<BYTE> buffer;
KEY_VALUE_FULL_INFORMATION* value_info = nullptr;
do {
buffer.resize(size_needed);
value_info = reinterpret_cast<KEY_VALUE_FULL_INFORMATION*>(buffer.data());
ntstatus = g_nt_query_value_key(key, &value_uni, KeyValueFullInformation,
value_info, size_needed, &size_needed);
} while (ntstatus == STATUS_BUFFER_TOO_SMALL && ++tries < kMaxTries);
// 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))
return false;
......@@ -967,8 +974,8 @@ bool SetRegKeyValue(HANDLE key,
ULONG type,
const BYTE* data,
DWORD data_size) {
if (!g_initialized)
InitNativeRegApi();
if (!g_initialized && !InitNativeRegApi())
return false;
NTSTATUS ntstatus = STATUS_UNSUCCESSFUL;
UNICODE_STRING value_uni = {};
......@@ -1092,27 +1099,99 @@ bool SetRegValueMULTISZ(ROOT_KEY root,
return true;
}
//------------------------------------------------------------------------------
// Enumeration Support
//------------------------------------------------------------------------------
bool QueryRegEnumerationInfo(HANDLE key, ULONG* out_subkey_count) {
if (!g_initialized && !InitNativeRegApi())
return false;
// Use a loop here, to be a little more tolerant of concurrent registry
// changes.
NTSTATUS ntstatus = STATUS_UNSUCCESSFUL;
int tries = 0;
// Start with sizeof the structure. It's very common for the variable sized
// "Class" element to be of length 0.
DWORD size_needed = sizeof(KEY_FULL_INFORMATION);
std::vector<BYTE> buffer;
KEY_FULL_INFORMATION* key_info = nullptr;
do {
buffer.resize(size_needed);
key_info = reinterpret_cast<KEY_FULL_INFORMATION*>(buffer.data());
ntstatus = g_nt_query_key(key, KeyFullInformation, key_info, size_needed,
&size_needed);
} while (ntstatus == STATUS_BUFFER_TOO_SMALL && ++tries < kMaxTries);
if (!NT_SUCCESS(ntstatus))
return false;
// Move desired information to out variables.
*out_subkey_count = key_info->SubKeys;
return true;
}
bool QueryRegSubkey(HANDLE key,
ULONG subkey_index,
std::wstring* out_subkey_name) {
if (!g_initialized && !InitNativeRegApi())
return false;
// Use a loop here, to be a little more tolerant of concurrent registry
// changes.
NTSTATUS ntstatus = STATUS_UNSUCCESSFUL;
int tries = 0;
// Start with sizeof the structure, plus 12 characters. It's very common for
// key names to be < 12 characters (without being inefficient as an initial
// allocation).
DWORD size_needed = sizeof(KEY_BASIC_INFORMATION) + (12 * sizeof(wchar_t));
std::vector<BYTE> buffer;
KEY_BASIC_INFORMATION* subkey_info = nullptr;
do {
buffer.resize(size_needed);
subkey_info = reinterpret_cast<KEY_BASIC_INFORMATION*>(buffer.data());
ntstatus = g_nt_enumerate_key(key, subkey_index, KeyBasicInformation,
subkey_info, size_needed, &size_needed);
} while (ntstatus == STATUS_BUFFER_TOO_SMALL && ++tries < kMaxTries);
if (!NT_SUCCESS(ntstatus))
return false;
// Move desired information to out variables.
// NOTE: NameLength is size of Name array in bytes. Name array is also
// NOT null terminated!
BYTE* name = reinterpret_cast<BYTE*>(subkey_info->Name);
std::vector<BYTE> content(name, name + subkey_info->NameLength);
EnsureTerminatedSZ(&content, false);
out_subkey_name->assign(reinterpret_cast<wchar_t*>(content.data()));
return true;
}
//------------------------------------------------------------------------------
// Utils
//------------------------------------------------------------------------------
const wchar_t* GetCurrentUserSidString() {
if (!g_initialized)
InitNativeRegApi();
if (!g_initialized && !InitNativeRegApi())
return nullptr;
return g_current_user_sid_string;
}
bool IsCurrentProcWow64() {
if (!g_initialized)
InitNativeRegApi();
if (!g_initialized && !InitNativeRegApi())
return false;
return g_wow64_proc;
}
bool SetTestingOverride(ROOT_KEY root, const std::wstring& new_path) {
if (!g_initialized)
InitNativeRegApi();
if (!g_initialized && !InitNativeRegApi())
return false;
std::wstring sani_new_path = new_path;
SanitizeSubkeyPath(&sani_new_path);
......@@ -1128,8 +1207,8 @@ bool SetTestingOverride(ROOT_KEY root, const std::wstring& new_path) {
}
std::wstring GetTestingOverride(ROOT_KEY root) {
if (!g_initialized)
InitNativeRegApi();
if (!g_initialized && !InitNativeRegApi())
return std::wstring();
if (root == HKCU || (root == AUTO && !g_system_install))
return g_HKCU_override;
......
......@@ -13,8 +13,10 @@
// Note that this API is currently lazy initialized. Any function that is
// NOT merely a wrapper function (i.e. any function that directly interacts with
// NTDLL) will immediately check:
// if (!g_initialized)
// InitNativeRegApi();
//
// if (!g_initialized && !InitNativeRegApi())
// return false;
//
// There is currently no multi-threading lock around the lazy initialization,
// as the main client for this API (chrome_elf) does not introduce
// a multi-threading concern. This can easily be changed if needed.
......@@ -73,11 +75,13 @@ bool OpenRegKey(ROOT_KEY root,
// Delete a registry key.
// - Caller must still call CloseRegKey after the delete.
// - Non-recursive. Must have no subkeys.
bool DeleteRegKey(HANDLE key);
// Delete a registry key.
// - WRAPPER: Function opens and closes the target key for caller.
// - Use |wow64_override| to force redirection behaviour, or pass nt::NONE.
// - Non-recursive. Must have no subkeys.
bool DeleteRegKey(ROOT_KEY root,
WOW64_OVERRIDE wow64_override,
const wchar_t* key_path);
......@@ -222,6 +226,25 @@ bool SetRegValueMULTISZ(ROOT_KEY root,
const wchar_t* value_name,
const std::vector<std::wstring>& values);
//------------------------------------------------------------------------------
// Enumeration Support
//------------------------------------------------------------------------------
// Query key information for subkey enumeration.
// - Key handle should have been opened with OpenRegKey (with at least
// KEY_ENUMERATE_SUB_KEYS access rights).
// - Currently only returns the number of subkeys. Use |subkey_count|
// in a loop for calling QueryRegSubkey.
bool QueryRegEnumerationInfo(HANDLE key, ULONG* out_subkey_count);
// Enumerate subkeys by index.
// - Key handle should have been opened with OpenRegKey (with at least
// KEY_ENUMERATE_SUB_KEYS access rights).
// - Get subkey count by calling QueryRegEnumerationInfo.
bool QueryRegSubkey(HANDLE key,
ULONG subkey_index,
std::wstring* out_subkey_name);
//------------------------------------------------------------------------------
// Utils
//------------------------------------------------------------------------------
......
......@@ -409,6 +409,33 @@ typedef VOID(WINAPI* RtlFreeUnicodeStringFunction)(
// -----------------------------------------------------------------------
// Registry
typedef enum _KEY_INFORMATION_CLASS {
KeyBasicInformation = 0,
KeyFullInformation = 2
} KEY_INFORMATION_CLASS,
*PKEY_INFORMATION_CLASS;
typedef struct _KEY_BASIC_INFORMATION {
LARGE_INTEGER LastWriteTime;
ULONG TitleIndex;
ULONG NameLength;
WCHAR Name[1];
} KEY_BASIC_INFORMATION, *PKEY_BASIC_INFORMATION;
typedef struct _KEY_FULL_INFORMATION {
LARGE_INTEGER LastWriteTime;
ULONG TitleIndex;
ULONG ClassOffset;
ULONG ClassLength;
ULONG SubKeys;
ULONG MaxNameLen;
ULONG MaxClassLen;
ULONG Values;
ULONG MaxValueNameLen;
ULONG MaxValueDataLen;
WCHAR Class[1];
} KEY_FULL_INFORMATION, *PKEY_FULL_INFORMATION;
typedef enum _KEY_VALUE_INFORMATION_CLASS {
KeyValueFullInformation = 1
} KEY_VALUE_INFORMATION_CLASS,
......@@ -449,6 +476,21 @@ typedef NTSTATUS (WINAPI *NtDeleteKeyFunction)(
typedef NTSTATUS(WINAPI* RtlFormatCurrentUserKeyPathFunction)(
OUT PUNICODE_STRING RegistryPath);
typedef NTSTATUS(WINAPI* NtQueryKeyFunction)(IN HANDLE KeyHandle,
IN KEY_INFORMATION_CLASS
KeyInformationClass,
OUT PVOID KeyInformation,
IN ULONG Length,
OUT PULONG ResultLength);
typedef NTSTATUS(WINAPI* NtEnumerateKeyFunction)(IN HANDLE KeyHandle,
IN ULONG Index,
IN KEY_INFORMATION_CLASS
KeyInformationClass,
OUT PVOID KeyInformation,
IN ULONG Length,
OUT PULONG ResultLength);
typedef NTSTATUS(WINAPI* NtQueryValueKeyFunction)(IN HANDLE KeyHandle,
IN PUNICODE_STRING ValueName,
IN KEY_VALUE_INFORMATION_CLASS
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment