78 lines
1.9 KiB
C++
78 lines
1.9 KiB
C++
#include "extra_protection/stubdrm.hpp"
|
|
|
|
#define WIN32_LEAN_AND_MEAN
|
|
#include <Windows.h>
|
|
|
|
#include <condition_variable>
|
|
#include <mutex>
|
|
#include <chrono>
|
|
#include <thread>
|
|
|
|
|
|
static std::mutex dll_unload_mtx{};
|
|
static std::condition_variable dll_unload_cv{};
|
|
static bool unload_dll = false;
|
|
|
|
static HMODULE my_hModule = nullptr;
|
|
static HANDLE unload_thread_handle = INVALID_HANDLE_VALUE;
|
|
|
|
|
|
static void send_unload_signal()
|
|
{
|
|
{
|
|
std::lock_guard lock(dll_unload_mtx);
|
|
unload_dll = true;
|
|
}
|
|
dll_unload_cv.notify_one();
|
|
}
|
|
|
|
DWORD WINAPI self_unload(LPVOID lpParameter)
|
|
{
|
|
constexpr const auto UNLOAD_TIMEOUT =
|
|
#ifdef _DEBUG
|
|
std::chrono::minutes(5)
|
|
#else
|
|
std::chrono::seconds(5)
|
|
#endif
|
|
;
|
|
|
|
{
|
|
std::unique_lock lock(dll_unload_mtx);
|
|
dll_unload_cv.wait_for(lock, UNLOAD_TIMEOUT, [](){ return unload_dll; });
|
|
}
|
|
unload_thread_handle = INVALID_HANDLE_VALUE;
|
|
FreeLibraryAndExitThread(my_hModule, 0);
|
|
}
|
|
|
|
BOOL APIENTRY DllMain(
|
|
HMODULE hModule,
|
|
DWORD reason,
|
|
LPVOID lpReserved)
|
|
{
|
|
switch (reason)
|
|
{
|
|
case DLL_PROCESS_ATTACH:
|
|
if (!stubdrm::patch()) {
|
|
// https://learn.microsoft.com/en-us/windows/win32/dlls/dllmain
|
|
// "The system immediately calls your entry-point function with DLL_PROCESS_DETACH and unloads the DLL"
|
|
unload_dll = true;
|
|
return FALSE;
|
|
}
|
|
my_hModule = hModule;
|
|
stubdrm::set_cleanup_cb(send_unload_signal);
|
|
unload_thread_handle = CreateThread(nullptr, 0, self_unload, nullptr, 0, nullptr);
|
|
break;
|
|
|
|
case DLL_PROCESS_DETACH:
|
|
if (!unload_dll) { // not unloaded yet, just an early exit, or thread timed out
|
|
stubdrm::restore();
|
|
if (unload_thread_handle != INVALID_HANDLE_VALUE && unload_thread_handle != NULL) {
|
|
TerminateThread(unload_thread_handle, 0);
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
|
|
return TRUE;
|
|
}
|