[vcpkg] Refactor RAII registry key manipulation

3 files changed, 35 insertions, 30 deletions

diff --git a/toolsrc/include/vcpkg_System.h b/toolsrc/include/vcpkg_System.h
index 14f39ca40..ecb98aba7 100644
--- a/toolsrc/include/vcpkg_System.h
+++ b/toolsrc/include/vcpkg_System.h
@@ -1,11 +1,14 @@
 #pragma once
 
+#include <Windows.h>
 #include "vcpkg_Strings.h"
 #include "filesystem_fs.h"
 #include "vcpkg_optional.h"
 
 namespace vcpkg::System
 {
+    optional<std::wstring> get_registry_string(HKEY base, const wchar_t* subkey, const wchar_t* valuename);
+
     fs::path get_exe_path_of_current_process();
 
     struct exit_code_and_output
diff --git a/toolsrc/src/metrics.cpp b/toolsrc/src/metrics.cpp
index 4ac43c5c6..b8e469f6e 100644
--- a/toolsrc/src/metrics.cpp
+++ b/toolsrc/src/metrics.cpp
@@ -227,38 +227,12 @@ true
 
     std::wstring GetSQMUser()
     {
-        LONG err;
+        auto hkcu_sqmclient = System::get_registry_string(HKEY_CURRENT_USER, LR"(Software\Microsoft\SQMClient)", L"UserId");
 
-        struct RAII_HKEY
-        {
-            HKEY hkey = nullptr;
-
-            ~RAII_HKEY()
-            {
-                if (hkey != nullptr)
-                    RegCloseKey(hkey);
-            }
-        } HKCU_SQMClient;
-
-        err = RegOpenKeyExW(HKEY_CURRENT_USER, LR"(Software\Microsoft\SQMClient)", NULL, KEY_READ, &HKCU_SQMClient.hkey);
-        if (err != ERROR_SUCCESS)
-        {
+        if (hkcu_sqmclient)
+            return std::move(*hkcu_sqmclient);
+        else
             return L"{}";
-        }
-
-        std::array<wchar_t, 128> buffer;
-        DWORD lType = 0;
-        DWORD dwBufferSize = static_cast<DWORD>(buffer.size() * sizeof(wchar_t));
-        err = RegQueryValueExW(HKCU_SQMClient.hkey, L"UserId", nullptr, &lType, reinterpret_cast<LPBYTE>(buffer.data()), &dwBufferSize);
-        if (err == ERROR_SUCCESS && lType == REG_SZ && dwBufferSize >= sizeof(wchar_t))
-        {
-            size_t sz = dwBufferSize / sizeof(wchar_t);
-            if (buffer[sz - 1] == '\0')
-                --sz;
-            return std::wstring(buffer.begin(), buffer.begin() + sz);
-        }
-
-        return L"{}";
     }
 
     void SetUserInformation(const std::string& user_id, const std::string& first_use_time)
diff --git a/toolsrc/src/vcpkg_System.cpp b/toolsrc/src/vcpkg_System.cpp
index 9609d819b..1730ff3ad 100644
--- a/toolsrc/src/vcpkg_System.cpp
+++ b/toolsrc/src/vcpkg_System.cpp
@@ -107,4 +107,32 @@ namespace vcpkg::System
     {
         _wputenv_s(varname, varvalue);
     }
+
+    static bool is_string_keytype(DWORD hkey_type)
+    {
+        return hkey_type == REG_SZ || hkey_type == REG_MULTI_SZ || hkey_type == REG_EXPAND_SZ;
+    }
+
+    optional<std::wstring> get_registry_string(HKEY base, const wchar_t* subKey, const wchar_t* valuename)
+    {
+        HKEY k = nullptr;
+        LSTATUS ec = RegOpenKeyExW(base, subKey, NULL, KEY_READ, &k);
+        if (ec != ERROR_SUCCESS)
+            return nullptr;
+
+        DWORD dwBufferSize = 0;
+        DWORD dwType = 0;
+        auto rc = RegQueryValueExW(k, valuename, nullptr, &dwType, nullptr, &dwBufferSize);
+        if (rc != ERROR_SUCCESS || !is_string_keytype(dwType) || dwBufferSize == 0 || dwBufferSize % sizeof(wchar_t) != 0)
+            return nullptr;
+        std::wstring ret;
+        ret.resize(dwBufferSize / sizeof(wchar_t));
+
+        rc = RegQueryValueExW(k, valuename, nullptr, &dwType, reinterpret_cast<LPBYTE>(ret.data()), &dwBufferSize);
+        if (rc != ERROR_SUCCESS || !is_string_keytype(dwType) || dwBufferSize != sizeof(wchar_t) * ret.size())
+            return nullptr;
+
+        ret.pop_back(); // remove extra trailing null byte
+        return std::make_unique<std::wstring>(std::move(ret));
+    }
 }