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 {
 
+ScopedRegKeyHANDLE::ScopedRegKeyHANDLE() : handle_(INVALID_HANDLE_VALUE) {}
+
+ScopedRegKeyHANDLE::ScopedRegKeyHANDLE(HANDLE handle) : handle_(handle) {}
+
+ScopedRegKeyHANDLE::~ScopedRegKeyHANDLE() {
+  if (!g_initialized)

[penny, 2016/11/21 21:35:46]

Let's first add a check for is_valid().  I'd rather not do any of this if
|handle_| is INVALID_HANDLE_VALUE or NULL.

I definitely don't want to clash with any debug or safety measures related to
calling NtClose on an invalid handle.

[scottmg, 2016/11/22 22:47:36]

Good catch! Done.

+    InitNativeRegApi();
+  g_nt_close(handle_);
+}
+
+HANDLE ScopedRegKeyHANDLE::release() {
+  HANDLE result = handle_;
+  handle_ = INVALID_HANDLE_VALUE;
+  return result;
+}
+
 //------------------------------------------------------------------------------
 // Create, open, delete, close functions
 //------------------------------------------------------------------------------
 
-bool CreateRegKey(ROOT_KEY root,
-                  const wchar_t* key_path,
-                  ACCESS_MASK access,
-                  HANDLE* out_handle OPTIONAL) {
+ScopedRegKeyHANDLE 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 ScopedRegKeyHANDLE();
 
   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 ScopedRegKeyHANDLE();
 
   // 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;
+    return ScopedRegKeyHANDLE();
 
   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]);
@@ skipping 24 matching lines @@
     // 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 ScopedRegKeyHANDLE();
 
-  // See if caller wants the handle left open.
-  if (out_handle)
-    *out_handle = last_handle;
-  else
-    g_nt_close(last_handle);
-
-  return true;
+  return ScopedRegKeyHANDLE(last_handle);
 }
 
-bool OpenRegKey(ROOT_KEY root,
-                const wchar_t* key_path,
-                ACCESS_MASK access,
-                HANDLE* out_handle,
-                NTSTATUS* error_code OPTIONAL) {
+ScopedRegKeyHANDLE OpenRegKey(ROOT_KEY root,
+                              const wchar_t* key_path,
+                              ACCESS_MASK access,
+                              NTSTATUS* error_code OPTIONAL) {
   // |key_path| can be null or empty, but it can't be longer than

[penny, 2016/11/21 21:35:47]

Above this, let's set a default error code (for if we want to bail out before
call to NtOpenKey).

if (error_code)
  *error_code = STATUS_UNSUCCESSFUL;

[scottmg, 2016/11/22 22:47:36]

Done.

   // |g_kRegMaxPathLen| at this point.
   if (key_path != nullptr &&
       ::wcsnlen(key_path, g_kRegMaxPathLen + 1) == g_kRegMaxPathLen + 1)
-    return false;
+    return ScopedRegKeyHANDLE();
 
   if (!g_initialized)
     InitNativeRegApi();
 
   NTSTATUS status = STATUS_UNSUCCESSFUL;

[penny, 2016/11/21 21:35:46]

Delete this initialization.

[scottmg, 2016/11/22 22:47:36]

Done.

   UNICODE_STRING key_path_uni = {};
   OBJECT_ATTRIBUTES obj = {};
-  *out_handle = INVALID_HANDLE_VALUE;
+  HANDLE result_handle = INVALID_HANDLE_VALUE;

[penny, 2016/11/21 21:35:47]

Move initialization of HANDLE result_handle down to just above the call to
g_nt_open_key_ex (line 704).

[scottmg, 2016/11/22 22:47:35]

Done.

 
   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);
+  status = g_nt_open_key_ex(&result_handle, access, &obj, 0);

[penny, 2016/11/21 21:35:47]

NTSTATUS status = ...

[scottmg, 2016/11/22 22:47:36]

Done.

   // See if caller wants the NTSTATUS.
   if (error_code)

[penny, 2016/11/21 21:35:46]

Let's move this down a tad, so that we only set |error_code| on failure.

if (!NT_SUCCESS(status)) {
  // See if caller wants the NTSTATUS.
  if (error_code)
    *error_code = status;
  return ScopedHANDLE();
}

return ScopedHANDLE(result_handle);


The above should replace line 705 to 712.

[scottmg, 2016/11/22 22:47:35]

Done.

     *error_code = status;
 
   if (NT_SUCCESS(status))
-    return true;
+    return ScopedRegKeyHANDLE(result_handle);
 
-  return false;
+  return ScopedRegKeyHANDLE();
 }
 
 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));

[penny, 2016/11/21 21:35:47]

Do my ocd a favour and add an empty line above this return.  :D

[scottmg, 2016/11/22 22:47:36]

Ugly! If you insist.

 }
 
 // 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))
+  ScopedRegKeyHANDLE 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());

[penny, 2016/11/21 21:35:47]

Do my ocd a favour and add an empty line above this return.  :D

[scottmg, 2016/11/22 22:47:35]

Done.

 }
 
 //------------------------------------------------------------------------------
 // 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))
+  ScopedRegKeyHANDLE 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);

[penny, 2016/11/21 21:35:47]

Do my ocd a favour and add an empty line above this return.  :D

[scottmg, 2016/11/22 22:47:35]

Done.

 }
 
 // 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))
+  ScopedRegKeyHANDLE 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);

[penny, 2016/11/21 21:35:47]

Do my ocd a favour and add an empty line above this return.  :D

[scottmg, 2016/11/22 22:47:36]

Done.

 }
 
 // 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))
+  ScopedRegKeyHANDLE 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);

[penny, 2016/11/21 21:35:47]

Do my ocd a favour and add an empty line above this return.  :D

[scottmg, 2016/11/22 22:47:36]

Done.

 }
 
 //------------------------------------------------------------------------------
 // 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))
+  ScopedRegKeyHANDLE 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);

[penny, 2016/11/21 21:35:47]

Do my ocd a favour and add an empty line above this return.  :D

[scottmg, 2016/11/22 22:47:35]

Done.

 }
 
 // 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))
+  ScopedRegKeyHANDLE 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);

[penny, 2016/11/21 21:35:46]

Do my ocd a favour and add an empty line above this return.  :D

[scottmg, 2016/11/22 22:47:36]

Done.

 }
 
 // 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
     // (not allowed in REG_MULTI_SZ), ignore it.
     if (!string.empty()) {
       for (const wchar_t& w : string) {
         builder.push_back(w);
       }
       builder.push_back(L'\0');
     }
   }
   // Add second null terminator to end REG_MULTI_SZ.
   builder.push_back(L'\0');
   // Handle rare case where the vector passed in was empty,
   // or only had an empty string.
   if (builder.size() == 1)
     builder.push_back(L'\0');
 
   if (std::numeric_limits<DWORD>::max() < builder.size())
     return false;
 
   return SetRegKeyValue(
-      key, value_name, REG_MULTI_SZ, reinterpret_cast<BYTE*>(builder.data()),
+      key, value_name, REG_MULTI_SZ,
+      reinterpret_cast<BYTE*>(builder.data()),

[penny, 2016/11/21 21:35:47]

I'm just curious why this got moved down... it wasn't over 80 chars.

(I honestly prefer one argument per line, but that's not our coding standard.)

[scottmg, 2016/11/22 22:47:36]

Forgot to reformat, done.
Agreed. (Crashpad dissented in this way and has a customized clang-format file.
Shhhh!)

       (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))
+  ScopedRegKeyHANDLE 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);

[penny, 2016/11/21 21:35:46]

Do my ocd a favour and add an empty line above this return.  :D

[scottmg, 2016/11/22 22:47:35]

Done.

 }
 
 //------------------------------------------------------------------------------
 // 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