[chrome_elf] NTRegistry - added wow64 redirection support.

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"
 
+#include <assert.h>
+#include <stdlib.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;
+bool g_wow64_proc = false;
 bool g_reg_redirection = false;
 const size_t g_kMaxPathLen = 255;
 wchar_t g_kRegPathHKLM[] = L"\\Registry\\Machine\\";
 wchar_t g_kRegPathHKCU[g_kMaxPathLen] = L"";
 wchar_t g_current_user_sid_string[g_kMaxPathLen] = L"";
 wchar_t g_override_path[g_kMaxPathLen] = L"";
 
-// Not using install_util, to prevent circular dependency.
+//------------------------------------------------------------------------------
+// Initialization - LOCAL
+//------------------------------------------------------------------------------
+
+// Not using install_static, 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 IsThisProcWow64() {

[robertshield, 2016/09/28 23:49:23]

This function looks an awful lot like
https://codesearch.chromium.org/chromium/src/chrome_elf/hook_util/hook_util.c...

Could this be commonalizified (totally a word) somehow?

[penny, 2016/09/30 05:28:22]

Commonalizified.  I've exposed this information with a new nt_registry function
IsCurrentProcWow64().  hook_util now calls that API.  This is nice because the
info only should need to be checked once, and then can be stored.

+  // Using BOOL type for compat with IsWow64Process() system API.
+  BOOL is_wow64 = FALSE;
+
+  // API might not exist, so dynamic lookup.
+  using IsWow64ProcessFunction = decltype(&IsWow64Process);
+  IsWow64ProcessFunction is_wow64_process =
+      reinterpret_cast<IsWow64ProcessFunction>(::GetProcAddress(
+          ::GetModuleHandle(L"kernel32.dll"), "IsWow64Process"));
+  if (!is_wow64_process)
+    return false;
+  if (!is_wow64_process(::GetCurrentProcess(), &is_wow64))
+    return false;
+  return is_wow64 ? true : 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"));
 
@@ skipping 36 matching lines @@
   // Finish setting up global HKCU path.
   ::wcsncat(g_kRegPathHKCU, current_user_reg_path.Buffer, (g_kMaxPathLen - 1));
   ::wcsncat(g_kRegPathHKCU, L"\\",
             (g_kMaxPathLen - ::wcslen(g_kRegPathHKCU) - 1));
   // Keep the sid string as well.
   wchar_t* ptr = ::wcsrchr(current_user_reg_path.Buffer, L'\\');
   ptr++;
   ::wcsncpy(g_current_user_sid_string, ptr, (g_kMaxPathLen - 1));
   rtl_free_unicode_str(&current_user_reg_path);
 
-  // Figure out if we're a system or user install.
+  // Figure out if this is a system or user install.
   g_system_install = IsThisProcSystem();
 
+  // Figure out if this is a WOW64 process.
+  g_wow64_proc = IsThisProcWow64();
+
   g_initialized = true;
   return true;
 }
 
+//------------------------------------------------------------------------------
+// Reg WOW64 Redirection - LOCAL
+//
+// How registry redirection works directly calling NTDLL APIs:
+// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+// - NOTE: On >= Win7, reflection support was removed.
+// -
+// https://msdn.microsoft.com/en-us/library/windows/desktop/aa384253(v=vs.85).aspx
+//
+// - 1) 32-bit / WOW64 process:
+//     a) Default access WILL be redirected to WOW64.
+//     b) KEY_WOW64_32KEY access WILL be redirected to WOW64.
+//     c) KEY_WOW64_64KEY access will NOT be redirected to WOW64.
+//
+// - 2) 64-bit process:
+//     a) Default access will NOT be redirected to WOW64.
+//     b) KEY_WOW64_32KEY access will NOT be redirected to WOW64.
+//     c) KEY_WOW64_64KEY access will NOT be redirected to WOW64.
+//
+// - Key point from above is that NTDLL redirects and respects access
+//   overrides for WOW64 calling processes.  But does NOT do any of that if the
+//   calling process is 64-bit.  2b is surprising and troublesome.
+//
+// How registry redirection works using these nt_registry APIs:
+// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+// - These APIs will behave the same as NTDLL above, EXCEPT for 2b.
+//   nt_registry APIs will respect the override access flags for all processes.
+//
+// - How the WOW64 redirection decision trees / Nodes work below:
+//
+//   The HKLM and HKCU descision trees represent the information at the MSDN

[robertshield, 2016/09/28 23:49:23]

nit: decision

[penny, 2016/09/30 05:28:22]

Thank you.  I dislike spelling or grammatical errors!

+//   link above... but in a way that generates a decision about whether a
+//   registry path should be subject to WOW64 redirection.  The tree is
+//   traversed as you scan along the registry path in question.
+//
+//    - Each Node contains a chunk of registry subkey(s) to match.
+//    - If it is NOT matched, traversal is done.
+//    - If it is matched:
+//       - Current state of |redirection_type| for the whole registry path is
+//         updated.
+//       - If |next| is empty, traversal is done.
+//       - Otherwise, |next| is an array of child Nodes to try to match against.
+//         Loop.
+//------------------------------------------------------------------------------
+
+// This enum defines states for how to handle redirection.
+// NOTE: When WOW64 redirection should happen, the redirect subkey can be either
+//       before or after the latest Node match.  Unfortunately not consistent.
+enum RedirectionType { SHARED = 0, REDIRECTED_BEFORE, REDIRECTED_AFTER };
+
+struct Node {
+  template <size_t len, size_t n_len>
+  constexpr Node(const wchar_t (&wcs)[len],
+                 RedirectionType rt,
+                 const Node (&n)[n_len])
+      : to_match(wcs),
+        to_match_len(len - 1),
+        redirection_type(rt),
+        next(n),
+        next_len(n_len) {}
+
+  template <size_t len>
+  constexpr Node(const wchar_t (&wcs)[len], RedirectionType rt)
+      : to_match(wcs),
+        to_match_len(len - 1),
+        redirection_type(rt),
+        next(nullptr),
+        next_len(0) {}
+
+  const wchar_t* to_match;
+  size_t to_match_len;
+  // If a match, this is the new state of how to redirect.
+  RedirectionType redirection_type;
+  // |next| is nullptr or an array of Nodes of length |array_len|.
+  const Node* next;
+  size_t next_len;
+};
+
+// HKLM or HKCU SOFTWARE\Classes is shared by default.  Specific subkeys under
+// Classes are redirected to SOFTWARE\WOW6432Node\Classes\<subkey> though.
+constexpr Node classes_subtree[] = {{L"CLSID", REDIRECTED_BEFORE},

[grt (UTC plus 2), 2016/09/29 10:30:22]

nit: kClassesSubtree

[penny, 2016/09/30 05:28:22]

Sorry - I should have changed these at the same time as the fixes in the test
file.

+                                    {L"DirectShow", REDIRECTED_BEFORE},
+                                    {L"Interface", REDIRECTED_BEFORE},
+                                    {L"Media Type", REDIRECTED_BEFORE},
+                                    {L"MediaFoundation", REDIRECTED_BEFORE}};
+
+// These specific HKLM\SOFTWARE subkeys are shared.  Specific
+// subkeys under Classes are redirected though... see classes_subtree.
+constexpr Node hklm_software_subtree[] = {

[grt (UTC plus 2), 2016/09/29 10:30:22]

nit: kHklmSoftwareSubtree

[penny, 2016/09/30 05:28:22]

Sorry - I should have changed these at the same time as the fixes in the test
file.

+    // TODO(pennymac): when MS fixes compiler bug, or bots are all using clang,
+    // remove the "Classes" subkeys below and replace with:
+    // {L"Classes", SHARED, classes_subtree},
+    // https://connect.microsoft.com/VisualStudio/feedback/details/3104499
+    {L"Classes\\CLSID", REDIRECTED_BEFORE},
+    {L"Classes\\DirectShow", REDIRECTED_BEFORE},
+    {L"Classes\\Interface", REDIRECTED_BEFORE},
+    {L"Classes\\Media Type", REDIRECTED_BEFORE},
+    {L"Classes\\MediaFoundation", REDIRECTED_BEFORE},
+    {L"Classes", SHARED},
+
+    {L"Clients", SHARED},
+    {L"Microsoft\\COM3", SHARED},
+    {L"Microsoft\\Cryptography\\Calais\\Current", SHARED},
+    {L"Microsoft\\Cryptography\\Calais\\Readers", SHARED},
+    {L"Microsoft\\Cryptography\\Services", SHARED},
+
+    {L"Microsoft\\CTF\\SystemShared", SHARED},
+    {L"Microsoft\\CTF\\TIP", SHARED},
+    {L"Microsoft\\DFS", SHARED},
+    {L"Microsoft\\Driver Signing", SHARED},
+    {L"Microsoft\\EnterpriseCertificates", SHARED},
+
+    {L"Microsoft\\EventSystem", SHARED},
+    {L"Microsoft\\MSMQ", SHARED},
+    {L"Microsoft\\Non-Driver Signing", SHARED},
+    {L"Microsoft\\Notepad\\DefaultFonts", SHARED},
+    {L"Microsoft\\OLE", SHARED},
+
+    {L"Microsoft\\RAS", SHARED},
+    {L"Microsoft\\RPC", SHARED},
+    {L"Microsoft\\SOFTWARE\\Microsoft\\Shared Tools\\MSInfo", SHARED},
+    {L"Microsoft\\SystemCertificates", SHARED},
+    {L"Microsoft\\TermServLicensing", SHARED},
+
+    {L"Microsoft\\Transaction Server", SHARED},
+    {L"Microsoft\\Windows\\CurrentVersion\\App Paths", SHARED},
+    {L"Microsoft\\Windows\\CurrentVersion\\Control Panel\\Cursors\\Schemes",
+     SHARED},
+    {L"Microsoft\\Windows\\CurrentVersion\\Explorer\\AutoplayHandlers", SHARED},
+    {L"Microsoft\\Windows\\CurrentVersion\\Explorer\\DriveIcons", SHARED},
+
+    {L"Microsoft\\Windows\\CurrentVersion\\Explorer\\KindMap", SHARED},
+    {L"Microsoft\\Windows\\CurrentVersion\\Group Policy", SHARED},
+    {L"Microsoft\\Windows\\CurrentVersion\\Policies", SHARED},
+    {L"Microsoft\\Windows\\CurrentVersion\\PreviewHandlers", SHARED},
+    {L"Microsoft\\Windows\\CurrentVersion\\Setup", SHARED},
+
+    {L"Microsoft\\Windows\\CurrentVersion\\Telephony\\Locations", SHARED},
+    {L"Microsoft\\Windows NT\\CurrentVersion\\Console", SHARED},
+    {L"Microsoft\\Windows NT\\CurrentVersion\\FontDpi", SHARED},
+    {L"Microsoft\\Windows NT\\CurrentVersion\\FontLink", SHARED},
+    {L"Microsoft\\Windows NT\\CurrentVersion\\FontMapper", SHARED},
+
+    {L"Microsoft\\Windows NT\\CurrentVersion\\Fonts", SHARED},
+    {L"Microsoft\\Windows NT\\CurrentVersion\\FontSubstitutes", SHARED},
+    {L"Microsoft\\Windows NT\\CurrentVersion\\Gre_Initialize", SHARED},
+    {L"Microsoft\\Windows NT\\CurrentVersion\\Image File Execution Options",
+     SHARED},
+    {L"Microsoft\\Windows NT\\CurrentVersion\\LanguagePack", SHARED},
+
+    {L"Microsoft\\Windows NT\\CurrentVersion\\NetworkCards", SHARED},
+    {L"Microsoft\\Windows NT\\CurrentVersion\\Perflib", SHARED},
+    {L"Microsoft\\Windows NT\\CurrentVersion\\Ports", SHARED},
+    {L"Microsoft\\Windows NT\\CurrentVersion\\Print", SHARED},
+    {L"Microsoft\\Windows NT\\CurrentVersion\\ProfileList", SHARED},
+
+    {L"Microsoft\\Windows NT\\CurrentVersion\\Time Zones", SHARED},
+    {L"Policies", SHARED},
+    {L"RegisteredApplications", SHARED}};
+
+// HKCU is entirely shared, except for a few specific Classes subkeys which
+// are redirected.  See |classes_subtree|.
+constexpr Node g_redirectionDecisionTreeHKCU = {L"SOFTWARE\\Classes", SHARED,

[grt (UTC plus 2), 2016/09/29 10:30:22]

nit: kRedirectionDecisionTreeHkcu

[penny, 2016/09/30 05:28:22]

Done.

+                                                classes_subtree};
+
+// HKLM\SOFTWARE is redirected by default to SOFTWARE\WOW6432Node.  Specific
+// subkeys under SOFTWARE are shared though... see |hklm_software_subtree|.
+constexpr Node g_redirectionDecisionTreeHKLM = {L"SOFTWARE", REDIRECTED_AFTER,

[grt (UTC plus 2), 2016/09/29 10:30:22]

nit: kRedirectionDecisionTreeHklm

[penny, 2016/09/30 05:28:22]

Done.

+                                                hklm_software_subtree};
+
+// Main redirection handler function.
+// If redirection is required, change is made to |subkey_path| in place.
+//
+// - This function should be called BEFORE concatenating |subkey_path| with the
+//   root hive or calling ParseFullRegPath().
+// - Also, |subkey_path| should be passed to SanitizeSubkeyPath() before calling

[grt (UTC plus 2), 2016/09/29 10:30:22]

how about enforcing this by adding to the top of the function:
  assert(subkey_path->empty() || subkey_path->front() != L'\\');
  assert(subkey_path->empty() || subkey_path->back() != L'\\');

better yet, can you make life easier on the callers by calling Sanitize in here
rather than making them do it?

[penny, 2016/09/30 05:28:22]

My take on this is that external APIs should be responsible for verifying and
sanitizing externally-supplied input.  They should gracefully handle any bad
input in debug or release builds.  I feel that private functions are fine to
simply debug assert on broken expectations as a red flag.

ProcessRedirection() is an internal function that will only be called from
within this file.  I've added the asserts you mentioned - just to catch
developer mistakes within this file.

By only sanitizing external inputs, we won't end up doing the same run-time
checks more than once.  (Private functions can easily be called from many
places.)

**I've also just added a |value_name| input check to QueryRegKeyValue() and
SetRegKeyValue().

[grt (UTC plus 2), 2016/09/30 09:32:16]

Your philosophy on when to sanitize external input makes total sense. My
observation was that (at the time of writing), SanitizeSubkeyPath was only
called immediately before ProcessRedirection, and ProcessRedirection was never
called without being immediately preceded by SanitizeSubkeyPath. Is sanitation
really needed, or can that step be removed by gently adjusting the redirector?

[penny, 2016/10/01 01:44:25]

I hear ya, and considered just changing ProcessRedirection to handle it. 
However, now I've got another external API for setting the testing override that
will use this sanitize function.

The lib on a whole assumes that the subkey paths it is storing and dealing with
won't have starting or trailing slashes - and I think it's nice to have a
consistent format expectation.

I've added an extra bit in there to ensure there are no sequences of more than
one backslash in a row as well.

[grt (UTC plus 2), 2016/10/02 20:10:53]

Do the underlying Nt registry functions choke on multiple \\\? Does advapi strip
them out?

[penny, 2016/10/03 19:18:48]

I'm not sure at the moment.  It (advapi or ntdll) certainly won't fail out - it
just handles them.  I resolved them in Sanitize() so that private internal
subkey string manipulation doesn't have to worry about this.

+//   this function.
+void ProcessRedirection(nt::ROOT_KEY root,
+                        ACCESS_MASK access,
+                        std::wstring* subkey_path) {
+  static constexpr wchar_t kRedirectBefore[] = L"WOW6432Node\\";
+  static constexpr wchar_t kRedirectAfter[] = L"\\WOW6432Node";
+
+  if (subkey_path == nullptr || subkey_path->empty() ||
+      (access & KEY_WOW64_32KEY && access & KEY_WOW64_64KEY))
+    return;
+
+  // Convert nt::AUTO to the appropriate root key.
+  nt::ROOT_KEY temp_root = root;
+  if (root == nt::AUTO)
+    temp_root = g_system_install ? nt::HKLM : nt::HKCU;
+
+  // No redirection during testing when there's already an override.
+  // Otherwise, the testing redirect directory Software\Chromium\TempTestKeys
+  // would get WOW64 redirected if root_key == HKLM in this function.
+  if ((temp_root == nt::HKCU && *nt::HKCU_override) ||
+      (temp_root == nt::HKLM && *nt::HKLM_override))
+    return;
+
+  // WOW64 redirection only supported on x64 architecture.  Return if x86.
+  SYSTEM_INFO system_info = {};
+  ::GetNativeSystemInfo(&system_info);
+  if (system_info.wProcessorArchitecture == PROCESSOR_ARCHITECTURE_INTEL)
+    return;
+
+  bool use_wow64 = g_wow64_proc;
+  // Consider KEY_WOW64_32KEY and KEY_WOW64_64KEY override access flags.
+  if (access & KEY_WOW64_32KEY)
+    use_wow64 = true;
+  if (access & KEY_WOW64_64KEY)
+    use_wow64 = false;
+
+  // If !use_wow64, there's nothing more to do.
+  if (!use_wow64)
+    return;
+
+  // Pick which decision tree to use.
+  const Node* current = (temp_root == nt::HKCU)
+                            ? &g_redirectionDecisionTreeHKCU
+                            : &g_redirectionDecisionTreeHKLM;
+
+  // The following loop works on the |subkey_path| from left to right.
+  // |position| tracks progress along |subkey_path|.
+  const wchar_t* position = subkey_path->c_str();
+  // Hold a count of chars left after position, for efficient calculations.
+  size_t chars_left = subkey_path->length();
+  // |redirect_state| holds the latest state of redirection requirement.
+  RedirectionType redirect_state = SHARED;
+  // |insertion_point| tracks latest spot for redirection subkey to be inserted.
+  const wchar_t* insertion_point = nullptr;
+  // |insert_string| tracks which redirection string would be inserted.
+  const wchar_t* insert_string = nullptr;
+
+  // The root of the tree is an array of 1.

[grt (UTC plus 2), 2016/09/29 10:30:22]

uber-nit: move the definition of |array_len| up so that it immediately precedes
the definition of |current| on line 352. this makes it more clear that it refers
to the length of |current|, and matches the order in which the two are updated
on line 402.

[penny, 2016/09/30 05:28:22]

Done.

+  size_t array_len = 1;

[grt (UTC plus 2), 2016/09/29 10:30:22]

nit: array_len -> node_array_len

[penny, 2016/09/30 05:28:23]

I like this, thanks.  More clear.  I've also changed current -> current_node.

+  size_t index = 0;

[grt (UTC plus 2), 2016/09/29 10:30:22]

nit: index -> node_index

[penny, 2016/09/30 05:28:23]

Done.

+  while (index < array_len) {
+    size_t len = current->to_match_len;

[grt (UTC plus 2), 2016/09/29 10:30:22]

nit: len -> to_match_len

[penny, 2016/09/30 05:28:22]

Done. current_to_match_len.  I don't want any confusion with the Node member
|to_match_len|.

+    // Make sure the remainder of the path is at least as long as the current
+    // subkey to match.
+    if (chars_left >= len) {
+      // Do case insensitive comparisons.
+      if (::wcsnicmp(position, current->to_match, len) == 0) {
+        // Make sure not to match on a substring.
+        if (*(position + len) == L'\\' || *(position + len) == L'\0') {
+          // MATCH!
+          // -------------------------------------------------------------------
+          // 1) Update state of redirection.
+          redirect_state = current->redirection_type;
+          // 1.5) If new state is to redirect, the new insertion point will be
+          //      either right before or right after this match.
+          if (redirect_state == REDIRECTED_BEFORE) {
+            insertion_point = position;
+            insert_string = kRedirectBefore;
+          } else if (redirect_state == REDIRECTED_AFTER) {
+            insertion_point = position + len;
+            insert_string = kRedirectAfter;
+          }
+          // 2) Adjust |position| along the subkey path.
+          position += len;
+          chars_left -= len;
+          // 2.5) Increment the position, to move past path seperator(s).
+          while (*position == L'\\') {
+            ++position;
+            --chars_left;
+          }
+          // 3) Move our loop parameters to the |next| array of Nodes.
+          array_len = current->next_len;
+          current = current->next;
+          index = 0;
+          // 4) Finish this loop and start on new array.
+          continue;
+        }
+      }
+    }
+
+    // Move to the next node in the array if we didn't match this loop.
+    ++current;
+    ++index;
+  }
+
+  if (redirect_state == SHARED)
+    return;
+
+  // Insert the redirection into |subkey_path|, at |insertion_point|.
+  subkey_path->insert((insertion_point - subkey_path->data()), insert_string);

[grt (UTC plus 2), 2016/09/29 10:30:22]

nit: data -> c_str for consistency with the code above

[penny, 2016/09/30 05:28:22]

Done.

+}
+
+//------------------------------------------------------------------------------
+// Reg Path Utilities - LOCAL
+//------------------------------------------------------------------------------
+
 const wchar_t* ConvertRootKey(nt::ROOT_KEY root) {
   nt::ROOT_KEY key = root;
 
   if (!root) {

[robertshield, 2016/09/28 23:49:23]

should this be if (root == nt::AUTO)?

[grt (UTC plus 2), 2016/09/29 10:30:22]

yes, please

[penny, 2016/09/30 05:28:22]

Done.  Sorry - missed this one!

     // AUTO
     key = g_system_install ? nt::HKLM : nt::HKCU;
   }
 
-  if ((key == nt::HKCU) && (::wcslen(nt::HKCU_override) != 0)) {
+  if (key == nt::HKCU && *nt::HKCU_override) {
     std::wstring temp(g_kRegPathHKCU);
     temp.append(nt::HKCU_override);
     temp.append(L"\\");
     ::wcsncpy(g_override_path, temp.c_str(), g_kMaxPathLen - 1);
     g_reg_redirection = true;
     return g_override_path;
-  } else if ((key == nt::HKLM) && (::wcslen(nt::HKLM_override) != 0)) {
+  } else if (key == nt::HKLM && *nt::HKLM_override) {
     std::wstring temp(g_kRegPathHKCU);
     temp.append(nt::HKLM_override);
     temp.append(L"\\");
     ::wcsncpy(g_override_path, temp.c_str(), g_kMaxPathLen - 1);
     g_reg_redirection = true;
     return g_override_path;
   }
 
   g_reg_redirection = false;
   if (key == nt::HKCU)
     return g_kRegPathHKCU;
   else
     return g_kRegPathHKLM;
 }
 
+// This utility should be called on an externally provided subkey path.
+// - Ensures there are no starting or trailing backslashes.
+// - Note from MSDN: "Key names cannot include the backslash character (\),
+//   but any other printable character can be used."
+void SanitizeSubkeyPath(std::wstring* input) {
+  if (input == nullptr || input->empty())

[grt (UTC plus 2), 2016/09/29 10:30:22]

two comments here:

since this is a private function and you know the callers never pass in nullptr,
replace the first part of this condition with an assert.

if it is true that |input| is not expected to be empty in the general case,
remove the second part of this and let empty inputs be handled below. this
avoids having an extra branch for the ordinary case to optimize for the rare
case.

[penny, 2016/09/30 05:28:22]

I fully agree.  Thank you.

+    return;
+
+  // Remove any trailing backslashes.
+  while (input->back() == L'\\')

[grt (UTC plus 2), 2016/09/29 10:30:22]

this provokes undefined behavior if the string contains only backslashes. how
about:
  size_t last_valid_pos = input->find_last_not_of(L'\\');
  if (last_valid_pos == std::wstring::npos) {
    // The string is all backslashes. Clear it and abort.
    input->clear();
    return;
  }
  // Chop of the trailing backslashes.
  input->resize(last_valid_pos + 1);

  // Remove leading backslashes.
  input->erase(0, input->find_first_not_of(L'\\'));

and then please add a unittest for this. :-)

[penny, 2016/09/30 05:28:22]

TIL, don't call front/back on an empty std::wstring, and at throws exceptions! 
Thanks!
How's this change?

[grt (UTC plus 2), 2016/09/30 09:32:16]

I think it's much more clear and safe to use the string API rather than rolling
your own. I also think you should avoid substr since it will create a new
instance and make a whole new allocation. erase will effectively be a memmove,
which i think is better than malloc->memcpy->free.

[penny, 2016/10/01 01:44:25]

I'm fully agreeing with you now.  This is much nicer and more efficient.  

+    input->pop_back();
+
+  // Remove any starting backslashes.
+  size_t index = 0;
+  while (input->at(index) == L'\\')

[grt (UTC plus 2), 2016/09/29 10:30:22]

fyi: at() does bounds checking and throws exceptions, so avoid it at all costs.

[penny, 2016/09/30 05:28:22]

TIL.  Thanks.  Reading the docs, I think [] operator will work fine.

+    ++index;
+  if (!index)
+    *input = input->substr(index);
+}
+
 // Turns a root and subkey path into the registry base hive and the rest of the
 // subkey tokens.
 // - |converted_root| should come directly out of ConvertRootKey function.
+// - |subkey_path| should be passed to SanitizeSubkeyPath() first.
 // - E.g. base hive: "\Registry\Machine\", "\Registry\User\<SID>\".
 bool ParseFullRegPath(const wchar_t* converted_root,
                       const wchar_t* subkey_path,
                       std::wstring* out_base,
                       std::vector<std::wstring>* subkeys) {
   out_base->clear();
   subkeys->clear();
   std::wstring temp = L"";
 
   if (g_reg_redirection) {
@@ skipping 39 matching lines @@
       out_base->push_back(L'\\');
     }
     subkeys->erase(subkeys->begin(), subkeys->begin() + num_base_tokens);
   } else {
     out_base->assign(converted_root);
   }
 
   return true;
 }
 
+//------------------------------------------------------------------------------
+// Misc wrapper functions - LOCAL
+//------------------------------------------------------------------------------
+
 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 {
 
 const size_t g_kRegMaxPathLen = 255;
 wchar_t HKLM_override[g_kRegMaxPathLen] = L"";

[grt (UTC plus 2), 2016/09/29 10:30:22]

wdyt of adding 1 to the length here and below so that g_kRegMaxPathLen doesn't
include the terminator? then you can get rid of the " - 1" in the test.

[penny, 2016/09/30 05:28:22]

Done.  Thanks.  I've adjusted this everywhere to be accurate for Windows
registry.  255 is max length without null character.

 wchar_t HKCU_override[g_kRegMaxPathLen] = L"";
 
 //------------------------------------------------------------------------------
 // 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();
 
+  std::wstring redirected_key_path;
+  if (key_path) {
+    redirected_key_path = key_path;
+    SanitizeSubkeyPath(&redirected_key_path);
+    ProcessRedirection(root, access, &redirected_key_path);
+  }
+
   std::wstring current_path;
   std::vector<std::wstring> subkeys;
-  if (!ParseFullRegPath(ConvertRootKey(root), key_path, &current_path,
-                        &subkeys))
+  if (!ParseFullRegPath(ConvertRootKey(root), redirected_key_path.c_str(),

[grt (UTC plus 2), 2016/09/29 10:30:22]

pass redirected_key_path directly as a const std::wstring&

[penny, 2016/09/30 05:28:22]

Done.

+                        &current_path, &subkeys))
     return false;
 
   // Open the base hive first.  It should always exist already.
   HANDLE last_handle = INVALID_HANDLE_VALUE;
   NTSTATUS status =
       CreateKeyWrapper(current_path, access, &last_handle, nullptr);
   if (!NT_SUCCESS(status))
     return false;
 
   size_t subkeys_size = subkeys.size();
@@ skipping 56 matching lines @@
                 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);
+  std::wstring full_path;
+  if (key_path) {
+    full_path = key_path;
+    SanitizeSubkeyPath(&full_path);
+    ProcessRedirection(root, access, &full_path);
+  }
+  full_path.insert(0, ConvertRootKey(root));
 
   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;
 
@@ skipping 11 matching lines @@
 
   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) {
+bool DeleteRegKey(ROOT_KEY root,
+                  WOW64_OVERRIDE wow64_override,
+                  const wchar_t* key_path) {
   HANDLE key = INVALID_HANDLE_VALUE;
 
-  if (!OpenRegKey(root, key_path, DELETE, &key, nullptr))
+  if (!OpenRegKey(root, key_path, DELETE | wow64_override, &key, nullptr))
     return false;
 
   if (!DeleteRegKey(key)) {
     CloseRegKey(key);
     return false;
   }
 
   CloseRegKey(key);
   return true;
 }
@@ skipping 61 matching lines @@
     return false;
 
   *out_dword = *(reinterpret_cast<DWORD*>(value_bytes));
 
   delete[] value_bytes;
   return true;
 }
 
 // wrapper function
 bool QueryRegValueDWORD(ROOT_KEY root,
+                        WOW64_OVERRIDE wow64_override,
                         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))
+  if (!OpenRegKey(root, key_path, KEY_QUERY_VALUE | wow64_override, &key, NULL))
     return false;
 
   if (!QueryRegValueDWORD(key, value_name, out_dword)) {
     CloseRegKey(key);
     return false;
   }
 
   CloseRegKey(key);
   return true;
 }
@@ skipping 11 matching lines @@
     return false;
 
   *out_sz = reinterpret_cast<wchar_t*>(value_bytes);
 
   delete[] value_bytes;
   return true;
 }
 
 // wrapper function
 bool QueryRegValueSZ(ROOT_KEY root,
+                     WOW64_OVERRIDE wow64_override,
                      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))
+  if (!OpenRegKey(root, key_path, KEY_QUERY_VALUE | wow64_override, &key, NULL))
     return false;
 
   if (!QueryRegValueSZ(key, value_name, out_sz)) {
     CloseRegKey(key);
     return false;
   }
 
   CloseRegKey(key);
   return true;
 }
@@ skipping 27 matching lines @@
   // 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,
+                          WOW64_OVERRIDE wow64_override,
                           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))
+  if (!OpenRegKey(root, key_path, KEY_QUERY_VALUE | wow64_override, &key, NULL))
     return false;
 
   if (!QueryRegValueMULTISZ(key, value_name, out_multi_sz)) {
     CloseRegKey(key);
     return false;
   }
 
   CloseRegKey(key);
   return true;
 }
@@ skipping 25 matching lines @@
 }
 
 // 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,
+                      WOW64_OVERRIDE wow64_override,
                       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))
+  if (!OpenRegKey(root, key_path, KEY_SET_VALUE | wow64_override, &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,
+                   WOW64_OVERRIDE wow64_override,
                    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))
+  if (!OpenRegKey(root, key_path, KEY_SET_VALUE | wow64_override, &key, NULL))
     return false;
 
   if (!SetRegValueSZ(key, value_name, value)) {
     CloseRegKey(key);
     return false;
   }
 
   return true;
 }
 
@@ skipping 23 matching lines @@
   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,
+                        WOW64_OVERRIDE wow64_override,
                         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))
+  if (!OpenRegKey(root, key_path, KEY_SET_VALUE | wow64_override, &key, NULL))
     return false;
 
   if (!SetRegValueMULTISZ(key, value_name, values)) {
     CloseRegKey(key);
     return false;
   }
 
   return true;
 }
 
 //------------------------------------------------------------------------------
 // Utils
 //------------------------------------------------------------------------------
 
 const wchar_t* GetCurrentUserSidString() {
   if (!g_initialized)
     InitNativeRegApi();
 
   return g_current_user_sid_string;
 }
 
 };  // namespace nt