Use resumable leaf frames in CET hijack and in GC stress.

src/coreclr/inc/gcinfodecoder.h

@@ -530,6 +530,7 @@ class GcInfoDecoder
     bool IsSafePoint();
     bool AreSafePointsInterruptible();
     bool IsInterruptibleSafePoint();
+    bool CouldBeInterruptibleSafePoint();
 
     // This is used for gccoverage
     bool IsSafePoint(UINT32 codeOffset);

src/coreclr/vm/eetwain.cpp

@@ -1423,32 +1423,6 @@ bool EECodeManager::EnumGcRefs( PREGDISPLAY     pRD,
 
     GCInfoToken gcInfoToken = pCodeInfo->GetGCInfoToken();
 
-#if defined(STRESS_HEAP) && defined(PARTIALLY_INTERRUPTIBLE_GC_SUPPORTED)
-    // When we simulate a hijack during gcstress
-    //  we start with ActiveStackFrame and the offset
-    //  after the call
-    // We need to make it look like a non-leaf frame
-    //  so that it's treated like a regular hijack
-    if (flags & ActiveStackFrame)
-    {
-        GcInfoDecoder _gcInfoDecoder(
-                            gcInfoToken,
-                            DECODE_INTERRUPTIBILITY,
-                            curOffs
-                            );
-        if(!_gcInfoDecoder.IsInterruptible() &&
-            !(InterruptibleSafePointsEnabled() && _gcInfoDecoder.IsInterruptibleSafePoint()))
-        {
-            // This must be the offset after a call
-#ifdef _DEBUG
-            GcInfoDecoder _safePointDecoder(gcInfoToken, (GcInfoDecoderFlags)0, 0);
-            _ASSERTE(_safePointDecoder.IsSafePoint(curOffs));
-#endif
-            flags &= ~((unsigned)ActiveStackFrame);
-        }
-    }
-#endif // STRESS_HEAP && PARTIALLY_INTERRUPTIBLE_GC_SUPPORTED
-
 #ifdef _DEBUG
     if (flags & ActiveStackFrame)
     {
@@ -1458,7 +1432,7 @@ bool EECodeManager::EnumGcRefs( PREGDISPLAY     pRD,
                             curOffs
                             );
         _ASSERTE(_gcInfoDecoder.IsInterruptible() ||
-            (InterruptibleSafePointsEnabled() && _gcInfoDecoder.IsInterruptibleSafePoint()));
+            (InterruptibleSafePointsEnabled() && _gcInfoDecoder.CouldBeInterruptibleSafePoint()));
     }
 #endif
 
@@ -1561,7 +1535,7 @@ bool EECodeManager::EnumGcRefs( PREGDISPLAY     pRD,
                 curOffs - 1
             );
 
-            _ASSERTE((InterruptibleSafePointsEnabled() && gcInfoDecoder.IsInterruptibleSafePoint()));
+            _ASSERTE((InterruptibleSafePointsEnabled() && gcInfoDecoder.CouldBeInterruptibleSafePoint()));
         }
     }
 

src/coreclr/vm/excep.cpp

@@ -6770,39 +6770,33 @@ VEH_ACTION WINAPI CLRVectoredExceptionHandler(PEXCEPTION_POINTERS pExceptionInfo
             return VEH_CONTINUE_SEARCH;
         }
 
-        HijackArgs hijackArgs;
-        hijackArgs.Rax = pExceptionInfo->ContextRecord->Rax;
-        hijackArgs.Rsp = pExceptionInfo->ContextRecord->Rsp;
-
+        PCONTEXT interruptedContext = pExceptionInfo->ContextRecord;
         bool areShadowStacksEnabled = Thread::AreShadowStacksEnabled();
         if (areShadowStacksEnabled)
         {
-            // When the CET is enabled, the return address is still on stack, so we need to set the Rsp as
-            // if it was popped.
-            hijackArgs.Rsp += 8;
+            // OS should have fixed the SP value to the same as we`ve stashed for the hijacked thread
+            _ASSERTE(*(size_t *)interruptedContext->Rsp == (uintptr_t)pThread->GetHijackedReturnAddress());
+
+            // When the CET is enabled, the interruption happens on the ret instruction in the calee.
+            // We need to "pop" rsp to the caller, as if the ret has consumed it.
+            interruptedContext->Rsp += 8;
         }
-        hijackArgs.Rip = 0 ; // The OnHijackWorker sets this
-        #define CALLEE_SAVED_REGISTER(regname) hijackArgs.Regs.regname = pExceptionInfo->ContextRecord->regname;
-        ENUM_CALLEE_SAVED_REGISTERS();
-        #undef CALLEE_SAVED_REGISTER
 
-        OnHijackWorker(&hijackArgs);
+        // Change the IP to be at the original return site, as if we have returned to the caller.
+        // That IP is an interruptible safe point, so we can suspend right there.
+        uintptr_t origIp = interruptedContext->Rip;
+        interruptedContext->Rip = (uintptr_t)pThread->GetHijackedReturnAddress();
 
-        #define CALLEE_SAVED_REGISTER(regname) pExceptionInfo->ContextRecord->regname = hijackArgs.Regs.regname;
-        ENUM_CALLEE_SAVED_REGISTERS();
-        #undef CALLEE_SAVED_REGISTER
-        pExceptionInfo->ContextRecord->Rax = hijackArgs.Rax;
+        FrameWithCookie<ResumableFrame> frame(pExceptionInfo->ContextRecord);
+        frame.Push(pThread);
+        CommonTripThread();
+        frame.Pop(pThread);
 
         if (areShadowStacksEnabled)
         {
-            // The context refers to the return instruction
-            // Set the return address on the stack to the original one
-            *(size_t *)pExceptionInfo->ContextRecord->Rsp = hijackArgs.ReturnAddress;
-        }
-        else
-        {
-            // The context refers to the location after the return was processed
-            pExceptionInfo->ContextRecord->Rip = hijackArgs.ReturnAddress;
+            // Undo the "pop", so that the ret could now succeed.
+            interruptedContext->Rsp = interruptedContext->Rsp - 8;
+            interruptedContext->Rip = origIp;
         }
 
         return VEH_CONTINUE_EXECUTION;

src/coreclr/vm/frames.h

@@ -769,12 +769,13 @@ class Frame : public FrameBase
 
 
 //-----------------------------------------------------------------------------
-// This frame provides context for a frame that
-// took an exception that is going to be resumed.
+// This frame provides a context for a code location at which
+// execution was interrupted and may be resumed,
+// such as asynchronous suspension or handling of an exception.
 //
 // It is necessary to create this frame if garbage
-// collection may happen during handling of the
-// exception.  The FRAME_ATTR_RESUMABLE flag tells
+// collection may happen during the interruption.
+// The FRAME_ATTR_RESUMABLE flag tells
 // the GC that the preceding frame needs to be treated
 // like the top of stack (with the important implication that
 // caller-save-registers will be potential roots).
@@ -822,35 +823,15 @@ class ResumableFrame : public Frame
     }
 #endif
 
-protected:
-    PTR_CONTEXT m_Regs;
-
-    // Keep as last entry in class
-    DEFINE_VTABLE_GETTER_AND_CTOR_AND_DTOR(ResumableFrame)
-};
-
-
-//-----------------------------------------------------------------------------
-// RedirectedThreadFrame
-//-----------------------------------------------------------------------------
-
-class RedirectedThreadFrame : public ResumableFrame
-{
-    VPTR_VTABLE_CLASS(RedirectedThreadFrame, ResumableFrame)
-    VPTR_UNIQUE(VPTR_UNIQUE_RedirectedThreadFrame)
-
-public:
-#ifndef DACCESS_COMPILE
-    RedirectedThreadFrame(T_CONTEXT *regs) : ResumableFrame(regs) {
-        LIMITED_METHOD_CONTRACT;
-    }
-
-    virtual void ExceptionUnwind();
-#endif
-
     virtual void GcScanRoots(promote_func* fn, ScanContext* sc)
     {
         WRAPPER_NO_CONTRACT;
+
+        // The captured context may be provided by OS or by our own
+        // capture routine. The context may not necessary be on the
+        // stack or could be outside of reported stack range.
+        // To be sure that the registers in the context are reported
+        // in conservative root reporting, just report them here.
 #if defined(FEATURE_CONSERVATIVE_GC) && !defined(DACCESS_COMPILE)
         if (sc->promotion && g_pConfig->GetGCConservative())
         {
@@ -884,6 +865,32 @@ class RedirectedThreadFrame : public ResumableFrame
 #endif
     }
 
+protected:
+    PTR_CONTEXT m_Regs;
+
+    // Keep as last entry in class
+    DEFINE_VTABLE_GETTER_AND_CTOR_AND_DTOR(ResumableFrame)
+};
+
+
+//-----------------------------------------------------------------------------
+// RedirectedThreadFrame
+//-----------------------------------------------------------------------------
+
+class RedirectedThreadFrame : public ResumableFrame
+{
+    VPTR_VTABLE_CLASS(RedirectedThreadFrame, ResumableFrame)
+    VPTR_UNIQUE(VPTR_UNIQUE_RedirectedThreadFrame)
+
+public:
+#ifndef DACCESS_COMPILE
+    RedirectedThreadFrame(T_CONTEXT *regs) : ResumableFrame(regs) {
+        LIMITED_METHOD_CONTRACT;
+    }
+
+    virtual void ExceptionUnwind();
+#endif
+
     // Keep as last entry in class
     DEFINE_VTABLE_GETTER_AND_CTOR_AND_DTOR(RedirectedThreadFrame)
 };

src/coreclr/vm/gccover.cpp

@@ -1816,37 +1816,20 @@ void DoGcStress (PCONTEXT regs, NativeCodeVersion nativeCodeVersion)
         frame.Push(pThread);
     }
 
-    DWORD_PTR retValRegs[2] = { 0 };
+    // The legacy X86 GC encoder does not encode the state of return registers at
+    // call sites, so we must add an extra frame to protect returns.
+#ifdef TARGET_X86
+    DWORD_PTR retValRegs[1] = { 0 };
     UINT  numberOfRegs = 0;
 
     if (afterCallProtect[0])
     {
-#if defined(TARGET_AMD64)
-        retValRegs[numberOfRegs++] = regs->Rax;
-#elif defined(TARGET_X86)
         retValRegs[numberOfRegs++] = regs->Eax;
-#elif  defined(TARGET_ARM)
-        retValRegs[numberOfRegs++] = regs->R0;
-#elif defined(TARGET_ARM64)
-        retValRegs[numberOfRegs++] = regs->X0;
-#elif defined(TARGET_LOONGARCH64)
-        retValRegs[numberOfRegs++] = regs->A0;
-#elif defined(TARGET_RISCV64)
-        retValRegs[numberOfRegs++] = regs->A0;
-#endif // TARGET_ARM64
-    }
-
-    if (afterCallProtect[1])
-    {
-#if defined(TARGET_AMD64) && defined(TARGET_UNIX)
-        retValRegs[numberOfRegs++] = regs->Rdx;
-#else // !TARGET_AMD64 || !TARGET_UNIX
-        _ASSERTE(!"Not expected multi reg return with pointers.");
-#endif // !TARGET_AMD64 || !TARGET_UNIX
     }
 
     _ASSERTE(sizeof(OBJECTREF) == sizeof(DWORD_PTR));
     GCFrame gcFrame(pThread, (OBJECTREF*)retValRegs, numberOfRegs, TRUE);
+#endif
 
     MethodDesc *pMD = nativeCodeVersion.GetMethodDesc();
     LOG((LF_GCROOTS, LL_EVERYTHING, "GCCOVER: Doing GC at method %s::%s offset 0x%x\n",
@@ -1872,36 +1855,15 @@ void DoGcStress (PCONTEXT regs, NativeCodeVersion nativeCodeVersion)
 
     CONSISTENCY_CHECK(!pThread->HasPendingGCStressInstructionUpdate());
 
+#ifdef TARGET_X86
     if (numberOfRegs != 0)
     {
         if (afterCallProtect[0])
         {
-#if defined(TARGET_AMD64)
-            regs->Rax = retValRegs[0];
-#elif defined(TARGET_X86)
             regs->Eax = retValRegs[0];
-#elif defined(TARGET_ARM)
-            regs->R0 = retValRegs[0];
-#elif defined(TARGET_ARM64)
-            regs->X[0] = retValRegs[0];
-#elif defined(TARGET_LOONGARCH64)
-            regs->A0 = retValRegs[0];
-#elif defined(TARGET_RISCV64)
-            regs->A0 = retValRegs[0];
-#else
-            PORTABILITY_ASSERT("DoGCStress - return register");
-#endif
-        }
-
-        if (afterCallProtect[1])
-        {
-#if defined(TARGET_AMD64) && defined(TARGET_UNIX)
-            regs->Rdx = retValRegs[numberOfRegs - 1];
-#else // !TARGET_AMD64 || !TARGET_UNIX
-            _ASSERTE(!"Not expected multi reg return with pointers.");
-#endif // !TARGET_AMD64 || !TARGET_UNIX
         }
     }
+#endif
 
     if (!Thread::UseRedirectForGcStress())
     {

src/coreclr/vm/gcinfodecoder.cpp

@@ -425,6 +425,16 @@ bool GcInfoDecoder::IsInterruptibleSafePoint()
     return IsSafePoint() && AreSafePointsInterruptible();
 }
 
+bool GcInfoDecoder::CouldBeInterruptibleSafePoint()
+{
+    // This is used in asserts. Ideally it would return false
+    // if current location canot possibly be a safepoint.
+    // However we optimize away "boring" callsites when no variables are tracked.
+    // So there is no way to tell currently tat a pint is indeed not a safe point.
+    // Thus we do what we can here, but this could be better if we could have more data
+    return AreSafePointsInterruptible() && m_NumInterruptibleRanges == 0;
+}
+
 bool GcInfoDecoder::HasMethodDescGenericsInstContext()
 {
     _ASSERTE( m_Flags & DECODE_GENERICS_INST_CONTEXT );

src/coreclr/vm/threads.h

@@ -3362,6 +3362,13 @@ class Thread
         return *retAddrSlot;
     }
 
+#ifdef FEATURE_HIJACK
+    void* GetHijackedReturnAddress()
+    {
+        return m_pvHJRetAddr;
+    }
+#endif
+
 #ifdef _DEBUG
 private:
     // When we create an object, or create an OBJECTREF, or create an Interior Pointer, or enter EE from managed