mirror of
https://github.com/snowie2000/mactype.git
synced 2025-01-05 10:17:02 +08:00
112 lines
3.3 KiB
C++
112 lines
3.3 KiB
C++
#include "dynCodeHelper.h"
|
|
|
|
/*
|
|
* class AutoEnableDynamicCodeGen
|
|
*/
|
|
|
|
typedef
|
|
BOOL
|
|
(WINAPI *PGET_PROCESS_MITIGATION_POLICY_PROC)(
|
|
_In_ HANDLE hProcess,
|
|
_In_ PROCESS_MITIGATION_POLICY MitigationPolicy,
|
|
_Out_ PVOID lpBuffer,
|
|
_In_ SIZE_T dwLength
|
|
);
|
|
|
|
AutoEnableDynamicCodeGen::PSET_THREAD_INFORMATION_PROC AutoEnableDynamicCodeGen::SetThreadInformationProc = nullptr;
|
|
AutoEnableDynamicCodeGen::PGET_THREAD_INFORMATION_PROC AutoEnableDynamicCodeGen::GetThreadInformationProc = nullptr;
|
|
PROCESS_MITIGATION_DYNAMIC_CODE_POLICY AutoEnableDynamicCodeGen::processPolicy;
|
|
volatile bool AutoEnableDynamicCodeGen::processPolicyObtained = false;
|
|
|
|
AutoEnableDynamicCodeGen::AutoEnableDynamicCodeGen(bool enable) : enabled(false)
|
|
{
|
|
if (enable == false)
|
|
{
|
|
return;
|
|
}
|
|
|
|
//
|
|
// Snap the dynamic code generation policy for this process so that we
|
|
// don't need to resolve APIs and query it each time. We expect the policy
|
|
// to have been established upfront.
|
|
//
|
|
|
|
if (processPolicyObtained == false)
|
|
{
|
|
CCriticalSectionLock __lock(CCriticalSectionLock::CS_VIRTMEM);
|
|
|
|
if (processPolicyObtained == false)
|
|
{
|
|
PGET_PROCESS_MITIGATION_POLICY_PROC GetProcessMitigationPolicyProc = nullptr;
|
|
|
|
HMODULE module = GetModuleHandleW(_T("api-ms-win-core-processthreads-l1-1-3.dll"));
|
|
|
|
if (module != nullptr)
|
|
{
|
|
GetProcessMitigationPolicyProc = (PGET_PROCESS_MITIGATION_POLICY_PROC)GetProcAddress(module, "GetProcessMitigationPolicy");
|
|
SetThreadInformationProc = (PSET_THREAD_INFORMATION_PROC)GetProcAddress(module, "SetThreadInformation");
|
|
GetThreadInformationProc = (PGET_THREAD_INFORMATION_PROC)GetProcAddress(module, "GetThreadInformation");
|
|
}
|
|
|
|
if ((GetProcessMitigationPolicyProc == nullptr) ||
|
|
(!GetProcessMitigationPolicyProc(GetCurrentProcess(), ProcessDynamicCodePolicy, (PPROCESS_MITIGATION_DYNAMIC_CODE_POLICY)&processPolicy, sizeof(processPolicy))))
|
|
{
|
|
processPolicy.ProhibitDynamicCode = 0;
|
|
}
|
|
|
|
processPolicyObtained = true;
|
|
}
|
|
}
|
|
|
|
//
|
|
// The process is not prohibiting dynamic code or does not allow threads
|
|
// to opt out. In either case, return to the caller.
|
|
//
|
|
// N.B. It is OK that this policy is mutable at runtime. If a process
|
|
// really does not allow thread opt-out, then the call below will fail
|
|
// benignly.
|
|
//
|
|
|
|
if ((processPolicy.ProhibitDynamicCode == 0) || (processPolicy.AllowThreadOptOut == 0))
|
|
{
|
|
return;
|
|
}
|
|
|
|
if (SetThreadInformationProc == nullptr || GetThreadInformationProc == nullptr)
|
|
{
|
|
return;
|
|
}
|
|
|
|
//
|
|
// If dynamic code is already allowed for this thread, then don't attempt to allow it again.
|
|
//
|
|
|
|
DWORD threadPolicy;
|
|
|
|
if ((GetThreadInformationProc(GetCurrentThread(), ThreadDynamicCodePolicy, &threadPolicy, sizeof(DWORD))) &&
|
|
(threadPolicy == THREAD_DYNAMIC_CODE_ALLOW))
|
|
{
|
|
return;
|
|
}
|
|
|
|
threadPolicy = THREAD_DYNAMIC_CODE_ALLOW;
|
|
|
|
BOOL result = SetThreadInformationProc(GetCurrentThread(), ThreadDynamicCodePolicy, &threadPolicy, sizeof(DWORD));
|
|
Assert(result);
|
|
|
|
enabled = true;
|
|
}
|
|
|
|
AutoEnableDynamicCodeGen::~AutoEnableDynamicCodeGen()
|
|
{
|
|
if (enabled)
|
|
{
|
|
DWORD threadPolicy = 0;
|
|
|
|
BOOL result = SetThreadInformationProc(GetCurrentThread(), ThreadDynamicCodePolicy, &threadPolicy, sizeof(DWORD));
|
|
Assert(result);
|
|
|
|
enabled = false;
|
|
}
|
|
}
|