[NativeAOT/x86] Load RtlRestoreContext dynamically and add fallback for old OSes (#99813)

* Load RtlRestoreContext dynamically and add fallback for old OSes

* Port the RtlRestoreContext SEH fallback logic from CoreCLR

* Remove unnecessary change

* Remove ARM32 mention from a comment

* Simplify implementations of GetCurrentSEHRecord, SetCurrentSEHRecord, and PopSEHRecords. Use the same implementation logic in CoreCLR and NativeAOT

* Update comments

Author: filipnavara

src/coreclr/nativeaot/Runtime/thread.cpp

@@ -1107,6 +1107,44 @@ void Thread::SetActivationPending(bool isPending)
     }
 }
 
+#ifdef TARGET_X86
+
+void Thread::SetPendingRedirect(PCODE eip)
+{
+    m_LastRedirectIP = eip;
+    m_SpinCount = 0;
+}
+
+bool Thread::CheckPendingRedirect(PCODE eip)
+{
+    if (eip == m_LastRedirectIP)
+    {
+        // We need to test for an infinite loop in assembly, as this will break the heuristic we
+        // are using.
+        const BYTE short_jmp = 0xeb;    // Machine code for a short jump.
+        const BYTE self = 0xfe;         // -2.  Short jumps are calculated as [ip]+2+[second_byte].
+
+        // If we find that we are in an infinite loop, we'll set the last redirected IP to 0 so that we will
+        // redirect the next time we attempt it.  Delaying one interation allows us to narrow the window of
+        // the race we are working around in this corner case.
+        BYTE *ip = (BYTE *)m_LastRedirectIP;
+        if (ip[0] == short_jmp && ip[1] == self)
+            m_LastRedirectIP = 0;
+
+        // We set a hard limit of 5 times we will spin on this to avoid any tricky race which we have not
+        // accounted for.
+        m_SpinCount++;
+        if (m_SpinCount >= 5)
+            m_LastRedirectIP = 0;
+
+        return true;
+    }
+
+    return false;
+}
+
+#endif // TARGET_X86
+
 #endif // !DACCESS_COMPILE
 
 void Thread::ValidateExInfoStack()

src/coreclr/nativeaot/Runtime/thread.h

@@ -97,6 +97,10 @@ struct ThreadBuffer
     uintptr_t               m_uHijackedReturnValueFlags;
     PTR_ExInfo              m_pExInfoStackHead;
     Object*                 m_threadAbortException;                 // ThreadAbortException instance -set only during thread abort
+#ifdef TARGET_X86
+    PCODE                   m_LastRedirectIP;
+    uint64_t                m_SpinCount;
+#endif
     Object*                 m_pThreadLocalStatics;
     InlinedThreadStaticRoot* m_pInlinedThreadLocalStatics;
     GCFrameRegistration*    m_pGCFrameRegistrations;
@@ -317,6 +321,11 @@ class Thread : private ThreadBuffer
 
     bool                IsActivationPending();
     void                SetActivationPending(bool isPending);
+
+#ifdef TARGET_X86
+    void                SetPendingRedirect(PCODE eip);
+    bool                CheckPendingRedirect(PCODE eip);
+#endif
 };
 
 #ifndef __GCENV_BASE_INCLUDED__

src/coreclr/nativeaot/Runtime/windows/PalRedhawkMinWin.cpp

@@ -28,6 +28,7 @@
 #include "gcconfig.h"
 
 #include "thread.h"
+#include "threadstore.h"
 
 #define REDHAWK_PALEXPORT extern "C"
 #define REDHAWK_PALAPI __stdcall
@@ -322,10 +323,120 @@ REDHAWK_PALEXPORT HANDLE REDHAWK_PALAPI PalCreateEventW(_In_opt_ LPSECURITY_ATTR
     return CreateEventW(pEventAttributes, manualReset, initialState, pName);
 }
 
+#ifdef TARGET_X86
+
+#define EXCEPTION_HIJACK  0xe0434f4e    // 0xe0000000 | 'COM'+1
+
+PEXCEPTION_REGISTRATION_RECORD GetCurrentSEHRecord()
+{
+    return (PEXCEPTION_REGISTRATION_RECORD)__readfsdword(0);
+}
+
+VOID SetCurrentSEHRecord(EXCEPTION_REGISTRATION_RECORD *pSEH)
+{
+    __writefsdword(0, (DWORD)pSEH);
+}
+
+VOID PopSEHRecords(LPVOID pTargetSP)
+{
+    PEXCEPTION_REGISTRATION_RECORD currentContext = GetCurrentSEHRecord();
+    // The last record in the chain is EXCEPTION_CHAIN_END which is defined as maxiumum
+    // pointer value so it cannot satisfy the loop condition.
+    while (currentContext < pTargetSP)
+    {
+        currentContext = currentContext->Next;
+    }
+    SetCurrentSEHRecord(currentContext);
+}
+
+// This will check who caused the exception.  If it was caused by the redirect function,
+// the reason is to resume the thread back at the point it was redirected in the first
+// place.  If the exception was not caused by the function, then it was caused by the call
+// out to the I[GC|Debugger]ThreadControl client and we need to determine if it's an
+// exception that we can just eat and let the runtime resume the thread, or if it's an
+// uncatchable exception that we need to pass on to the runtime.
+int RtlRestoreContextFallbackExceptionFilter(PEXCEPTION_POINTERS pExcepPtrs, CONTEXT *pCtx, Thread *pThread)
+{
+    if (pExcepPtrs->ExceptionRecord->ExceptionCode == STATUS_STACK_OVERFLOW)
+    {
+        return EXCEPTION_CONTINUE_SEARCH;
+    }
+
+    // Get the thread handle
+    _ASSERTE(pExcepPtrs->ExceptionRecord->ExceptionCode == EXCEPTION_HIJACK);
+
+    // Copy everything in the saved context record into the EH context.
+    // Historically the EH context has enough space for every enabled context feature.
+    // That may not hold for the future features beyond AVX, but this codepath is
+    // supposed to be used only on OSes that do not have RtlRestoreContext.
+    CONTEXT* pTarget = pExcepPtrs->ContextRecord;
+    if (!CopyContext(pTarget, pCtx->ContextFlags, pCtx))
+    {
+        PalPrintFatalError("Could not set context record.\n");
+        RhFailFast();
+    }
+
+    DWORD espValue = pCtx->Esp;
+
+    // NOTE: Ugly, ugly workaround.
+    // We need to resume the thread into the managed code where it was redirected,
+    // and the corresponding ESP is below the current one.  But C++ expects that
+    // on an EXCEPTION_CONTINUE_EXECUTION that the ESP will be above where it has
+    // installed the SEH handler.  To solve this, we need to remove all handlers
+    // that reside above the resumed ESP, but we must leave the OS-installed
+    // handler at the top, so we grab the top SEH handler, call
+    // PopSEHRecords which will remove all SEH handlers above the target ESP and
+    // then link the OS handler back in with SetCurrentSEHRecord.
+
+    // Get the special OS handler and save it until PopSEHRecords is done
+    EXCEPTION_REGISTRATION_RECORD *pCurSEH = GetCurrentSEHRecord();
+
+    // Unlink all records above the target resume ESP
+    PopSEHRecords((LPVOID)(size_t)espValue);
+
+    // Link the special OS handler back in to the top
+    pCurSEH->Next = GetCurrentSEHRecord();
+
+    // Register the special OS handler as the top handler with the OS
+    SetCurrentSEHRecord(pCurSEH);
+
+    // Resume execution at point where thread was originally redirected
+    return EXCEPTION_CONTINUE_EXECUTION;
+}
+
+EXTERN_C VOID __cdecl RtlRestoreContextFallback(PCONTEXT ContextRecord, struct _EXCEPTION_RECORD* ExceptionRecord)
+{
+    Thread *pThread = ThreadStore::GetCurrentThread();
+
+    // A counter to avoid a nasty case where an
+    // up-stack filter throws another exception
+    // causing our filter to be run again for
+    // some unrelated exception.
+    int filter_count = 0;
+
+    __try
+    {
+        // Save the instruction pointer where we redirected last.  This does not race with the check
+        // against this variable because the GC will not attempt to redirect the thread until the
+        // instruction pointer of this thread is back in managed code.
+        pThread->SetPendingRedirect(ContextRecord->Eip);
+        RaiseException(EXCEPTION_HIJACK, 0, 0, NULL);
+    }
+    __except (++filter_count == 1
+            ? RtlRestoreContextFallbackExceptionFilter(GetExceptionInformation(), ContextRecord, pThread)
+            : EXCEPTION_CONTINUE_SEARCH)
+    {
+        _ASSERTE(!"Reached body of __except in RtlRestoreContextFallback");
+    }
+}
+
+#endif // TARGET_X86
+
 typedef BOOL(WINAPI* PINITIALIZECONTEXT2)(PVOID Buffer, DWORD ContextFlags, PCONTEXT* Context, PDWORD ContextLength, ULONG64 XStateCompactionMask);
 PINITIALIZECONTEXT2 pfnInitializeContext2 = NULL;
 
 #ifdef TARGET_X86
+EXTERN_C VOID __cdecl RtlRestoreContextFallback(PCONTEXT ContextRecord, struct _EXCEPTION_RECORD* ExceptionRecord);
 typedef VOID(__cdecl* PRTLRESTORECONTEXT)(PCONTEXT ContextRecord, struct _EXCEPTION_RECORD* ExceptionRecord);
 PRTLRESTORECONTEXT pfnRtlRestoreContext = NULL;
 
@@ -356,6 +467,11 @@ REDHAWK_PALEXPORT CONTEXT* PalAllocateCompleteOSContext(_Out_ uint8_t** contextB
     {
         HMODULE hm = GetModuleHandleW(_T("ntdll.dll"));
         pfnRtlRestoreContext = (PRTLRESTORECONTEXT)GetProcAddress(hm, "RtlRestoreContext");
+        if (pfnRtlRestoreContext == NULL)
+        {
+            // Fallback to the internal implementation if OS doesn't provide one.
+            pfnRtlRestoreContext = RtlRestoreContextFallback;
+        }
     }
 #endif //TARGET_X86
 
@@ -438,7 +554,12 @@ REDHAWK_PALEXPORT _Success_(return) bool REDHAWK_PALAPI PalSetThreadContext(HAND
 REDHAWK_PALEXPORT void REDHAWK_PALAPI PalRestoreContext(CONTEXT * pCtx)
 {
     __asan_handle_no_return();
+#ifdef TARGET_X86
+    _ASSERTE(pfnRtlRestoreContext != NULL);
+    pfnRtlRestoreContext(pCtx, NULL);
+#else
     RtlRestoreContext(pCtx, NULL);
+#endif //TARGET_X86
 }
 
 REDHAWK_PALIMPORT void REDHAWK_PALAPI PopulateControlSegmentRegisters(CONTEXT* pContext)
@@ -568,16 +689,41 @@ REDHAWK_PALEXPORT void REDHAWK_PALAPI PalHijack(HANDLE hThread, _In_opt_ void* p
 
     if (GetThreadContext(hThread, &win32ctx))
     {
+        bool isSafeToRedirect = true;
+
+#ifdef TARGET_X86
+        // Workaround around WOW64 problems. Only do this workaround if a) this is x86, and b) the OS does
+        // not support trap frame reporting.
+        if ((win32ctx.ContextFlags & CONTEXT_EXCEPTION_REPORTING) == 0)
+        {
+            // This code fixes a race between GetThreadContext and NtContinue.  If we redirect managed code
+            // at the same place twice in a row, we run the risk of reading a bogus CONTEXT when we redirect
+            // the second time.  This leads to access violations on x86 machines.  To fix the problem, we
+            // never redirect at the same instruction pointer that we redirected at on the previous GC.
+            if (((Thread*)pThreadToHijack)->CheckPendingRedirect(win32ctx.Eip))
+            {
+                isSafeToRedirect = false;
+            }
+        }
+#else
+        // In some cases Windows will not set the CONTEXT_EXCEPTION_REPORTING flag if the thread is executing
+        // in kernel mode (i.e. in the middle of a syscall or exception handling). Therefore, we should treat
+        // the absence of the CONTEXT_EXCEPTION_REPORTING flag as an indication that it is not safe to
+        // manipulate with the current state of the thread context.
+        isSafeToRedirect = (win32ctx.ContextFlags & CONTEXT_EXCEPTION_REPORTING) != 0;
+#endif
+
         // The CONTEXT_SERVICE_ACTIVE and CONTEXT_EXCEPTION_ACTIVE output flags indicate we suspended the thread
         // at a point where the kernel cannot guarantee a completely accurate context. We'll fail the request in
         // this case (which should force our caller to resume the thread and try again -- since this is a fairly
         // narrow window we're highly likely to succeed next time).
-        // Note: in some cases (x86 WOW64, ARM32 on ARM64) the OS will not set the CONTEXT_EXCEPTION_REPORTING flag
-        // if the thread is executing in kernel mode (i.e. in the middle of a syscall or exception handling).
-        // Therefore, we should treat the absence of the CONTEXT_EXCEPTION_REPORTING flag as an indication that
-        // it is not safe to manipulate with the current state of the thread context.
         if ((win32ctx.ContextFlags & CONTEXT_EXCEPTION_REPORTING) != 0 &&
-            ((win32ctx.ContextFlags & (CONTEXT_SERVICE_ACTIVE | CONTEXT_EXCEPTION_ACTIVE)) == 0))
+            ((win32ctx.ContextFlags & (CONTEXT_SERVICE_ACTIVE | CONTEXT_EXCEPTION_ACTIVE)) != 0))
+        {
+            isSafeToRedirect = false;
+        }
+
+        if (isSafeToRedirect)
         {
             g_pHijackCallback(&win32ctx, pThreadToHijack);
         }

src/coreclr/vm/excep.h

@@ -518,7 +518,7 @@ EXCEPTION_HANDLER_DECL(COMPlusFrameHandlerRevCom);
 #endif // FEATURE_COMINTEROP
 
 // Pop off any SEH handlers we have registered below pTargetSP
-VOID __cdecl PopSEHRecords(LPVOID pTargetSP);
+VOID PopSEHRecords(LPVOID pTargetSP);
 
 #ifdef DEBUGGING_SUPPORTED
 VOID UnwindExceptionTrackerAndResumeInInterceptionFrame(ExInfo* pExInfo, EHContext* context);

src/coreclr/vm/i386/excepx86.cpp

@@ -1849,39 +1849,7 @@ PEXCEPTION_REGISTRATION_RECORD GetCurrentSEHRecord()
 {
     WRAPPER_NO_CONTRACT;
 
-    LPVOID fs0 = (LPVOID)__readfsdword(0);
-
-#if 0  // This walk is too expensive considering we hit it every time we a CONTRACT(NOTHROW)
-#ifdef _DEBUG
-    EXCEPTION_REGISTRATION_RECORD *pEHR = (EXCEPTION_REGISTRATION_RECORD *)fs0;
-    LPVOID spVal;
-    __asm {
-        mov spVal, esp
-    }
-
-    // check that all the eh frames are all greater than the current stack value. If not, the
-    // stack has been updated somehow w/o unwinding the SEH chain.
-
-    // LOG((LF_EH, LL_INFO1000000, "ER Chain:\n"));
-    while (pEHR != NULL && pEHR != EXCEPTION_CHAIN_END) {
-        // LOG((LF_EH, LL_INFO1000000, "\tp: prev:p handler:%x\n", pEHR, pEHR->Next, pEHR->Handler));
-        if (pEHR < spVal) {
-            if (gLastResumedExceptionFunc != 0)
-                _ASSERTE(!"Stack is greater than start of SEH chain - possible missing leave in handler. See gLastResumedExceptionHandler & gLastResumedExceptionFunc for info");
-            else
-                _ASSERTE(!"Stack is greater than start of SEH chain (FS:0)");
-        }
-        if (pEHR->Handler == (void *)-1)
-            _ASSERTE(!"Handler value has been corrupted");
-
-            _ASSERTE(pEHR < pEHR->Next);
-
-        pEHR = pEHR->Next;
-    }
-#endif
-#endif // 0
-
-    return (EXCEPTION_REGISTRATION_RECORD*) fs0;
+    return (PEXCEPTION_REGISTRATION_RECORD)__readfsdword(0);
 }
 
 PEXCEPTION_REGISTRATION_RECORD GetFirstCOMPlusSEHRecord(Thread *pThread) {
@@ -1917,29 +1885,23 @@ PEXCEPTION_REGISTRATION_RECORD GetPrevSEHRecord(EXCEPTION_REGISTRATION_RECORD *n
 VOID SetCurrentSEHRecord(EXCEPTION_REGISTRATION_RECORD *pSEH)
 {
     WRAPPER_NO_CONTRACT;
-    *GetThread()->GetExceptionListPtr() = pSEH;
+
+    __writefsdword(0, (DWORD)pSEH);
 }
 
-// Note that this logic is copied below, in PopSEHRecords
-__declspec(naked)
-VOID __cdecl PopSEHRecords(LPVOID pTargetSP)
+VOID PopSEHRecords(LPVOID pTargetSP)
 {
-    // No CONTRACT possible on naked functions
     STATIC_CONTRACT_NOTHROW;
     STATIC_CONTRACT_GC_NOTRIGGER;
 
-    __asm{
-        mov     ecx, [esp+4]        ;; ecx <- pTargetSP
-        mov     eax, fs:[0]         ;; get current SEH record
-  poploop:
-        cmp     eax, ecx
-        jge     done
-        mov     eax, [eax]          ;; get next SEH record
-        jmp     poploop
-  done:
-        mov     fs:[0], eax
-        retn
+    PEXCEPTION_REGISTRATION_RECORD currentContext = GetCurrentSEHRecord();
+    // The last record in the chain is EXCEPTION_CHAIN_END which is defined as maxiumum
+    // pointer value so it cannot satisfy the loop condition.
+    while (currentContext < pTargetSP)
+    {
+        currentContext = currentContext->Next;
     }
+    SetCurrentSEHRecord(currentContext);
 }
 
 //

src/coreclr/vm/threadsuspend.cpp

@@ -1117,11 +1117,11 @@ BOOL Thread::IsContextSafeToRedirect(const CONTEXT* pContext)
 #ifndef TARGET_UNIX
 
 #if !defined(TARGET_X86)
-    // In some cases (x86 WOW64, ARM32 on ARM64) Windows will not set the CONTEXT_EXCEPTION_REPORTING flag
-    // if the thread is executing in kernel mode (i.e. in the middle of a syscall or exception handling).
-    // Therefore, we should treat the absence of the CONTEXT_EXCEPTION_REPORTING flag as an indication that
-    // it is not safe to manipulate with the current state of the thread context.
-    // Note: the x86 WOW64 case is already handled in GetSafelyRedirectableThreadContext; in addition, this
+    // In some cases Windows will not set the CONTEXT_EXCEPTION_REPORTING flag if the thread is executing
+    // in kernel mode (i.e. in the middle of a syscall or exception handling). Therefore, we should treat
+    // the absence of the CONTEXT_EXCEPTION_REPORTING flag as an indication that it is not safe to
+    // manipulate with the current state of the thread context.
+    // Note: The x86 WOW64 case is already handled in GetSafelyRedirectableThreadContext; in addition, this
     // flag is never set on Windows7 x86 WOW64. So this check is valid for non-x86 architectures only.
     isSafeToRedirect = (pContext->ContextFlags & CONTEXT_EXCEPTION_REPORTING) != 0;
 #endif // !defined(TARGET_X86)