[Chrome ELF] New NT registry API.

chrome_elf/blacklist/blacklist.cc

Index: chrome_elf/blacklist/blacklist.cc
diff --git a/chrome_elf/blacklist/blacklist.cc b/chrome_elf/blacklist/blacklist.cc
index 5cdf64562836fccec2d4a545df620914075c24e3..932374600d4a31b7d842eba67df4d7c68b36bcf1 100644
--- a/chrome_elf/blacklist/blacklist.cc
+++ b/chrome_elf/blacklist/blacklist.cc
@@ -11,6 +11,7 @@
 
 #include "chrome_elf/blacklist/blacklist_interceptions.h"
 #include "chrome_elf/chrome_elf_constants.h"
+#include "chrome_elf/chrome_elf_reg.h"
 #include "chrome_elf/chrome_elf_util.h"
 #include "chrome_elf/thunk_getter.h"
 #include "sandbox/win/src/interception_internal.h"
@@ -88,54 +89,6 @@ namespace {
 // determine if the blacklist is enabled for them.
 bool g_blacklist_initialized = false;
 
-// Helper to set DWORD registry values.
-DWORD SetDWValue(HKEY* key, const wchar_t* property, DWORD value) {
-  return ::RegSetValueEx(*key,
-                         property,
-                         0,
-                         REG_DWORD,
-                         reinterpret_cast<LPBYTE>(&value),
-                         sizeof(value));
-}
-
-bool GenerateStateFromBeaconAndAttemptCount(HKEY* key, DWORD blacklist_state) {
-  LONG result = 0;
-  if (blacklist_state == blacklist::BLACKLIST_ENABLED) {
-    // If the blacklist succeeded on the previous run reset the failure
-    // counter.
-    return (SetDWValue(key,
-                       blacklist::kBeaconAttemptCount,
-                       static_cast<DWORD>(0)) == ERROR_SUCCESS);
-  } else {
-    // Some part of the blacklist setup failed last time.  If this has occured
-    // blacklist::kBeaconMaxAttempts times in a row we switch the state to
-    // failed and skip setting up the blacklist.
-    DWORD attempt_count = 0;
-    DWORD attempt_count_size = sizeof(attempt_count);
-    result = ::RegQueryValueEx(*key,
-                               blacklist::kBeaconAttemptCount,
-                               0,
-                               NULL,
-                               reinterpret_cast<LPBYTE>(&attempt_count),
-                               &attempt_count_size);
-
-    if (result == ERROR_FILE_NOT_FOUND)
-      attempt_count = 0;
-    else if (result != ERROR_SUCCESS)
-      return false;
-
-    ++attempt_count;
-    SetDWValue(key, blacklist::kBeaconAttemptCount, attempt_count);
-
-    if (attempt_count >= blacklist::kBeaconMaxAttempts) {
-      blacklist_state = blacklist::BLACKLIST_SETUP_FAILED;
-      SetDWValue(key, blacklist::kBeaconState, blacklist_state);
-    }
-
-    return false;
-  }
-}
-
 }  // namespace
 
 namespace blacklist {
@@ -148,86 +101,79 @@ namespace blacklist {
 #endif
 
 bool LeaveSetupBeacon() {
-  HKEY key = NULL;
-  DWORD disposition = 0;
-  LONG result = ::RegCreateKeyEx(HKEY_CURRENT_USER,
-                                 kRegistryBeaconPath,
-                                 0,
-                                 NULL,
-                                 REG_OPTION_NON_VOLATILE,
-                                 KEY_QUERY_VALUE | KEY_SET_VALUE,
-                                 NULL,
-                                 &key,
-                                 &disposition);
-  if (result != ERROR_SUCCESS)
+  HANDLE key_handle = INVALID_HANDLE_VALUE;
+
+  if (!nt::CreateRegKey(nt::HKCU, kRegistryBeaconPath,
+                        KEY_QUERY_VALUE | KEY_SET_VALUE, &key_handle))
     return false;
 
-  // Retrieve the current blacklist state.
   DWORD blacklist_state = BLACKLIST_STATE_MAX;
-  DWORD blacklist_state_size = sizeof(blacklist_state);
-  DWORD type = 0;
-  result = ::RegQueryValueEx(key,
-                             kBeaconState,
-                             0,
-                             &type,
-                             reinterpret_cast<LPBYTE>(&blacklist_state),
-                             &blacklist_state_size);
-
-  if (result != ERROR_SUCCESS || blacklist_state == BLACKLIST_DISABLED ||
-      type != REG_DWORD) {
-    ::RegCloseKey(key);
+  if (!nt::GetRegValue_DWORD(key_handle, kBeaconState, &blacklist_state) ||
+      blacklist_state == BLACKLIST_DISABLED) {
+    nt::CloseRegKey(key_handle);
     return false;
   }
 
-  if (!GenerateStateFromBeaconAndAttemptCount(&key, blacklist_state)) {
-    ::RegCloseKey(key);
-    return false;
-  }
+  // Handle attempt count.
+  // Only return true if BL is enabled and succeeded on previous run.
+  bool success = false;
+  if (blacklist_state == BLACKLIST_ENABLED) {
+    // If the blacklist succeeded on the previous run reset the failure
+    // counter.  Then update the beacon state.
+    if (nt::SetRegValue_DWORD(key_handle, kBeaconAttemptCount,
+                              static_cast<DWORD>(0))) {
+      if (nt::SetRegValue_DWORD(key_handle, kBeaconState,
+                                BLACKLIST_SETUP_RUNNING))
+        success = true;
+    }
+  } else {
+    // Some part of the blacklist setup failed last time.  If this has occured
+    // blacklist::kBeaconMaxAttempts times in a row we switch the state to
+    // failed and skip setting up the blacklist.
+    DWORD attempt_count = 0;
 
-  result = SetDWValue(&key, kBeaconState, BLACKLIST_SETUP_RUNNING);
-  ::RegCloseKey(key);
+    nt::GetRegValue_DWORD(key_handle, blacklist::kBeaconAttemptCount,
+                          &attempt_count);
+    ++attempt_count;
+    nt::SetRegValue_DWORD(key_handle, blacklist::kBeaconAttemptCount,
+                          attempt_count);
 
-  return (result == ERROR_SUCCESS);
+    if (attempt_count >= blacklist::kBeaconMaxAttempts) {
+      blacklist_state = blacklist::BLACKLIST_SETUP_FAILED;
+      nt::SetRegValue_DWORD(key_handle, blacklist::kBeaconState,
+                            blacklist_state);
+    }
+  }
+
+  nt::CloseRegKey(key_handle);
+  return success;
 }
 
 bool ResetBeacon() {
-  HKEY key = NULL;
-  DWORD disposition = 0;
-  LONG result = ::RegCreateKeyEx(HKEY_CURRENT_USER,
-                                 kRegistryBeaconPath,
-                                 0,
-                                 NULL,
-                                 REG_OPTION_NON_VOLATILE,
-                                 KEY_QUERY_VALUE | KEY_SET_VALUE,
-                                 NULL,
-                                 &key,
-                                 &disposition);
-  if (result != ERROR_SUCCESS)
+  HANDLE key_handle = INVALID_HANDLE_VALUE;
+
+  if (!nt::CreateRegKey(nt::HKCU, kRegistryBeaconPath,
+                        KEY_QUERY_VALUE | KEY_SET_VALUE, &key_handle))
     return false;
 
   DWORD blacklist_state = BLACKLIST_STATE_MAX;
-  DWORD blacklist_state_size = sizeof(blacklist_state);
-  DWORD type = 0;
-  result = ::RegQueryValueEx(key,
-                             kBeaconState,
-                             0,
-                             &type,
-                             reinterpret_cast<LPBYTE>(&blacklist_state),
-                             &blacklist_state_size);
-
-  if (result != ERROR_SUCCESS || type != REG_DWORD) {
-    ::RegCloseKey(key);
+  if (!nt::GetRegValue_DWORD(key_handle, kBeaconState, &blacklist_state)) {
+    nt::CloseRegKey(key_handle);
     return false;
   }
 
   // Reaching this point with the setup running state means the setup did not
   // crash, so we reset to enabled.  Any other state indicates that setup was
   // skipped; in that case we leave the state alone for later recording.
-  if (blacklist_state == BLACKLIST_SETUP_RUNNING)
-    result = SetDWValue(&key, kBeaconState, BLACKLIST_ENABLED);
+  if (blacklist_state == BLACKLIST_SETUP_RUNNING) {
+    if (!nt::SetRegValue_DWORD(key_handle, kBeaconState, BLACKLIST_ENABLED)) {
+      nt::CloseRegKey(key_handle);
+      return false;
+    }
+  }
 
-  ::RegCloseKey(key);
-  return (result == ERROR_SUCCESS);
+  nt::CloseRegKey(key_handle);
+  return true;
 }
 
 int BlacklistSize() {
@@ -416,41 +362,18 @@ bool Initialize(bool force) {
 }
 
 void AddDllsFromRegistryToBlacklist() {
-  HKEY key = NULL;
-  LONG result = ::RegOpenKeyEx(HKEY_CURRENT_USER,
-                               kRegistryFinchListPath,
-                               0,
-                               KEY_QUERY_VALUE | KEY_SET_VALUE,
-                               &key);
-
-  if (result != ERROR_SUCCESS)
+  std::vector<base::string16> dlls;
+
+  if (!nt::GetRegValue_MULTI_SZ(nt::HKCU, kRegistryFinchListPath,
+                                kRegistryFinchListValueName, &dlls) ||
+      dlls.size() == 0)

[robertshield, 2016/04/20 05:16:09]

dlls.empty()

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

Done.

     return;
 
-  // We add dlls from the registry to the blacklist.
-  DWORD value_len;
-  DWORD name_len = MAX_PATH;
-  std::vector<wchar_t> name_buffer(name_len);
-  for (int i = 0; result == ERROR_SUCCESS; ++i) {
-    name_len = MAX_PATH;
-    value_len = 0;
-    result = ::RegEnumValue(
-        key, i, &name_buffer[0], &name_len, NULL, NULL, NULL, &value_len);
-    if (result != ERROR_SUCCESS)
-      break;
-
-    name_len = name_len + 1;
-    value_len = value_len + 1;
-    std::vector<wchar_t> value_buffer(value_len);
-    result = ::RegEnumValue(key, i, &name_buffer[0], &name_len, NULL, NULL,
-                            reinterpret_cast<BYTE*>(&value_buffer[0]),
-                            &value_len);
-    if (result != ERROR_SUCCESS)
-      break;
-    value_buffer[value_len - 1] = L'\0';
-    AddDllToBlacklist(&value_buffer[0]);
+  // Add each DLL to the BL in memory
+  for (auto name : dlls) {
+    AddDllToBlacklist(name.c_str());
   }
 
-  ::RegCloseKey(key);
   return;
 }