sync/es3n1n-defendnot
1
0
Fork 0
mirror of https://github.com/es3n1n/defendnot.git synced 2026-08-02 10:32:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #6 from es3n1n/com-ptr-startup

Browse Source 
refactor(autorun): get rid of defers
This commit is contained in:
Arsenii es3n1n
2025-05-14 08:17:47 +09:00
committed by GitHub
parent 01351bd406 db4158da3b
commit 782246d642
1 changed files with 104 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files

217
defendnot-loader/core/autorun.cpp
View File

@@ -1,11 +1,12 @@
#include "core/core.hpp"
#include "shared/ctx.hpp"
#include "shared/defer.hpp"
#include "shared/names.hpp"
#include <memory>
#include <print>
#include <stdexcept>
#include <type_traits>
#include <comdef.h>
#include <taskschd.h>
@@ -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 <typename Ty>
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 <typename Callable>
[[nodiscard]] bool with_service(Callable&& callback) {
HRESULT hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
co_initialize();
ComPtr<ITaskService> service;
auto hr =
CoCreateInstance(CLSID_TaskScheduler, nullptr, CLSCTX_INPROC_SERVER, IID_ITaskService, reinterpret_cast<void**>(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<void**>(&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<ITaskFolder> 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<ITaskDefinition> 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(&reg_info);
ComPtr<IRegistrationInfo> reg_info;
hr = task->get_RegistrationInfo(reg_info.ref_to_ptr());
if (FAILED(hr)) {
return false;
}
defer->void {
reg_info->Release();
};
ComPtr<IPrincipal> principal;
hr = task->get_Principal(principal.ref_to_ptr());
if (FAILED(hr)) {
return false;
}
ComPtr<ITriggerCollection> trigger_collection;
hr = task->get_Triggers(trigger_collection.ref_to_ptr());
if (FAILED(hr)) {
return false;
}
ComPtr<ITrigger> trigger;
hr = trigger_collection->Create(TASK_TRIGGER_LOGON, trigger.ref_to_ptr());
if (FAILED(hr)) {
return false;
}
ComPtr<IActionCollection> action_collection;
hr = task->get_Actions(action_collection.ref_to_ptr());
if (FAILED(hr)) {
return false;
}
ComPtr<IAction> action;
hr = action_collection->Create(TASK_ACTION_EXEC, action.ref_to_ptr());
if (FAILED(hr)) {
return false;
}
ComPtr<IExecAction> exec_action;
hr = action->QueryInterface(IID_IExecAction, reinterpret_cast<void**>(exec_action.ref_to_ptr()));
if (FAILED(hr)) {
return false;
}
ComPtr<ITaskSettings> 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""), &registered_task);
defer->void {
registered_task->Release();
};
ComPtr<IRegisteredTask> 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);
});
}