Enable JIT clamping

lib/Parser/rterrors.h

@@ -331,5 +331,6 @@ RT_ERROR_MSG(JSERR_SimdInt16x8TypeMismatch, 5636, "SIMD.Int16x8.%s: Invalid SIMD
 RT_ERROR_MSG(JSERR_SimdUint32x4TypeMismatch, 5637, "SIMD.Uint32x4.%s: Invalid SIMD types for operation", "Expecting UInt32x4 values", kjstTypeError, 0)
 RT_ERROR_MSG(JSERR_SimdUint16x8TypeMismatch, 5638, "SIMD.Uint16x8.%s: Invalid SIMD types for operation", "Expecting UInt16x8 values", kjstTypeError, 0)
 RT_ERROR_MSG(JSERR_SimdUint8x16TypeMismatch, 5639, "SIMD.Uint8x16.%s: Invalid SIMD types for operation", "Expecting UInt8x16 values", kjstTypeError, 0)
+
 RT_ERROR_MSG(JSERR_FunctionArgument_FirstCannotBeRegExp, 5640, "%s: first argument cannot be a RegExp", "First argument cannot be a RegExp", kjstTypeError, 0)
 RT_ERROR_MSG(JSERR_RegExpExecInvalidReturnType, 5641, "%s: Return value of RegExp 'exec' is not an Object and is not null", "Return value of RegExp 'exec' is not an Object and is not null", kjstTypeError, 0)

lib/Runtime/Base/DelayLoadLibrary.cpp

@@ -401,6 +401,12 @@ namespace Js
         _In_ ULONG NumberOfOffsets,
         _In_reads_(NumberOfOffsets) PCFG_CALL_TARGET_INFO OffsetInformation)
     {
+#if defined(ENABLE_JIT_CLAMP)
+        // Ensure that dynamic code generation is allowed for this thread as
+        // this is required for the call to SetProcessValidCallTargets to
+        // succeed.
+        AutoEnableDynamicCodeGen enableCodeGen;
+#endif
 
 #if defined(DELAYLOAD_SET_CFG_TARGET)
         if (m_hModule)

lib/common/CommonDefines.h

@@ -154,6 +154,7 @@
 #define ENABLE_DOM_FAST_PATH
 #define ENABLE_JS_ETW                               // ETW support
 #define EDIT_AND_CONTINUE
+#define ENABLE_JIT_CLAMP
 #endif
 
 // Telemetry flags

lib/common/Memory/PageAllocator.cpp

@@ -2239,6 +2239,12 @@ HeapPageAllocator<T>::ProtectPages(__in char* address, size_t pageCount, __in vo
     }
     *dwOldVirtualProtectFlags = memBasicInfo.Protect;
 
+#if defined(ENABLE_JIT_CLAMP)
+    bool makeExecutable = (dwVirtualProtectFlags & (PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE)) ? true : false;
+
+    AutoEnableDynamicCodeGen enableCodeGen(makeExecutable);
+#endif
+
     DWORD oldProtect; // this is only for first page
     BOOL retVal = ::VirtualProtect(address, pageCount * AutoSystemInfo::PageSize, dwVirtualProtectFlags, &oldProtect);
     Assert(oldProtect == *dwOldVirtualProtectFlags);

lib/common/Memory/VirtualAllocWrapper.cpp

@@ -13,6 +13,22 @@ LPVOID VirtualAllocWrapper::Alloc(LPVOID lpAddress, size_t dwSize, DWORD allocat
     Assert(this == nullptr);
     LPVOID address = nullptr;
 
+#if defined(ENABLE_JIT_CLAMP)
+    bool makeExecutable;
+
+    if ((isCustomHeapAllocation) || 
+        (protectFlags & (PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE)))
+    {
+        makeExecutable = true;
+    }
+    else
+    {
+        makeExecutable = false;
+    }
+
+    AutoEnableDynamicCodeGen enableCodeGen(makeExecutable);
+#endif
+
 #if defined(_CONTROL_FLOW_GUARD)
     DWORD oldProtectFlags;
     if (AutoSystemInfo::Data.IsCFGEnabled() && isCustomHeapAllocation)
@@ -27,6 +43,7 @@ LPVOID VirtualAllocWrapper::Alloc(LPVOID lpAddress, size_t dwSize, DWORD allocat
     {
         address = VirtualAlloc(lpAddress, dwSize, allocationType, protectFlags);
     }
+
     return address;
 }
 
@@ -65,6 +82,7 @@ VirtualAllocWrapper::IsInRange(void * address)
     return false;
 }
 
+
 /*
 * class PreReservedVirtualAllocWrapper
 */
@@ -236,6 +254,11 @@ LPVOID PreReservedVirtualAllocWrapper::Alloc(LPVOID lpAddress, size_t dwSize, DW
         AssertMsg(dwSize % AutoSystemInfo::PageSize == 0, "COMMIT is managed at AutoSystemInfo::PageSize granularity");
 
         char * committedAddress = nullptr;
+
+#if defined(ENABLE_JIT_CLAMP)
+        AutoEnableDynamicCodeGen enableCodeGen;
+#endif
+
 #if defined(_CONTROL_FLOW_GUARD)
         if (AutoSystemInfo::Data.IsCFGEnabled())
         {
@@ -251,7 +274,6 @@ LPVOID PreReservedVirtualAllocWrapper::Alloc(LPVOID lpAddress, size_t dwSize, DW
             committedAddress = (char *) VirtualAlloc(addressToCommit, dwSize, MEM_COMMIT, protectFlags);
         }
 
-
         //Keep track of the committed pages within the preReserved Memory Region
         if (lpAddress == nullptr && committedAddress != nullptr)
         {
@@ -317,3 +339,118 @@ PreReservedVirtualAllocWrapper::Free(LPVOID lpAddress, size_t dwSize, DWORD dwFr
         return success;
     }
 }
+
+#if defined(ENABLE_JIT_CLAMP)
+/*
+* 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;
+CriticalSection AutoEnableDynamicCodeGen::processPolicyCS;
+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)
+    {
+        AutoCriticalSection autocs(&processPolicyCS);
+
+        if (processPolicyObtained == false)
+        {
+            PGET_PROCESS_MITIGATION_POLICY_PROC GetProcessMitigationPolicyProc = nullptr;
+
+            HMODULE module = GetModuleHandleW(L"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;
+    }
+}
+
+#endif // defined(ENABLE_JIT_CLAMP)
+

lib/common/Memory/VirtualAllocWrapper.h

@@ -55,6 +55,43 @@ private:
     CriticalSection                                 cs;
 };
 
+#if defined(ENABLE_JIT_CLAMP)
+
+class AutoEnableDynamicCodeGen
+{
+public:
+    AutoEnableDynamicCodeGen(bool enable = true);
+    ~AutoEnableDynamicCodeGen();
+
+private:
+    bool enabled;
+
+    typedef
+    BOOL
+    (WINAPI *PSET_THREAD_INFORMATION_PROC)(
+        _In_ HANDLE                   hThread,
+        _In_ THREAD_INFORMATION_CLASS ThreadInformationClass,
+        _In_reads_bytes_(ThreadInformationSize) PVOID ThreadInformation,
+        _In_ DWORD                    ThreadInformationSize
+    );
+
+    typedef
+    BOOL
+    (WINAPI *PGET_THREAD_INFORMATION_PROC)(
+        _In_ HANDLE                   hThread,
+        _In_ THREAD_INFORMATION_CLASS ThreadInformationClass,
+        _Out_writes_bytes_(ThreadInformationSize) PVOID ThreadInformation,
+        _In_ DWORD                    ThreadInformationSize
+    );
+
+    static PSET_THREAD_INFORMATION_PROC SetThreadInformationProc;
+    static PGET_THREAD_INFORMATION_PROC GetThreadInformationProc;
+    static PROCESS_MITIGATION_DYNAMIC_CODE_POLICY processPolicy;
+    static CriticalSection processPolicyCS;
+    static volatile bool processPolicyObtained;
+};
+#endif
+
 template<typename TVirtualAlloc = VirtualAllocWrapper>
 class PageAllocatorBase;