Make nt_registry Create/OpenRegKey return a scoped object

chrome_elf/blacklist/blacklist.cc

 // Copyright 2013 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/blacklist/blacklist.h"
 
 #include <assert.h>
 #include <string.h>
 
 #include <vector>
@@ skipping 83 matching lines @@
 namespace blacklist {
 
 #if defined(_WIN64)
 // Allocate storage for the pointer to the old NtMapViewOfSectionFunction.
 #pragma section(".oldntmap", write, read)
 __declspec(allocate(".oldntmap"))
     NtMapViewOfSectionFunction g_nt_map_view_of_section_func = NULL;
 #endif
 
 bool LeaveSetupBeacon() {
-  HANDLE key_handle = INVALID_HANDLE_VALUE;
+  nt::ScopedHANDLE key_handle = nt::CreateRegKey(
+      nt::HKCU, kRegistryBeaconPath, KEY_QUERY_VALUE | KEY_SET_VALUE);
 
-  if (!nt::CreateRegKey(nt::HKCU, kRegistryBeaconPath,
-                        KEY_QUERY_VALUE | KEY_SET_VALUE, &key_handle))
+  if (!key_handle.is_valid())
     return false;
 
   DWORD blacklist_state = BLACKLIST_STATE_MAX;
-  if (!nt::QueryRegValueDWORD(key_handle, kBeaconState, &blacklist_state) ||
+  if (!nt::QueryRegValueDWORD(key_handle.get(), kBeaconState,
+                              &blacklist_state) ||
       blacklist_state == BLACKLIST_DISABLED) {
-    nt::CloseRegKey(key_handle);
     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::SetRegValueDWORD(key_handle, kBeaconAttemptCount,
+    if (nt::SetRegValueDWORD(key_handle.get(), kBeaconAttemptCount,
                              static_cast<DWORD>(0))) {
-      if (nt::SetRegValueDWORD(key_handle, kBeaconState,
+      if (nt::SetRegValueDWORD(key_handle.get(), 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;
 
-    nt::QueryRegValueDWORD(key_handle, blacklist::kBeaconAttemptCount,
+    nt::QueryRegValueDWORD(key_handle.get(), blacklist::kBeaconAttemptCount,
                            &attempt_count);
     ++attempt_count;
-    nt::SetRegValueDWORD(key_handle, blacklist::kBeaconAttemptCount,
+    nt::SetRegValueDWORD(key_handle.get(), blacklist::kBeaconAttemptCount,
                          attempt_count);
 
     if (attempt_count >= blacklist::kBeaconMaxAttempts) {
       blacklist_state = blacklist::BLACKLIST_SETUP_FAILED;
-      nt::SetRegValueDWORD(key_handle, blacklist::kBeaconState,
+      nt::SetRegValueDWORD(key_handle.get(), blacklist::kBeaconState,
                            blacklist_state);
     }
   }
 
-  nt::CloseRegKey(key_handle);
   return success;
 }
 
 bool ResetBeacon() {
-  HANDLE key_handle = INVALID_HANDLE_VALUE;
-
-  if (!nt::CreateRegKey(nt::HKCU, kRegistryBeaconPath,
-                        KEY_QUERY_VALUE | KEY_SET_VALUE, &key_handle))
+  nt::ScopedHANDLE key_handle = nt::CreateRegKey(
+      nt::HKCU, kRegistryBeaconPath, KEY_QUERY_VALUE | KEY_SET_VALUE);
+  if (!key_handle.is_valid())
     return false;
 
   DWORD blacklist_state = BLACKLIST_STATE_MAX;
-  if (!nt::QueryRegValueDWORD(key_handle, kBeaconState, &blacklist_state)) {
-    nt::CloseRegKey(key_handle);
+  if (!nt::QueryRegValueDWORD(key_handle.get(), kBeaconState, &blacklist_state))
     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) {
-    if (!nt::SetRegValueDWORD(key_handle, kBeaconState, BLACKLIST_ENABLED)) {
-      nt::CloseRegKey(key_handle);
+    if (!nt::SetRegValueDWORD(key_handle.get(), kBeaconState,
+                              BLACKLIST_ENABLED)) {
       return false;
     }
   }
 
-  nt::CloseRegKey(key_handle);
   return true;
 }
 
 int BlacklistSize() {
   int size = -1;
   while (blacklist::g_troublesome_dlls[++size] != NULL) {
   }
 
   return size;
 }
@@ skipping 158 matching lines @@
 
   // Mark the thunk storage as executable and prevent any future writes to it.
   page_executable = page_executable &&
                     VirtualProtect(&g_thunk_storage, sizeof(g_thunk_storage),
                                    PAGE_EXECUTE_READ, &old_protect);
 
   return NT_SUCCESS(ret) && page_executable;
 }
 
 }  // namespace blacklist