[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_reg_redirection = false;
-const size_t g_kMaxPathLen = 255;
+bool g_wow64_proc = false;
 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"";
+wchar_t g_kRegPathHKCU[nt::g_kRegMaxPathLen + 1] = L"";
+wchar_t g_current_user_sid_string[nt::g_kRegMaxPathLen + 1] = L"";
 
-// Not using install_util, to prevent circular dependency.
+// For testing only.
+wchar_t g_HKLM_override[nt::g_kRegMaxPathLen + 1] = L"";
+wchar_t g_HKCU_override[nt::g_kRegMaxPathLen + 1] = L"";
+
+//------------------------------------------------------------------------------
+// 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() {
+  // 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 27 matching lines @@
           ::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_kMaxPathLen - 1));
+  ::wcsncat(g_kRegPathHKCU, current_user_reg_path.Buffer, nt::g_kRegMaxPathLen);
   ::wcsncat(g_kRegPathHKCU, L"\\",
-            (g_kMaxPathLen - ::wcslen(g_kRegPathHKCU) - 1));
+            (nt::g_kRegMaxPathLen - ::wcslen(g_kRegPathHKCU)));
   // 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));
+  ::wcsncpy(g_current_user_sid_string, ptr, nt::g_kRegMaxPathLen);
   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;
 }
 
-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) && (::wcslen(nt::HKCU_override) != 0)) {
-    std::wstring temp(g_kRegPathHKCU);
-    temp.append(nt::HKCU_override);
+//------------------------------------------------------------------------------
+// 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 decision trees represent the information at the MSDN
+//   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 kClassesSubtree[] = {{L"CLSID", REDIRECTED_BEFORE},
+                                    {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 kHklmSoftwareSubtree[] = {
+    // TODO(pennymac): when MS fixes compiler bug, or bots are all using clang,
+    // remove the "Classes" subkeys below and replace with:
+    // {L"Classes", SHARED, kClassesSubtree},
+    // 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 kRedirectionDecisionTreeHkcu = {L"SOFTWARE\\Classes", SHARED,
+                                               kClassesSubtree};
+
+// HKLM\SOFTWARE is redirected by default to SOFTWARE\WOW6432Node.  Specific
+// subkeys under SOFTWARE are shared though... see |hklm_software_subtree|.
+constexpr Node kRedirectionDecisionTreeHklm = {L"SOFTWARE", REDIRECTED_AFTER,
+                                               kHklmSoftwareSubtree};
+
+// 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
+//   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";
+
+  assert(subkey_path != nullptr);
+  assert(subkey_path->empty() || subkey_path->front() != L'\\');
+  assert(subkey_path->empty() || subkey_path->back() != L'\\');
+  assert(root != nt::AUTO);
+
+  // |subkey_path| could legitimately be empty.
+  if (subkey_path->empty() ||
+      (access & KEY_WOW64_32KEY && access & KEY_WOW64_64KEY))
+    return;
+
+  // 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 (root == nt::HKCU ? *g_HKCU_override : *g_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;
+
+  // The root of the decision trees are an array of 1.
+  size_t node_array_len = 1;
+  // Pick which decision tree to use.
+  const Node* current_node = (root == nt::HKCU) ? &kRedirectionDecisionTreeHkcu
+                                                : &kRedirectionDecisionTreeHklm;
+
+  // 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;
+
+  size_t node_index = 0;
+  while (node_index < node_array_len) {
+    size_t current_to_match_len = current_node->to_match_len;
+    // Make sure the remainder of the path is at least as long as the current
+    // subkey to match.
+    if (chars_left >= current_to_match_len) {
+      // Do case insensitive comparisons.
+      if (!::wcsnicmp(position, current_node->to_match, current_to_match_len)) {
+        // Make sure not to match on a substring.
+        if (*(position + current_to_match_len) == L'\\' ||
+            *(position + current_to_match_len) == L'\0') {
+          // MATCH!
+          // -------------------------------------------------------------------
+          // 1) Update state of redirection.
+          redirect_state = current_node->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 + current_to_match_len;
+            insert_string = kRedirectAfter;
+          }
+          // 2) Adjust |position| along the subkey path.
+          position += current_to_match_len;
+          chars_left -= current_to_match_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.
+          node_array_len = current_node->next_len;
+          current_node = current_node->next;
+          node_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_node;
+    ++node_index;
+  }
+
+  if (redirect_state == SHARED)
+    return;
+
+  // Insert the redirection into |subkey_path|, at |insertion_point|.
+  subkey_path->insert((insertion_point - subkey_path->c_str()), insert_string);
+}
+
+//------------------------------------------------------------------------------
+// Reg Path Utilities - LOCAL
+//------------------------------------------------------------------------------
+
+std::wstring ConvertRootKey(nt::ROOT_KEY root) {
+  assert(root != nt::AUTO);
+
+  if (root == nt::HKCU && *g_HKCU_override) {
+    std::wstring temp = g_kRegPathHKCU;
+    temp.append(g_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)) {
-    std::wstring temp(g_kRegPathHKCU);
-    temp.append(nt::HKLM_override);
+    return temp;
+  } else if (root == nt::HKLM && *g_HKLM_override) {
+    // Yes, HKLM override goes into HKCU.  This is not a typo.
+    std::wstring temp = g_kRegPathHKCU;
+    temp.append(g_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;
+    return temp;
+  }
+
+  return std::wstring((root == nt::HKCU) ? g_kRegPathHKCU : g_kRegPathHKLM);
+}
+
+// This utility should be called on an externally provided subkey path.
+// - Ensures there are no starting or trailing backslashes, and no more than
+// - one backslash in a row.
+// - Note from MSDN: "Key names cannot include the backslash character (\),
+//   but any other printable character can be used.  Value names and data can
+//   include the backslash character."
+void SanitizeSubkeyPath(std::wstring* input) {
+  assert(input != nullptr);
+
+  // Remove trailing backslashes.
+  size_t last_valid_pos = input->find_last_not_of(L'\\');
+  if (last_valid_pos == std::wstring::npos) {
+    // The string is all backslashes, or it's empty.  Clear and abort.
+    input->clear();
+    return;
+  }
+  // Chop off the trailing backslashes.
+  input->resize(last_valid_pos + 1);
+
+  // Remove leading backslashes.
+  input->erase(0, input->find_first_not_of(L'\\'));
+
+  // Replace any occurances of more than 1 backslash in a row with just 1.
+  size_t index = input->find_first_of(L"\\");
+  while (index != std::wstring::npos) {
+    // Remove a second consecutive backslash, and leave index where it is,
+    // or move to the next backslash in the string.
+    if ((*input)[index + 1] == L'\\')
+      input->erase(index + 1, 1);
+    else
+      index = input->find_first_of(L"\\", index + 1);
+  }
 }
 
 // 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,
+bool ParseFullRegPath(const std::wstring& converted_root,
+                      const std::wstring& subkey_path,
                       std::wstring* out_base,
                       std::vector<std::wstring>* subkeys) {
   out_base->clear();
   subkeys->clear();
-  std::wstring temp = L"";
-
-  if (g_reg_redirection) {
+  std::wstring temp_path;
+
+  // Special case if there is testing redirection set up.
+  if (*g_HKCU_override || *g_HKLM_override) {
     // Why process |converted_root|?  To handle reg redirection used by tests.
     // E.g.:
     // |converted_root| = "\REGISTRY\USER\S-1-5-21-39260824-743453154-142223018-
     // 716772\Software\Chromium\TempTestKeys\13110669370890870$94c6ed9d-bc34-
     // 44f3-a0b3-9eee2d3f2f82\".
     // |subkey_path| = "SOFTWARE\Google\Chrome\BrowserSec".
-    temp.append(converted_root);
-  }
-  if (subkey_path != nullptr)
-    temp.append(subkey_path);
+    //
+    // Note: bypassing the starting backslash in the |converted_root|.
+    temp_path.append(converted_root, 1, converted_root.size() - 1);
+  }
+  temp_path.append(subkey_path);
 
   // Tokenize the full path.
   size_t find_start = 0;
-  size_t delimiter = temp.find_first_of(L'\\');
+  size_t delimiter = temp_path.find_first_of(L'\\');
   while (delimiter != std::wstring::npos) {
-    std::wstring token = temp.substr(find_start, delimiter - find_start);
-    if (!token.empty())
-      subkeys->push_back(token);
+    subkeys->emplace_back(temp_path, find_start, delimiter - find_start);
+    // Move past the backslash.
     find_start = delimiter + 1;
-    delimiter = temp.find_first_of(L'\\', find_start);
-  }
-  if (!temp.empty() && find_start < temp.length())
-    // Get the last token.
-    subkeys->push_back(temp.substr(find_start));
-
-  if (g_reg_redirection) {
+    delimiter = temp_path.find_first_of(L'\\', find_start);
+  }
+  // Get the last token if there is one.
+  if (!temp_path.empty())
+    subkeys->emplace_back(temp_path, find_start);
+
+  // Special case if there is testing redirection set up.
+  if (*g_HKCU_override || *g_HKLM_override) {
     // The base hive for HKCU needs to include the user SID.
     uint32_t num_base_tokens = 2;
-    const wchar_t* hkcu = L"\\REGISTRY\\USER\\";
-    if (0 == ::wcsnicmp(converted_root, hkcu, ::wcslen(hkcu)))
+    if (0 == temp_path.compare(0, 14, L"REGISTRY\\USER\\"))
       num_base_tokens = 3;
 
     if (subkeys->size() < num_base_tokens)
       return false;
 
     // Pull out the base hive tokens.
     out_base->push_back(L'\\');
-    for (size_t i = 0; i < num_base_tokens; i++) {
-      out_base->append((*subkeys)[i].c_str());
+    for (size_t i = 0; i < num_base_tokens; ++i) {
+      out_base->append((*subkeys)[i]);
       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"";
-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) {
+  // |key_path| can be null or empty, but it can't be longer than
+  // |g_kRegMaxPathLen| at this point.
+  if (key_path != nullptr &&
+      ::wcsnlen(key_path, g_kRegMaxPathLen + 1) == g_kRegMaxPathLen + 1)
+    return false;
+
   if (!g_initialized)
     InitNativeRegApi();
 
+  if (root == nt::AUTO)
+    root = g_system_install ? nt::HKLM : nt::HKCU;
+
+  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,
+                        &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 48 matching lines @@
     g_nt_close(last_handle);
 
   return true;
 }
 
 bool OpenRegKey(ROOT_KEY root,
                 const wchar_t* key_path,
                 ACCESS_MASK access,
                 HANDLE* out_handle,
                 NTSTATUS* error_code OPTIONAL) {
+  // |key_path| can be null or empty, but it can't be longer than
+  // |g_kRegMaxPathLen| at this point.
+  if (key_path != nullptr &&
+      ::wcsnlen(key_path, g_kRegMaxPathLen + 1) == g_kRegMaxPathLen + 1)
+    return false;
+
   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);
+  if (root == nt::AUTO)
+    root = g_system_install ? nt::HKLM : nt::HKCU;
+
+  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;
 }
 
+bool IsCurrentProcWow64() {
+  if (!g_initialized)
+    InitNativeRegApi();
+
+  return g_wow64_proc;
+}
+
+bool SetTestingOverride(ROOT_KEY root, const std::wstring& new_path) {
+  if (!g_initialized)
+    InitNativeRegApi();
+
+  std::wstring sani_new_path = new_path;
+  SanitizeSubkeyPath(&sani_new_path);
+  if (sani_new_path.length() > g_kRegMaxPathLen)
+    return false;
+
+  if (root == HKCU || (root == AUTO && !g_system_install))
+    ::wcsncpy(g_HKCU_override, sani_new_path.c_str(), nt::g_kRegMaxPathLen);
+  else
+    ::wcsncpy(g_HKLM_override, sani_new_path.c_str(), nt::g_kRegMaxPathLen);
+
+  return true;
+}
+
+std::wstring GetTestingOverride(ROOT_KEY root) {
+  if (!g_initialized)
+    InitNativeRegApi();
+
+  if (root == HKCU || (root == AUTO && !g_system_install))
+    return std::wstring(g_HKCU_override);
+  else

[grt (UTC plus 2), 2016/10/03 20:17:45]

nit: no "else" after "return" in an "if" body

[penny, 2016/10/03 20:46:03]

oh, right.

+    return std::wstring(g_HKLM_override);
+}
+
 };  // namespace nt