Fix SSP restoring in edge cases

src/coreclr/System.Private.CoreLib/src/System/Runtime/ExceptionServices/AsmOffsets.cs

@@ -72,7 +72,7 @@ class AsmOffsets
     public const int OFFSETOF__REGDISPLAY__SP = 0x1a70;
     public const int OFFSETOF__REGDISPLAY__ControlPC = 0x1a78;
 #else // TARGET_UNIX
-    public const int SIZEOF__REGDISPLAY = 0xbe0;
+    public const int SIZEOF__REGDISPLAY = 0xbf0;
     public const int OFFSETOF__REGDISPLAY__SP = 0xbd0;
     public const int OFFSETOF__REGDISPLAY__ControlPC = 0xbd8;
 #endif // TARGET_UNIX

src/coreclr/inc/regdisp.h

@@ -45,6 +45,10 @@ struct REGDISPLAY_BASE {
 
     TADDR SP;
     TADDR ControlPC; // LOONGARCH: use RA for PC
+
+#if defined(TARGET_AMD64) && defined(TARGET_WINDOWS)
+    TADDR SSP;
+#endif
 };
 
 inline PCODE GetControlPC(const REGDISPLAY_BASE *pRD) {
@@ -510,6 +514,12 @@ inline void FillRegDisplay(const PREGDISPLAY pRD, PT_CONTEXT pctx, PT_CONTEXT pC
     // This will setup the PC and SP
     SyncRegDisplayToCurrentContext(pRD);
 
+#if !defined(DACCESS_COMPILE)
+#if defined(TARGET_AMD64) && defined(TARGET_WINDOWS)
+    pRD->SSP = GetSSP(pctx);
+#endif
+#endif // !DACCESS_COMPILE
+
     if (fLightUnwind)
         return;
 

src/coreclr/vm/amd64/Context.S

@@ -44,7 +44,14 @@ NESTED_ENTRY ClrRestoreNonvolatileContextWorker, _TEXT, NoHandler
         rdsspq  rax
         sub     r11, rax
         shr     r11, 3
-        incsspq r11
+        // the incsspq instruction uses only the lowest 8 bits of the argument, so we need to loop in case the increment is larger than 255
+        mov     rax, 255
+    Update_Loop:
+        cmp     r11, rax
+        cmovb   rax, r11
+        incsspq rax
+        sub     r11, rax
+        ja      Update_Loop
     No_Ssp_Update:
 
         // When user-mode shadow stacks are enabled, and for example the intent is to continue execution in managed code after

src/coreclr/vm/amd64/Context.asm

@@ -53,7 +53,14 @@ NESTED_ENTRY ClrRestoreNonvolatileContextWorker, _TEXT
         rdsspq  rax
         sub     r11, rax
         shr     r11, 3
-        incsspq r11
+        ; the incsspq instruction uses only the lowest 8 bits of the argument, so we need to loop in case the increment is larger than 255
+        mov     rax, 255
+    Update_Loop:
+        cmp     r11, rax
+        cmovb   rax, r11
+        incsspq rax
+        sub     r11, rax
+        ja      Update_Loop
     No_Ssp_Update:
 
         ; When user-mode shadow stacks are enabled, and for example the intent is to continue execution in managed code after

src/coreclr/vm/amd64/cgenamd64.cpp

@@ -108,13 +108,20 @@ void InlinedCallFrame::UpdateRegDisplay(const PREGDISPLAY pRD, bool updateFloats
     if (updateFloats)
     {
         UpdateFloatingPointRegisters(pRD);
+        // The float updating unwinds the stack so the pRD->pCurrentContext->Rip contains correct unwound Rip
+        // This is used for exception handling and the Rip extracted from m_pCallerReturnAddress is slightly
+        // off, which causes problem with searching for the return address on shadow stack on x64, so
+        // we keep the value from the unwind.
     }
+    else
 #endif // DACCESS_COMPILE
+    {
+        pRD->pCurrentContext->Rip = *(DWORD64 *)&m_pCallerReturnAddress;
+    }
 
     pRD->IsCallerContextValid = FALSE;
     pRD->IsCallerSPValid      = FALSE;        // Don't add usage of this field.  This is only temporary.
 
-    pRD->pCurrentContext->Rip = *(DWORD64 *)&m_pCallerReturnAddress;
     pRD->pCurrentContext->Rsp = *(DWORD64 *)&m_pCallSiteSP;
     pRD->pCurrentContext->Rbp = *(DWORD64 *)&m_pCalleeSavedFP;
 

src/coreclr/vm/exceptionhandling.cpp

@@ -7673,27 +7673,23 @@ UINT_PTR GetEstablisherFrame(REGDISPLAY* pvRegDisplay, ExInfo* exInfo)
 #endif
 }
 
-size_t GetSSPForFrameOnCurrentStack(CONTEXT *pContext)
-{
-    size_t *targetSSP = 0;
-
 #if defined(HOST_AMD64) && defined(HOST_WINDOWS)
-    targetSSP = (size_t *)_rdsspq();
-    // Find the shadow stack pointer for the frame we are going to restore to.
+size_t GetSSPForFrameOnCurrentStack(TADDR ip)
+{
+    size_t *targetSSP = (size_t *)_rdsspq();
     // The SSP we search is pointing to the return address of the frame represented
-    // by the passed in context. So we search for the instruction pointer from 
+    // by the passed in IP. So we search for the instruction pointer from 
     // the context and return one slot up from there.
     if (targetSSP != NULL)
     {
-        size_t ip = GetIP(pContext);
         while (*targetSSP++ != ip)
         {
         }
     }
-#endif // HOST_AMD64 && HOST_WINDOWS
 
     return (size_t)targetSSP;
 }
+#endif // HOST_AMD64 && HOST_WINDOWS
 
 extern "C" void * QCALLTYPE CallCatchFunclet(QCall::ObjectHandleOnStack exceptionObj, BYTE* pHandlerIP, REGDISPLAY* pvRegDisplay, ExInfo* exInfo)
 {
@@ -7711,7 +7707,12 @@ extern "C" void * QCALLTYPE CallCatchFunclet(QCall::ObjectHandleOnStack exceptio
     exInfo->m_ScannedStackRange.ExtendUpperBound(exInfo->m_frameIter.m_crawl.GetRegisterSet()->SP);
     DWORD_PTR dwResumePC = 0;
     UINT_PTR callerTargetSp = 0;
-    size_t targetSSP = GetSSPForFrameOnCurrentStack(pvRegDisplay->pCurrentContext);
+#if defined(HOST_AMD64) && defined(HOST_WINDOWS)
+    size_t targetSSP = exInfo->m_frameIter.m_crawl.GetRegisterSet()->SSP;
+    _ASSERTE(targetSSP == 0 || (*(size_t*)(targetSSP-8) == exInfo->m_frameIter.m_crawl.GetRegisterSet()->ControlPC));
+#else
+    size_t targetSSP = 0;
+#endif
 
     if (pHandlerIP != NULL)
     {
@@ -7929,6 +7930,12 @@ extern "C" void QCALLTYPE ResumeAtInterceptionLocation(REGDISPLAY* pvRegDisplay)
 
     pThread->GetExceptionState()->GetDebuggerState()->GetDebuggerInterceptInfo(&pInterceptMD, NULL, (PBYTE*)&(sfInterceptStackFrame.SP), &ulRelOffset, NULL);
 
+#if defined(HOST_AMD64) && defined(HOST_WINDOWS)
+    TADDR targetSSP = pExInfo->m_frameIter.m_crawl.GetRegisterSet()->SSP;
+#else
+    TADDR targetSSP = 0;
+#endif
+
     ExInfo::PopExInfos(pThread, (void*)targetSp);
 
     PCODE pStartAddress = pInterceptMD->GetNativeCode();
@@ -7941,8 +7948,6 @@ extern "C" void QCALLTYPE ResumeAtInterceptionLocation(REGDISPLAY* pvRegDisplay)
     _ASSERTE(FitsIn<DWORD>(ulRelOffset));
     uResumePC = codeInfo.GetJitManager()->GetCodeAddressForRelOffset(codeInfo.GetMethodToken(), static_cast<DWORD>(ulRelOffset));
 
-    size_t targetSSP = GetSSPForFrameOnCurrentStack(pvRegDisplay->pCurrentContext);
-
     SetIP(pvRegDisplay->pCurrentContext, uResumePC);
     ClrRestoreNonvolatileContext(pvRegDisplay->pCurrentContext, targetSSP);
 }
@@ -8416,6 +8421,14 @@ extern "C" bool QCALLTYPE SfiInit(StackFrameIterator* pThis, CONTEXT* pStackwalk
     if (result)
     {
         TADDR controlPC = pThis->m_crawl.GetRegisterSet()->ControlPC;
+
+#if defined(HOST_AMD64) && defined(HOST_WINDOWS)
+        // Get the SSP for the first managed frame. It is incremented during the stack walk so that
+        // when we reach the handling frame, it contains correct SSP to set when resuming after
+        // the catch handler.
+        pThis->m_crawl.GetRegisterSet()->SSP = GetSSPForFrameOnCurrentStack(controlPC);
+#endif
+
         if (!pThis->m_crawl.HasFaulted() && !pThis->m_crawl.IsIPadjusted())
         {
             controlPC -= STACKWALK_CONTROLPC_ADJUST_OFFSET;

src/coreclr/vm/stackwalk.cpp

@@ -524,6 +524,12 @@ UINT_PTR Thread::VirtualUnwindCallFrame(PREGDISPLAY pRD, EECodeInfo* pCodeInfo /
         VirtualUnwindCallFrame(pRD->pCurrentContext, pRD->pCurrentContextPointers, pCodeInfo);
     }
 
+#if defined(TARGET_AMD64) && defined(TARGET_WINDOWS)
+    if (pRD->SSP != 0)
+    {
+        pRD->SSP += 8;
+    }
+#endif // TARGET_AMD64 && TARGET_WINDOWS
     SyncRegDisplayToCurrentContext(pRD);
     pRD->IsCallerContextValid = FALSE;
     pRD->IsCallerSPValid      = FALSE;        // Don't add usage of this field.  This is only temporary.
@@ -1555,6 +1561,9 @@ void StackFrameIterator::SkipTo(StackFrameIterator *pOtherStackFrameIterator)
     *pRD->pCurrentContextPointers = *pOtherRD->pCurrentContextPointers;
     SetIP(pRD->pCurrentContext, GetIP(pOtherRD->pCurrentContext));
     SetSP(pRD->pCurrentContext, GetSP(pOtherRD->pCurrentContext));
+#if defined(TARGET_AMD64) && defined(TARGET_WINDOWS)
+    pRD->SSP = pOtherRD->SSP;
+#endif
 
 #define CALLEE_SAVED_REGISTER(regname) pRD->pCurrentContext->regname = (pRD->pCurrentContextPointers->regname == NULL) ? pOtherRD->pCurrentContext->regname : *pRD->pCurrentContextPointers->regname;
     ENUM_CALLEE_SAVED_REGISTERS();

src/tests/Regressions/coreclr/GitHub_104820/test104820.cs

@@ -0,0 +1,45 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+using System;
+using System.Runtime.CompilerServices;
+using Xunit;
+
+public class Test104820
+{
+    // Test that SSP is updated properly when resuming after catch when the shadow stack
+    // contains the instruction pointer of the resume frame multiple times and the SSP needs
+    // to be restored to the location of one that's not the first one found on the shadow
+    // stack.
+    // This test fails with stack overflow of the shadow stack if the SSP is updated incorrectly.
+    static void ShadowStackPointerUpdateTest(int depth)
+    {
+        if (depth == 0)
+        {
+            throw new Exception();
+        }
+
+        for (int i = 0; i < 1000; i++)
+        {
+            try
+            {
+                try
+                {
+                    ShadowStackPointerUpdateTest(depth - 1);
+                }
+                catch (Exception)
+                {
+                    throw new Exception();
+                }
+            }
+            catch (Exception) when (depth == 100)
+            {
+            }
+        }
+    }
+
+    [Fact]
+    public static void TestEntryPoint()
+    {
+        ShadowStackPointerUpdateTest(100);
+    }
+}

src/tests/Regressions/coreclr/GitHub_104820/test104820.csproj

@@ -0,0 +1,11 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <CLRTestPriority>1</CLRTestPriority>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="test104820.cs" />
+  </ItemGroup>
+  <ItemGroup>
+    <ProjectReference Include="$(TestSourceDir)Common/CoreCLRTestLibrary/CoreCLRTestLibrary.csproj" />
+  </ItemGroup>
+</Project>