From 11634b9a6749a65409d2cfa286bbb3d9c10d07ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E9=80=9A=E6=B4=B2?= Date: Wed, 25 Jan 2023 15:49:31 +0800 Subject: [PATCH] Windows: add com helper --- CMakeLists.txt | 2 +- src/util/windows/com.cpp | 51 ++++++++++++++++++++++++++++++++++++++++ src/util/windows/com.hpp | 25 ++++++++++++++++++++ src/util/windows/wmi.cpp | 44 +++++++--------------------------- src/util/windows/wmi.hpp | 1 + 5 files changed, 87 insertions(+), 36 deletions(-) create mode 100644 src/util/windows/com.cpp create mode 100644 src/util/windows/com.hpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 83a3591fc..be8be1a71 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -511,6 +511,7 @@ elseif(WIN32) src/detection/wifi/wifi_windows.c src/detection/wmtheme/wmtheme_windows.c src/util/windows/getline.c + src/util/windows/com.cpp src/util/windows/registry.c src/util/windows/unicode.c src/util/windows/wmi.cpp @@ -697,7 +698,6 @@ elseif(WIN32) PRIVATE "ole32" PRIVATE "oleaut32" PRIVATE "opengl32" - PRIVATE "wbemuuid" PRIVATE "ws2_32" PRIVATE "ntdll" PRIVATE "version" diff --git a/src/util/windows/com.cpp b/src/util/windows/com.cpp new file mode 100644 index 000000000..c11cccf95 --- /dev/null +++ b/src/util/windows/com.cpp @@ -0,0 +1,51 @@ +#include "com.hpp" +#include "fastfetch.h" + +#include +#include + +//https://learn.microsoft.com/en-us/windows/win32/wmisdk/example--getting-wmi-data-from-the-local-computer +//https://learn.microsoft.com/en-us/windows/win32/cimwin32prov/computer-system-hardware-classes +static void CoUninitializeWrap() +{ + CoUninitialize(); +} + +static BOOL CALLBACK InitHandleFunction(FF_MAYBE_UNUSED PINIT_ONCE lpInitOnce, FF_MAYBE_UNUSED PVOID lpParameter, PVOID* lpContext) +{ + // Initialize COM + if (FAILED(CoInitializeEx(NULL, COINIT_MULTITHREADED))) + { + *lpContext = (PVOID) "CoInitializeEx() failed"; + return FALSE; + } + + // Set general COM security levels + if (FAILED(CoInitializeSecurity( + NULL, + -1, // COM authentication + NULL, // Authentication services + NULL, // Reserved + RPC_C_AUTHN_LEVEL_DEFAULT, // Default authentication + RPC_C_IMP_LEVEL_IMPERSONATE, // Default Impersonation + NULL, // Authentication info + EOAC_NONE, // Additional capabilities + NULL // Reserved + ))) + { + CoUninitialize(); + *lpContext = (PVOID) "CoInitializeSecurity() failed"; + return FALSE; + } + + atexit(CoUninitializeWrap); + return TRUE; +} + +const char* ffInitCom(void) +{ + const char* error = NULL; + static INIT_ONCE s_InitOnce; + InitOnceExecuteOnce(&s_InitOnce, &InitHandleFunction, NULL, (void**)&error); + return error; +} diff --git a/src/util/windows/com.hpp b/src/util/windows/com.hpp new file mode 100644 index 000000000..e0f1051ae --- /dev/null +++ b/src/util/windows/com.hpp @@ -0,0 +1,25 @@ +#pragma once + +#ifndef FF_INCLUDED_util_windows_com +#define FF_INCLUDED_util_windows_com + +#ifdef __cplusplus + +#include + +const char* ffInitCom(void); + +static inline void ffReleaseComObject(void* ppUnknown) +{ + IUnknown* pUnknown = *(IUnknown**) ppUnknown; + if (pUnknown) pUnknown->Release(); +} + +#define FF_AUTO_RELEASE_COM_OBJECT __attribute__((__cleanup__(ffReleaseComObject))) + +#else + // Win32 COM headers requires C++ compiler + #error Must be included in C++ source file +#endif + +#endif diff --git a/src/util/windows/wmi.cpp b/src/util/windows/wmi.cpp index 1da4ac0ff..65e54d422 100644 --- a/src/util/windows/wmi.cpp +++ b/src/util/windows/wmi.cpp @@ -1,4 +1,5 @@ #include "wmi.hpp" +#include "util/windows/com.hpp" #include #include @@ -19,40 +20,10 @@ namespace }; } -//https://learn.microsoft.com/en-us/windows/win32/wmisdk/example--getting-wmi-data-from-the-local-computer -//https://learn.microsoft.com/en-us/windows/win32/cimwin32prov/computer-system-hardware-classes -static void CoUninitializeWrap() -{ - CoUninitialize(); -} - static BOOL CALLBACK InitHandleFunction(PINIT_ONCE, PVOID lpParameter, PVOID* lpContext) { HRESULT hres; - // Initialize COM - hres = CoInitializeEx(nullptr, COINIT_MULTITHREADED); - - // Set general COM security levels - hres = CoInitializeSecurity( - nullptr, - -1, // COM authentication - nullptr, // Authentication services - nullptr, // Reserved - RPC_C_AUTHN_LEVEL_DEFAULT, // Default authentication - RPC_C_IMP_LEVEL_IMPERSONATE, // Default Impersonation - nullptr, // Authentication info - EOAC_NONE, // Additional capabilities - nullptr // Reserved - ); - - if (FAILED(hres)) - { - CoUninitialize(); - *((const char**)lpContext) = "Failed to initialize security"; - return FALSE; - } - // Obtain the initial locator to WMI IWbemLocator* pLoc = nullptr; hres = CoCreateInstance( @@ -64,7 +35,6 @@ static BOOL CALLBACK InitHandleFunction(PINIT_ONCE, PVOID lpParameter, PVOID* lp if (FAILED(hres)) { - CoUninitialize(); *((const char**)lpContext) = "Failed to create IWbemLocator object"; return FALSE; } @@ -90,7 +60,6 @@ static BOOL CALLBACK InitHandleFunction(PINIT_ONCE, PVOID lpParameter, PVOID* lp if (FAILED(hres)) { - CoUninitialize(); *((const char**)lpContext) = "Could not connect WMI server"; return FALSE; } @@ -110,21 +79,26 @@ static BOOL CALLBACK InitHandleFunction(PINIT_ONCE, PVOID lpParameter, PVOID* lp if (FAILED(hres)) { pSvc->Release(); - CoUninitialize(); *((const char**)lpContext) = "Could not set proxy blanket"; return FALSE; } *((IWbemServices**)lpContext) = pSvc; - atexit(CoUninitializeWrap); return TRUE; } FFWmiQuery::FFWmiQuery(const wchar_t* queryStr, FFstrbuf* error, FFWmiNamespace wmiNs) : pEnumerator(nullptr) { + const char* context = ffInitCom(); + if (context) + { + if (error) + ffStrbufAppendS(error, context); + return; + } + static INIT_ONCE s_InitOnce[(int) FFWmiNamespace::LAST] = {}; - const char* context; if (InitOnceExecuteOnce( &s_InitOnce[(int)wmiNs], &InitHandleFunction, diff --git a/src/util/windows/wmi.hpp b/src/util/windows/wmi.hpp index 25bef6e91..c4ece40f6 100644 --- a/src/util/windows/wmi.hpp +++ b/src/util/windows/wmi.hpp @@ -9,6 +9,7 @@ extern "C" { #include "util/FFstrbuf.h" } +#include #include enum class FFWmiNamespace {