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

refactor(autorun): get rid of defers

Browse Source
This commit is contained in:
es3n1n
2025-05-12 17:03:57 +09:00
parent 928be9a45f
commit db4158da3b
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 "core/core.hpp"
#include "shared/ctx.hpp" #include "shared/ctx.hpp"
#include "shared/defer.hpp"
#include "shared/names.hpp" #include "shared/names.hpp"
#include <memory>
#include <print> #include <print>
#include <stdexcept> #include <stdexcept>
#include <type_traits>
#include <comdef.h> #include <comdef.h>
#include <taskschd.h> #include <taskschd.h>
@@ -18,44 +19,74 @@ namespace loader {
namespace { namespace {
constexpr std::string_view kTaskName = names::kProjectName; 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> template <typename Callable>
[[nodiscard]] bool with_service(Callable&& callback) { [[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)) { if (FAILED(hr)) {
return false; 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{}); hr = service->Connect(VARIANT{}, VARIANT{}, VARIANT{}, VARIANT{});
if (FAILED(hr)) { if (FAILED(hr)) {
return false; return false;
} }
ITaskFolder* root_folder = nullptr; ComPtr<ITaskFolder> root_folder;
hr = service->GetFolder(BSTR(L"\\"), &root_folder); hr = service->GetFolder(BSTR(L"\\"), root_folder.ref_to_ptr());
if (FAILED(hr)) { if (FAILED(hr)) {
return false; return false;
} }
defer->void { /// Cleanup our task, we will recreate it in the callback if needed
root_folder->Release();
};
root_folder->DeleteTask(BSTR(kTaskName.data()), 0); root_folder->DeleteTask(BSTR(kTaskName.data()), 0);
return callback(service, root_folder); return callback(service.get(), root_folder.get());
} }
} // namespace } // namespace
@@ -63,110 +94,70 @@ namespace loader {
const auto bin_path = shared::get_this_module_path(); const auto bin_path = shared::get_this_module_path();
return with_service([bin_path](ITaskService* service, ITaskFolder* folder) -> bool { return with_service([bin_path](ITaskService* service, ITaskFolder* folder) -> bool {
ITaskDefinition* task = nullptr; ComPtr<ITaskDefinition> task;
auto hr = service->NewTask(0, &task); auto hr = service->NewTask(0, task.ref_to_ptr());
if (FAILED(hr)) { if (FAILED(hr)) {
return false; return false;
} }
defer->void { ComPtr<IRegistrationInfo> reg_info;
task->Release(); hr = task->get_RegistrationInfo(reg_info.ref_to_ptr());
};
IRegistrationInfo* reg_info = nullptr;
hr = task->get_RegistrationInfo(&reg_info);
if (FAILED(hr)) { if (FAILED(hr)) {
return false; return false;
} }
defer->void { ComPtr<IPrincipal> principal;
reg_info->Release(); 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())); 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_DisallowStartIfOnBatteries(VARIANT_FALSE);
settings->put_StopIfGoingOnBatteries(VARIANT_FALSE); settings->put_StopIfGoingOnBatteries(VARIANT_FALSE);
exec_action->put_Path(_bstr_t(bin_path.string().c_str())); exec_action->put_Path(_bstr_t(bin_path.string().c_str()));
exec_action->put_Arguments(_bstr_t("--from-autorun")); exec_action->put_Arguments(_bstr_t("--from-autorun"));
IRegisteredTask* registered_task = nullptr; ComPtr<IRegisteredTask> registered_task;
hr = folder->RegisterTaskDefinition(_bstr_t(kTaskName.data()), task, TASK_CREATE_OR_UPDATE, VARIANT{}, VARIANT{}, TASK_LOGON_INTERACTIVE_TOKEN, hr = folder->RegisterTaskDefinition(_bstr_t(kTaskName.data()), task.get(), TASK_CREATE_OR_UPDATE, VARIANT{}, VARIANT{},
_variant_t(L""), &registered_task); TASK_LOGON_INTERACTIVE_TOKEN, _variant_t(L""), registered_task.ref_to_ptr());
defer->void {
registered_task->Release();
};
return SUCCEEDED(hr); return SUCCEEDED(hr);
}); });
} }