diff --git a/defendnot-loader/core/autorun.cpp b/defendnot-loader/core/autorun.cpp index 4bc5248..cfe4a25 100644 --- a/defendnot-loader/core/autorun.cpp +++ b/defendnot-loader/core/autorun.cpp @@ -1,11 +1,12 @@ #include "core/core.hpp" #include "shared/ctx.hpp" -#include "shared/defer.hpp" #include "shared/names.hpp" +#include #include #include +#include #include #include @@ -18,44 +19,74 @@ namespace loader { namespace { constexpr std::string_view kTaskName = names::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 { + const auto result = CoInitializeEx(nullptr, COINIT_MULTITHREADED); + + if (FAILED(result)) { + throw std::runtime_error("failed to CoInitializeEx"); + } + }); + } + template [[nodiscard]] bool with_service(Callable&& callback) { - HRESULT hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED); + co_initialize(); + + ComPtr service; + auto hr = + CoCreateInstance(CLSID_TaskScheduler, nullptr, CLSCTX_INPROC_SERVER, IID_ITaskService, reinterpret_cast(service.ref_to_ptr())); if (FAILED(hr)) { return false; } - defer->void { - CoUninitialize(); - }; - - ITaskService* service = nullptr; - hr = CoCreateInstance(CLSID_TaskScheduler, nullptr, CLSCTX_INPROC_SERVER, IID_ITaskService, reinterpret_cast(&service)); - if (FAILED(hr)) { - return false; - } - - defer->void { - service->Release(); - }; - hr = service->Connect(VARIANT{}, VARIANT{}, VARIANT{}, VARIANT{}); if (FAILED(hr)) { return false; } - ITaskFolder* root_folder = nullptr; - hr = service->GetFolder(BSTR(L"\\"), &root_folder); + ComPtr root_folder; + hr = service->GetFolder(BSTR(L"\\"), root_folder.ref_to_ptr()); if (FAILED(hr)) { return false; } - defer->void { - root_folder->Release(); - }; - + /// Cleanup our task, we will recreate it in the callback if needed root_folder->DeleteTask(BSTR(kTaskName.data()), 0); - return callback(service, root_folder); + return callback(service.get(), root_folder.get()); } } // namespace @@ -63,110 +94,70 @@ namespace loader { const auto bin_path = shared::get_this_module_path(); return with_service([bin_path](ITaskService* service, ITaskFolder* folder) -> bool { - ITaskDefinition* task = nullptr; - auto hr = service->NewTask(0, &task); + ComPtr task; + auto hr = service->NewTask(0, task.ref_to_ptr()); if (FAILED(hr)) { return false; } - defer->void { - task->Release(); - }; - - IRegistrationInfo* reg_info = nullptr; - hr = task->get_RegistrationInfo(®_info); + ComPtr reg_info; + hr = task->get_RegistrationInfo(reg_info.ref_to_ptr()); if (FAILED(hr)) { return false; } - defer->void { - reg_info->Release(); - }; + ComPtr principal; + hr = task->get_Principal(principal.ref_to_ptr()); + if (FAILED(hr)) { + return false; + } + + ComPtr trigger_collection; + hr = task->get_Triggers(trigger_collection.ref_to_ptr()); + if (FAILED(hr)) { + return false; + } + + ComPtr trigger; + hr = trigger_collection->Create(TASK_TRIGGER_LOGON, trigger.ref_to_ptr()); + if (FAILED(hr)) { + return false; + } + + ComPtr action_collection; + hr = task->get_Actions(action_collection.ref_to_ptr()); + if (FAILED(hr)) { + return false; + } + + ComPtr action; + hr = action_collection->Create(TASK_ACTION_EXEC, action.ref_to_ptr()); + if (FAILED(hr)) { + return false; + } + + ComPtr exec_action; + hr = action->QueryInterface(IID_IExecAction, reinterpret_cast(exec_action.ref_to_ptr())); + if (FAILED(hr)) { + return false; + } + + ComPtr settings; + hr = task->get_Settings(settings.ref_to_ptr()); + if (FAILED(hr)) { + return false; + } + + principal->put_RunLevel(TASK_RUNLEVEL_HIGHEST); reg_info->put_Author(_bstr_t(names::kRepoUrl.data())); - - IPrincipal* pPrincipal = nullptr; - hr = task->get_Principal(&pPrincipal); - if (FAILED(hr)) { - return false; - } - - defer->void { - pPrincipal->Release(); - }; - - ITriggerCollection* trigger_collection = nullptr; - hr = task->get_Triggers(&trigger_collection); - if (FAILED(hr)) { - return false; - } - - defer->void { - trigger_collection->Release(); - }; - - ITrigger* trigger = nullptr; - hr = trigger_collection->Create(TASK_TRIGGER_LOGON, &trigger); - if (FAILED(hr)) { - return false; - } - - defer->void { - trigger->Release(); - }; - - IActionCollection* action_collection = nullptr; - hr = task->get_Actions(&action_collection); - if (FAILED(hr)) { - return false; - } - - defer->void { - action_collection->Release(); - }; - - IAction* action = nullptr; - hr = action_collection->Create(TASK_ACTION_EXEC, &action); - if (FAILED(hr)) { - return false; - } - - defer->void { - action->Release(); - }; - - IExecAction* exec_action = nullptr; - hr = action->QueryInterface(IID_IExecAction, (void**)&exec_action); - if (FAILED(hr)) { - return false; - } - - defer->void { - exec_action->Release(); - }; - - ITaskSettings* settings = nullptr; - hr = task->get_Settings(&settings); - if (FAILED(hr)) { - return false; - } - - defer->void { - settings->Release(); - }; - - pPrincipal->put_RunLevel(TASK_RUNLEVEL_HIGHEST); settings->put_DisallowStartIfOnBatteries(VARIANT_FALSE); settings->put_StopIfGoingOnBatteries(VARIANT_FALSE); exec_action->put_Path(_bstr_t(bin_path.string().c_str())); exec_action->put_Arguments(_bstr_t("--from-autorun")); - IRegisteredTask* registered_task = nullptr; - hr = folder->RegisterTaskDefinition(_bstr_t(kTaskName.data()), task, TASK_CREATE_OR_UPDATE, VARIANT{}, VARIANT{}, TASK_LOGON_INTERACTIVE_TOKEN, - _variant_t(L""), ®istered_task); - - defer->void { - registered_task->Release(); - }; + ComPtr registered_task; + hr = folder->RegisterTaskDefinition(_bstr_t(kTaskName.data()), task.get(), TASK_CREATE_OR_UPDATE, VARIANT{}, VARIANT{}, + TASK_LOGON_INTERACTIVE_TOKEN, _variant_t(L""), registered_task.ref_to_ptr()); return SUCCEEDED(hr); }); }