From dd13d9c6a0cee35760e96713543f83bc12decae1 Mon Sep 17 00:00:00 2001 From: es3n1n Date: Sat, 31 May 2025 13:30:48 +0200 Subject: [PATCH] refactor(com): move com utils to shared/ --- cxx-shared/cxx-shared.vcxitems | 1 + cxx-shared/shared/com.hpp | 121 ++++++++++++++++++++++++++++++ defendnot-loader/core/autorun.cpp | 57 ++++---------- defendnot/bootstrap/bootstrap.cpp | 18 ++--- defendnot/core/com.hpp | 65 +++------------- 5 files changed, 155 insertions(+), 107 deletions(-) create mode 100644 cxx-shared/shared/com.hpp diff --git a/cxx-shared/cxx-shared.vcxitems b/cxx-shared/cxx-shared.vcxitems index 1e2f684..2d57efa 100644 --- a/cxx-shared/cxx-shared.vcxitems +++ b/cxx-shared/cxx-shared.vcxitems @@ -14,6 +14,7 @@ + diff --git a/cxx-shared/shared/com.hpp b/cxx-shared/shared/com.hpp new file mode 100644 index 0000000..65f14d0 --- /dev/null +++ b/cxx-shared/shared/com.hpp @@ -0,0 +1,121 @@ +#pragma once +#include + +#include "shared/strings.hpp" + +#include +#include +#include +#include +#include + +namespace com { + inline HRESULT checked(HRESULT result, const std::source_location loc = std::source_location::current()) { + if (result == 0) { + return result; + } + + auto msg = std::format("Got HRESULT={:#x} at\n{}:{}", static_cast(result) & 0xFFFFFFFF, loc.function_name(), loc.line()); + throw std::runtime_error(msg); + } + + template + inline HRESULT retry_while_pending(Callable&& fn) { + bool delayed = false; + HRESULT status = 0; + do { + if (status != 0) { + delayed = true; + std::println("delaying for com retry..."); + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + } + + status = fn(); + } while (status == E_PENDING); + + return status; + } + + /// A very basic implementation, a lot of stuff is missing + template + class Ptr { + public: + Ptr() = default; + explicit Ptr(Ty* ptr): ptr_(ptr) { } + + ~Ptr() { + release(); + } + + /// No copying + Ptr(const Ptr&) = delete; + Ptr& operator=(const Ptr&) = delete; + + /// Move + Ptr(Ptr&& other) noexcept: ptr_(other.ptr_) { + other.ptr_ = nullptr; + } + Ptr& operator=(Ptr&& other) noexcept { + if (this != &other) { + release(); + ptr_ = other.ptr_; + other.ptr_ = nullptr; + } + return *this; + } + + public: + [[nodiscard]] Ty* get() const { + return ptr_; + } + + [[nodiscard]] Ty* operator->() const { + return ptr_; + } + + [[nodiscard]] Ty** ref_to_ptr() { + return &ptr_; + } + + private: + void release() { + if (ptr_ != nullptr) { + ptr_->Release(); + } + } + + Ty* ptr_ = nullptr; + }; + + template + class IBaseObject { + public: + static constexpr GUID kClsId = ClsId; + static constexpr GUID kIID = IID; + + private: + virtual HRESULT QueryInterface() = 0; + virtual std::uint32_t AddRef() = 0; + + public: + virtual std::uint32_t Release() = 0; + }; + + template + concept ComObject = requires { + Ty::kClsId; + Ty::kIID; + }; + + template + [[nodiscard]] Ptr query() { + Ptr result; + const auto status = CoCreateInstance(Ty::kClsId, 0, 1, Ty::kIID, reinterpret_cast(result.ref_to_ptr())); + if (status == REGDB_E_CLASSNOTREG) { + throw std::runtime_error(strings::wsc_unavailable_error().data()); + } + + checked(status); + return result; + } +} // namespace com diff --git a/defendnot-loader/core/autorun.cpp b/defendnot-loader/core/autorun.cpp index 3d07d2d..4b7db24 100644 --- a/defendnot-loader/core/autorun.cpp +++ b/defendnot-loader/core/autorun.cpp @@ -1,5 +1,6 @@ #include "core/core.hpp" +#include "shared/com.hpp" #include "shared/ctx.hpp" #include "shared/strings.hpp" @@ -19,38 +20,6 @@ namespace loader { namespace { constexpr std::string_view kTaskName = strings::kProjectName; - /// A very basic implementation, a lot of stuff is missing - template - class ComPtr { - public: - ComPtr() = default; - explicit ComPtr(Ty* ptr): ptr_(ptr) { } - - ~ComPtr() { - if (ptr_ != nullptr) { - ptr_->Release(); - } - } - - ComPtr(const ComPtr&) = delete; - ComPtr& operator=(const ComPtr&) = delete; - - [[nodiscard]] Ty* get() const { - return ptr_; - } - - [[nodiscard]] Ty* operator->() const { - return ptr_; - } - - [[nodiscard]] Ty** ref_to_ptr() { - return &ptr_; - } - - private: - Ty* ptr_ = nullptr; - }; - void co_initialize() { static std::once_flag fl; std::call_once(fl, []() -> void { @@ -66,7 +35,7 @@ namespace loader { [[nodiscard]] bool with_service(Callable&& callback) { co_initialize(); - ComPtr service; + com::Ptr service; auto hr = CoCreateInstance(CLSID_TaskScheduler, nullptr, CLSCTX_INPROC_SERVER, IID_ITaskService, reinterpret_cast(service.ref_to_ptr())); if (FAILED(hr)) { @@ -78,7 +47,7 @@ namespace loader { return false; } - ComPtr root_folder; + com::Ptr root_folder; hr = service->GetFolder(BSTR(L"\\"), root_folder.ref_to_ptr()); if (FAILED(hr)) { return false; @@ -104,55 +73,55 @@ namespace loader { user_id = bstr_sys; } - ComPtr task; + com::Ptr task; auto hr = service->NewTask(0, task.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr reg_info; + com::Ptr reg_info; hr = task->get_RegistrationInfo(reg_info.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr principal; + com::Ptr principal; hr = task->get_Principal(principal.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr trigger_collection; + com::Ptr trigger_collection; hr = task->get_Triggers(trigger_collection.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr trigger; + com::Ptr trigger; hr = trigger_collection->Create(task_trigger, trigger.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr action_collection; + com::Ptr action_collection; hr = task->get_Actions(action_collection.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr action; + com::Ptr action; hr = action_collection->Create(TASK_ACTION_EXEC, action.ref_to_ptr()); if (FAILED(hr)) { return false; } - ComPtr exec_action; + com::Ptr exec_action; hr = action->QueryInterface(IID_IExecAction, reinterpret_cast(exec_action.ref_to_ptr())); if (FAILED(hr)) { return false; } - ComPtr settings; + com::Ptr settings; hr = task->get_Settings(settings.ref_to_ptr()); if (FAILED(hr)) { return false; @@ -175,7 +144,7 @@ namespace loader { exec_action->put_Arguments(bstr_t("--from-autorun")); /// Register the task and we are done - ComPtr registered_task; + com::Ptr registered_task; hr = folder->RegisterTaskDefinition(bstr_t(kTaskName.data()), task.get(), TASK_CREATE_OR_UPDATE, VARIANT{}, VARIANT{}, TASK_LOGON_NONE, variant_t(L""), registered_task.ref_to_ptr()); return SUCCEEDED(hr); diff --git a/defendnot/bootstrap/bootstrap.cpp b/defendnot/bootstrap/bootstrap.cpp index e91d4bd..913c807 100644 --- a/defendnot/bootstrap/bootstrap.cpp +++ b/defendnot/bootstrap/bootstrap.cpp @@ -10,13 +10,13 @@ namespace defendnot { void startup() { /// Setup shared::ctx.deserialize(); - logln("init: {:#x}", com_checked(CoInitialize(nullptr))); + logln("init: {:#x}", com::checked(CoInitialize(nullptr))); /// Get the main WSC interface we will be dealing with - auto inst = IWscAVStatus::get(); + auto inst = com::query(); /// This can fail if we dont have any avs registered so no com_checked - logln("unregister: {:#x}", com_retry_while_pending([&inst]() -> HRESULT { return inst->Unregister(); }) & 0xFFFFFFFF); + logln("unregister: {:#x}", com::retry_while_pending([&inst]() -> HRESULT { return inst->Unregister(); }) & 0xFFFFFFFF); if (shared::ctx.state == shared::State::OFF) { return; } @@ -34,12 +34,10 @@ namespace defendnot { }; /// Register and activate our AV - logln("register: {:#x}", com_checked(inst->Register(name, name, 0, 0))); - logln("update: {:#x}", com_checked(inst->UpdateStatus(WSCSecurityProductState::ON, WSCSecurityProductState::ON))); - logln("scan_update: {:#x}", com_checked(inst->UpdateScanSubstatus(WSCSecurityProductSubStatus::NO_ACTION))); - logln("settings_update: {:#x}", com_checked(inst->UpdateSettingsSubstatus(WSCSecurityProductSubStatus::NO_ACTION))); - - /// \fixme @es3n1n: this seems to be giving no effect whatsoever. - logln("prot_update: {:#x}", com_checked(inst->UpdateProtectionUpdateSubstatus(WSCSecurityProductSubStatus::NO_ACTION))); + logln("register: {:#x}", com::checked(inst->Register(name, name, 0, 0))); + logln("update: {:#x}", com::checked(inst->UpdateStatus(WSCSecurityProductState::ON, WSCSecurityProductState::ON))); + logln("scan_update: {:#x}", com::checked(inst->UpdateScanSubstatus(WSCSecurityProductSubStatus::NO_ACTION))); + logln("settings_update: {:#x}", com::checked(inst->UpdateSettingsSubstatus(WSCSecurityProductSubStatus::NO_ACTION))); + logln("prot_update: {:#x}", com::checked(inst->UpdateProtectionUpdateSubstatus(WSCSecurityProductSubStatus::NO_ACTION))); } } // namespace defendnot diff --git a/defendnot/core/com.hpp b/defendnot/core/com.hpp index 0788bc1..1874f83 100644 --- a/defendnot/core/com.hpp +++ b/defendnot/core/com.hpp @@ -5,14 +5,23 @@ #include #include "core/log.hpp" +#include "shared/com.hpp" #include "shared/strings.hpp" #include namespace defendnot { namespace detail { - inline GUID CLSID_IWscAVStatus = {0x0F2102C37, 0x90C3, 0x450C, {0x0B3, 0x0F6, 0x92, 0x0BE, 0x16, 0x93, 0x0BD, 0x0F2}}; - inline GUID IID_IWscAVStatus = {0x3901A765, 0x0AB91, 0x4BA9, {0xA5, 0x53, 0x5B, 0x85, 0x38, 0xDE, 0xB8, 0x40}}; + constexpr GUID CLSID_WscIsv = {0xF2102C37, 0x90C3, 0x450C, {0xB3, 0x0F6, 0x92, 0xBE, 0x16, 0x93, 0xBD, 0xF2}}; + + constexpr GUID IID_IWscFWStatus = {0x9B8F6C6E, 0x8A4A, 0x4891, {0xAF, 0x63, 0x1A, 0x2F, 0x50, 0x92, 0x40, 0x40}}; + constexpr GUID IID_IWscFWStatus2 = {0x62F698CB, 0x94A, 0x4C68, {0x94, 0x19, 0x8E, 0x8C, 0x49, 0x42, 0x0E, 0x59}}; + + constexpr GUID IID_IWscAVStatus = {0x3901A765, 0xAB91, 0x4BA9, {0xA5, 0x53, 0x5B, 0x85, 0x38, 0xDE, 0xB8, 0x40}}; + constexpr GUID IID_IWscAVStatus3 = {0xCF007CA2, 0xF5E3, 0x11E5, {0x9C, 0xE9, 0x5E, 0x55, 0x17, 0x50, 0x7C, 0x66}}; + constexpr GUID IID_IWscAVStatus4 = {0x4DCBAFAC, 0x29BA, 0x46B1, {0x80, 0xFC, 0xB8, 0xBD, 0xE3, 0xC0, 0xAE, 0x4D}}; + + constexpr GUID IID_IWscASStatus = {0x24E9756, 0xBA6C, 0x4AD1, {0x83, 0x21, 0x87, 0xBA, 0xE7, 0x8F, 0xD0, 0xE3}}; } // namespace detail enum class WSCSecurityProductState : std::uint32_t { @@ -29,43 +38,8 @@ namespace defendnot { ACTION_NEEDED = 3, }; - inline HRESULT com_checked(HRESULT result, const std::source_location loc = std::source_location::current()) { - if (result == 0) { - return result; - } - - auto msg = std::format("Got HRESULT={:#x} at\n{}:{}", static_cast(result) & 0xFFFFFFFF, loc.function_name(), loc.line()); - throw std::runtime_error(msg); - } - - template - inline HRESULT com_retry_while_pending(Callable&& fn) { - bool delayed = false; - HRESULT status = 0; - do { - if (status != 0) { - delayed = true; - logln("delaying for com retry..."); - std::this_thread::sleep_for(std::chrono::seconds(5)); - } - - status = fn(); - } while (status == E_PENDING); - - if (delayed) { - /// Sleep for additional 15 seconds to let WSC proceed all previous requests - /// \todo @es3n1n: there should be a better way to handle this - std::this_thread::sleep_for(std::chrono::seconds(15)); - } - - return status; - } - - class IWscAVStatus { + class IWscAVStatus4 : public com::IBaseObject { public: - virtual HRESULT QueryInterface() = 0; - virtual std::uint32_t AddRef() = 0; - virtual std::uint32_t Release() = 0; virtual HRESULT Register(BSTR path_to_signed_product_exe, BSTR display_name, std::uint32_t, std::uint32_t) = 0; virtual HRESULT Unregister() = 0; virtual HRESULT UpdateStatus(WSCSecurityProductState state, WSCSecurityProductState state2) = 0; @@ -87,20 +61,5 @@ namespace defendnot { virtual HRESULT RegisterAS(std::uint16_t*, std::uint16_t*, std::uint32_t, std::uint32_t) = 0; virtual HRESULT UnregisterAS() = 0; virtual HRESULT UpdateStatusAS(WSCSecurityProductState state, WSCSecurityProductState state2) = 0; - - private: - virtual void dtor() = 0; - - public: - static IWscAVStatus* get() { - IWscAVStatus* result = nullptr; - const auto status = CoCreateInstance(detail::CLSID_IWscAVStatus, 0, 1, detail::IID_IWscAVStatus, reinterpret_cast(&result)); - if (status == REGDB_E_CLASSNOTREG) { - throw std::runtime_error(strings::wsc_unavailable_error().data()); - } - - com_checked(status); - return result; - } }; } // namespace defendnot