[release/9.0-staging] Fix return address hijacking with CET (#109548)

* Fix return address hijacking with CET

There is a problematic case when return address is hijacked while in a
managed method that tail calls a GC write barrier and when CET is
enabled. The write barrier code can change while the handler for the
hijacked address is executed from the vectored exception handler.
When the vectored exception handler then returns to the write barrier to
re-execute the `ret` instruction that has triggered the vectored
exception handler due to the main stack containing a different address
than the shadow stack (now with the main stack fixed), the instruction
may no longer be `ret` due to the change of the write barrier change.

This change fixes it by setting the context to return to from the
vectored exception handler to point to the caller and setting the Rsp
and SSP to match that. That way, the write barrier code no longer
matters.

* Add equivalent change to nativeaot

* Add missing ifdef

---------

Co-authored-by: Jan Vorlicek (from Dev Box) <janvorli@microsoft.com>

src/coreclr/nativeaot/Runtime/EHHelpers.cpp

@@ -485,6 +485,9 @@ int32_t __stdcall RhpHardwareExceptionHandler(uintptr_t faultCode, uintptr_t fau
 
 #else // TARGET_UNIX
 
+uintptr_t GetSSP(CONTEXT *pContext);
+void SetSSP(CONTEXT *pContext, uintptr_t ssp);
+
 static bool g_ContinueOnFatalErrors = false;
 
 // Set the runtime to continue search when encountering an unhandled runtime exception. Once done it is forever.
@@ -539,22 +542,16 @@ int32_t __stdcall RhpVectoredExceptionHandler(PEXCEPTION_POINTERS pExPtrs)
             // 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->SetSp(interruptedContext->GetSp() + 8);
+            uintptr_t ssp = GetSSP(interruptedContext);
+            SetSSP(interruptedContext, ssp + 8);
         }
 
         // 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->GetIp();
         interruptedContext->SetIp((uintptr_t)pThread->GetHijackedReturnAddress());
 
         pThread->InlineSuspend(interruptedContext);
 
-        if (areShadowStacksEnabled)
-        {
-            // Undo the "pop", so that the ret could now succeed.
-            interruptedContext->SetSp(interruptedContext->GetSp() - 8);
-            interruptedContext->SetIp(origIp);
-        }
-
         ASSERT(!pThread->IsHijacked());
         return EXCEPTION_CONTINUE_EXECUTION;
     }

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

@@ -1006,3 +1006,25 @@ REDHAWK_PALEXPORT void PalFlushInstructionCache(_In_ void* pAddress, size_t size
     FlushInstructionCache(GetCurrentProcess(), pAddress, size);
 }
 
+#ifdef TARGET_AMD64
+uintptr_t GetSSP(CONTEXT *pContext)
+{
+    XSAVE_CET_U_FORMAT* pCET = (XSAVE_CET_U_FORMAT*)LocateXStateFeature(pContext, XSTATE_CET_U, NULL);
+    if ((pCET != NULL) && (pCET->Ia32CetUMsr != 0))
+    {
+        return pCET->Ia32Pl3SspMsr;
+    }
+
+    return 0;
+}
+
+void SetSSP(CONTEXT *pContext, uintptr_t ssp)
+{
+    XSAVE_CET_U_FORMAT* pCET = (XSAVE_CET_U_FORMAT*)LocateXStateFeature(pContext, XSTATE_CET_U, NULL);
+    if (pCET != NULL)
+    {
+        pCET->Ia32Pl3SspMsr = ssp;
+        pCET->Ia32CetUMsr = 1;
+    }
+}
+#endif // TARGET_AMD64

src/coreclr/vm/excep.cpp

@@ -6533,25 +6533,19 @@ VEH_ACTION WINAPI CLRVectoredExceptionHandler(PEXCEPTION_POINTERS pExceptionInfo
             // 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;
+            DWORD64 ssp = GetSSP(interruptedContext);
+            SetSSP(interruptedContext, ssp + 8);
         }
 
         // 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();
 
         FrameWithCookie<ResumableFrame> frame(pExceptionInfo->ContextRecord);
         frame.Push(pThread);
         CommonTripThread();
         frame.Pop(pThread);
 
-        if (areShadowStacksEnabled)
-        {
-            // Undo the "pop", so that the ret could now succeed.
-            interruptedContext->Rsp = interruptedContext->Rsp - 8;
-            interruptedContext->Rip = origIp;
-        }
-
         return VEH_CONTINUE_EXECUTION;
     }
 #endif