Merge pull request #168 from icnocop/master

Initializing COM before using CComPtr to resolve runtime application crash when running WMI query checks
dotnetinstaller · Sep 8, 2024 · 01d7f51 · 01d7f51
2 parents 9231303 + 71bef87
commit 01d7f51
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 37 deletions.
diff --git a/dotNetInstallerToolsLib/ComInitializer.h b/dotNetInstallerToolsLib/ComInitializer.h
@@ -5,19 +5,21 @@
 #include "ExceptionMacros.h"
 #include <comdef.h>
 
-class ComInitializer
+namespace DVLib
 {
-public:
-    ComInitializer()
+    class ComInitializer
     {
-        // initialize COM
-        CHECK_HR(::CoInitializeEx(0, COINIT_APARTMENTTHREADED),
-            L"Failed to initialize COM library.");
-    }
-
-    ~ComInitializer()
-    {
-        ::CoUninitialize();
-    }
-};
+    public:
+        ComInitializer()
+        {
+            // initialize COM
+            CHECK_HR(::CoInitializeEx(0, COINIT_APARTMENTTHREADED),
+                L"Failed to initialize COM library.");
+        }
 
+        ~ComInitializer()
+        {
+            ::CoUninitialize();
+        }
+    };
+}
diff --git a/dotNetInstallerToolsLib/WmiUtil.cpp b/dotNetInstallerToolsLib/WmiUtil.cpp
@@ -10,6 +10,9 @@
 
 bool DVLib::WmiUtil::Any(const std::wstring& query)
 {
+    // initialize COM
+    ComInitializer comInitializer;
+
     CComPtr<IWbemClassObject> pclsObj = nullptr;
     ULONG uReturn = FirstOrNull(query, pclsObj.p);
 
@@ -18,6 +21,9 @@ bool DVLib::WmiUtil::Any(const std::wstring& query)
 
 std::wstring DVLib::WmiUtil::Get(const std::wstring& query, const std::wstring& propertyName)
 {
+    // initialize COM
+    ComInitializer comInitializer;
+
     CComPtr<IWbemClassObject> pclsObj = nullptr;
     ULONG uReturn = FirstOrNull(query, pclsObj.p);
     if (uReturn == 0)
@@ -26,7 +32,7 @@ std::wstring DVLib::WmiUtil::Get(const std::wstring& query, const std::wstring&
         return L"";
     }
 
-    _variant_t vtProp;
+    _variant_t vtProp = {};
 
     // get the value of the specified property
     CHECK_HR(pclsObj->Get(propertyName.c_str(), 0, &vtProp, 0, 0),
@@ -37,37 +43,43 @@ std::wstring DVLib::WmiUtil::Get(const std::wstring& query, const std::wstring&
 
 ULONG DVLib::WmiUtil::FirstOrNull(const std::wstring& query, IWbemClassObject*& pclsObj)
 {
-    // initialize COM
-    ComInitializer comInitializer;
 
     // create an instance of the COM object
     CComPtr<IWbemLocator> pLocator = nullptr;
 
     CHECK_HR(::CoCreateInstance(
-        CLSID_WbemLocator,          // Class identifier
-        0,                        // IUnknown
-        CLSCTX_INPROC_SERVER,   // Context
-        IID_IWbemLocator,             // Interface identifier
-        (LPVOID*)&pLocator),           // Pointer to created interface
+        CLSID_WbemLocator,    // Class identifier
+        0,                    // IUnknown
+        CLSCTX_INPROC_SERVER, // Context
+        IID_IWbemLocator,     // Interface identifier
+        (LPVOID*)&pLocator),  // Pointer to created interface
         L"Failed to create IWbemLocator object.");
+    if (pLocator == nullptr)
+    {
+        throw std::exception("Failed to create IWbemLocator object.");
+    }
 
     // connect to the WMI namespace
     CComPtr<IWbemServices> pServices = nullptr;
 
     CHECK_HR(pLocator->ConnectServer(
-        _bstr_t(L"ROOT\\CIMV2"),  // WMI namespace
-        NULL,                                    // User name
-        NULL,                                    // Password
-        0,                                       // Locale
-        NULL,                                    // Security flags
-        0,                                       // Authority
-        0,                                       // Context
-        &pServices),                             // IWbemServices proxy
+        _bstr_t(L"ROOT\\CIMV2"), // WMI namespace
+        NULL,                    // User name
+        NULL,                    // Password
+        0,                       // Locale
+        NULL,                    // Security flags
+        0,                       // Authority
+        0,                       // Context
+        &pServices),             // IWbemServices proxy
         L"Failed to connect to WMI namespace.");
+    if (pServices == nullptr)
+    {
+        throw std::exception("Failed to connect to WMI namespace.");
+    }
 
     // set security levels on the proxy
     CHECK_HR(::CoSetProxyBlanket(
-        pServices,             // IWbemServices proxy
+        pServices,                   // IWbemServices proxy
         RPC_C_AUTHN_WINNT,           // Authentication service
         RPC_C_AUTHZ_NONE,            // Authorization service
         NULL,                        // Server principal name
@@ -81,11 +93,11 @@ ULONG DVLib::WmiUtil::FirstOrNull(const std::wstring& query, IWbemClassObject*&
     CComPtr<IEnumWbemClassObject> pEnumerator = nullptr;
 
     CHECK_HR(pServices->ExecQuery(
-        bstr_t("WQL"),                                  // Query language
-        bstr_t(query.c_str()),                                // Query
+        bstr_t("WQL"),                                         // Query language
+        bstr_t(query.c_str()),                                 // Query
         WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, // Flags
-        NULL,                                                        // Context
-        &pEnumerator),                                               // Enumerator
+        NULL,                                                  // Context
+        &pEnumerator),                                         // Enumerator
         L"Failed to execute the WMI query.");
     if (pEnumerator == nullptr)
     {
@@ -97,9 +109,9 @@ ULONG DVLib::WmiUtil::FirstOrNull(const std::wstring& query, IWbemClassObject*&
 
     CHECK_HR(pEnumerator->Next(
         WBEM_INFINITE, // Timeout
-        1,              // Count of objects to return
-        &pclsObj,     // Objects
-        &uReturn),   // Count of returned objects
+        1,             // Count of objects to return
+        &pclsObj,      // Objects
+        &uReturn),     // Count of returned objects
         L"Failed to get the WMI class object.");
 
     return uReturn;