Make nt_registry Create/OpenRegKey return a scoped object

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 {
@@ skipping 547 matching lines @@
                              nullptr);
 
   return g_nt_create_key(out_handle, access, &obj, 0, nullptr,
                          REG_OPTION_NON_VOLATILE, create_or_open);
 }
 
 }  // namespace
 
 namespace nt {
 
+ScopedHANDLE::ScopedHANDLE() : handle_(nullptr) {}
+
+ScopedHANDLE::ScopedHANDLE(HANDLE handle) : handle_(handle) {}
+
+ScopedHANDLE::~ScopedHANDLE() {
+  if (!g_initialized)
+    InitNativeRegApi();
+  if (is_valid())
+    g_nt_close(handle_);
+}
+
+bool ScopedHANDLE::is_valid() const {
+  // This class should only be used for HANDLEs that have an invalid value of
+  // nullptr.
+  assert(handle_ != INVALID_HANDLE_VALUE);
+
+  return handle_ != nullptr;
+}
+
+HANDLE ScopedHANDLE::release() {
+  HANDLE result = handle_;
+  handle_ = nullptr;
+  return result;
+}
+
 //------------------------------------------------------------------------------
 // Create, open, delete, close functions
 //------------------------------------------------------------------------------
 
-bool CreateRegKey(ROOT_KEY root,
-                  const wchar_t* key_path,
-                  ACCESS_MASK access,
-                  HANDLE* out_handle OPTIONAL) {
+ScopedHANDLE CreateRegKey(ROOT_KEY root,
+                          const wchar_t* key_path,
+                          ACCESS_MASK access) {
   // |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;
+    return ScopedHANDLE();
 
   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), redirected_key_path,
                         &current_path, &subkeys))
-    return false;
+    return ScopedHANDLE();
 
   // Open the base hive first.  It should always exist already.
-  HANDLE last_handle = INVALID_HANDLE_VALUE;
+  HANDLE last_handle = nullptr;
   NTSTATUS status =
       CreateKeyWrapper(current_path, access, &last_handle, nullptr);
   if (!NT_SUCCESS(status))
-    return false;
+    return ScopedHANDLE();
 
   size_t subkeys_size = subkeys.size();
   if (subkeys_size != 0)
     g_nt_close(last_handle);
 
   // Recursively open/create each subkey.
   std::vector<HANDLE> rollback;
   bool success = true;
   for (size_t i = 0; i < subkeys_size; i++) {
     current_path.append(subkeys[i]);
     current_path.push_back(L'\\');
 
     // Process the latest subkey.
     ULONG created = 0;
-    HANDLE key_handle = INVALID_HANDLE_VALUE;
+    HANDLE key_handle = nullptr;
     status =
         CreateKeyWrapper(current_path.c_str(), access, &key_handle, &created);
     if (!NT_SUCCESS(status)) {
       success = false;
       break;
     }
 
     if (i == subkeys_size - 1) {
       last_handle = key_handle;
     } else {
       // Save any subkey handle created, in case of rollback.
       if (created == REG_CREATED_NEW_KEY)
         rollback.push_back(key_handle);
       else
         g_nt_close(key_handle);
     }
   }
 
   if (!success) {
     // Delete any subkeys created.
     for (HANDLE handle : rollback) {
       g_nt_delete_key(handle);
     }
   }
   for (HANDLE handle : rollback) {
     // Close the rollback handles, on success or failure.
     g_nt_close(handle);
   }
   if (!success)
-    return false;
+    return ScopedHANDLE();
 
-  // See if caller wants the handle left open.
-  if (out_handle)
-    *out_handle = last_handle;
-  else
-    g_nt_close(last_handle);
-
-  return true;
+  return ScopedHANDLE(last_handle);
 }
 
-bool OpenRegKey(ROOT_KEY root,
-                const wchar_t* key_path,
-                ACCESS_MASK access,
-                HANDLE* out_handle,
-                NTSTATUS* error_code OPTIONAL) {
+ScopedHANDLE OpenRegKey(ROOT_KEY root,
+                        const wchar_t* key_path,
+                        ACCESS_MASK access,
+                        NTSTATUS* error_code OPTIONAL) {
+  if (error_code)
+    *error_code = STATUS_UNSUCCESSFUL;
+
   // |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;
+    return ScopedHANDLE();
 
   if (!g_initialized)
     InitNativeRegApi();
 
-  NTSTATUS status = STATUS_UNSUCCESSFUL;
   UNICODE_STRING key_path_uni = {};
   OBJECT_ATTRIBUTES obj = {};
-  *out_handle = INVALID_HANDLE_VALUE;
 
   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;
+  HANDLE result_handle = nullptr;
+  NTSTATUS status = g_nt_open_key_ex(&result_handle, access, &obj, 0);
 
-  if (NT_SUCCESS(status))
-    return true;
+  if (!NT_SUCCESS(status)) {
+    // See if caller wants the NTSTATUS.
+    if (error_code)
+      *error_code = status;
+    return ScopedHANDLE();
+  }
 
-  return false;
+  return ScopedHANDLE(result_handle);
 }
 
 bool DeleteRegKey(HANDLE key) {
   if (!g_initialized)
     InitNativeRegApi();
 
-  NTSTATUS status = STATUS_UNSUCCESSFUL;
-
-  status = g_nt_delete_key(key);
-
-  if (NT_SUCCESS(status))
-    return true;
-
-  return false;
+  return NT_SUCCESS(g_nt_delete_key(key));
 }
 
 // wrapper function
 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 | wow64_override, &key, nullptr))
+  ScopedHANDLE key =
+      OpenRegKey(root, key_path, DELETE | wow64_override, nullptr);
+  if (!key.is_valid())
     return false;
 
-  if (!DeleteRegKey(key)) {
-    CloseRegKey(key);
-    return false;
-  }
-
-  CloseRegKey(key);
-  return true;
-}
-
-void CloseRegKey(HANDLE key) {
-  if (!g_initialized)
-    InitNativeRegApi();
-  g_nt_close(key);
+  return DeleteRegKey(key.get());
 }
 
 //------------------------------------------------------------------------------
 // Getter functions
 //------------------------------------------------------------------------------
 
 bool QueryRegKeyValue(HANDLE key,
                       const wchar_t* value_name,
                       ULONG* out_type,
                       BYTE** out_buffer,
@@ skipping 50 matching lines @@
   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 | wow64_override, &key, NULL))
+  ScopedHANDLE key =
+      OpenRegKey(root, key_path, KEY_QUERY_VALUE | wow64_override, nullptr);
+  if (!key.is_valid())
     return false;
 
-  if (!QueryRegValueDWORD(key, value_name, out_dword)) {
-    CloseRegKey(key);
-    return false;
-  }
-
-  CloseRegKey(key);
-  return true;
+  return QueryRegValueDWORD(key.get(), value_name, out_dword);
 }
 
 // wrapper function
 bool QueryRegValueSZ(HANDLE key,
                      const wchar_t* value_name,
                      std::wstring* out_sz) {
   BYTE* value_bytes = nullptr;
   DWORD ret_size = 0;
   ULONG type = REG_NONE;
 
   if (!QueryRegKeyValue(key, value_name, &type, &value_bytes, &ret_size) ||
-      type != REG_SZ)
+      type != REG_SZ) {
     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 | wow64_override, &key, NULL))
+  ScopedHANDLE key =
+      OpenRegKey(root, key_path, KEY_QUERY_VALUE | wow64_override, nullptr);
+  if (!key.is_valid())
     return false;
 
-  if (!QueryRegValueSZ(key, value_name, out_sz)) {
-    CloseRegKey(key);
-    return false;
-  }
-
-  CloseRegKey(key);
-  return true;
+  return QueryRegValueSZ(key.get(), value_name, out_sz);
 }
 
 // wrapper function
 bool QueryRegValueMULTISZ(HANDLE key,
                           const wchar_t* value_name,
                           std::vector<std::wstring>* out_multi_sz) {
   BYTE* value_bytes = nullptr;
   DWORD ret_size = 0;
   ULONG type = REG_NONE;
 
@@ skipping 22 matching lines @@
   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 | wow64_override, &key, NULL))
+  ScopedHANDLE key =
+      OpenRegKey(root, key_path, KEY_QUERY_VALUE | wow64_override, nullptr);
+  if (!key.is_valid())
     return false;
 
-  if (!QueryRegValueMULTISZ(key, value_name, out_multi_sz)) {
-    CloseRegKey(key);
-    return false;
-  }
-
-  CloseRegKey(key);
-  return true;
+  return QueryRegValueMULTISZ(key.get(), value_name, out_multi_sz);
 }
 
 //------------------------------------------------------------------------------
 // Setter functions
 //------------------------------------------------------------------------------
 
 bool SetRegKeyValue(HANDLE key,
                     const wchar_t* value_name,
                     ULONG type,
                     const BYTE* data,
@@ skipping 20 matching lines @@
   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 | wow64_override, &key, NULL))
+  ScopedHANDLE key =
+      OpenRegKey(root, key_path, KEY_SET_VALUE | wow64_override, nullptr);
+  if (!key.is_valid())
     return false;
 
-  if (!SetRegValueDWORD(key, value_name, value)) {
-    CloseRegKey(key);
-    return false;
-  }
-
-  return true;
+  return SetRegValueDWORD(key.get(), value_name, value);
 }
 
 // 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 | wow64_override, &key, NULL))
+  ScopedHANDLE key =
+      OpenRegKey(root, key_path, KEY_SET_VALUE | wow64_override, nullptr);
+  if (!key.is_valid())
     return false;
 
-  if (!SetRegValueSZ(key, value_name, value)) {
-    CloseRegKey(key);
-    return false;
-  }
-
-  return true;
+  return SetRegValueSZ(key.get(), value_name, value);
 }
 
 // wrapper function
 bool SetRegValueMULTISZ(HANDLE key,
                         const wchar_t* value_name,
                         const std::vector<std::wstring>& values) {
   std::vector<wchar_t> builder;
 
   for (auto& string : values) {
     // Just in case someone is passing in an illegal empty string
@@ skipping 19 matching lines @@
       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 | wow64_override, &key, NULL))
+  ScopedHANDLE key =
+      OpenRegKey(root, key_path, KEY_SET_VALUE | wow64_override, nullptr);
+  if (!key.is_valid())
     return false;
 
-  if (!SetRegValueMULTISZ(key, value_name, values)) {
-    CloseRegKey(key);
-    return false;
-  }
-
-  return true;
+  return SetRegValueMULTISZ(key.get(), value_name, values);
 }
 
 //------------------------------------------------------------------------------
 // Utils
 //------------------------------------------------------------------------------
 
 const wchar_t* GetCurrentUserSidString() {
   if (!g_initialized)
     InitNativeRegApi();
 
@@ skipping 28 matching lines @@
   if (!g_initialized)
     InitNativeRegApi();
 
   if (root == HKCU || (root == AUTO && !g_system_install))
     return g_HKCU_override;
 
   return g_HKLM_override;
 }
 
 };  // namespace nt