Fix SSP restoring in edge cases (#104820)

* Fix SSP restoring in edge cases

There are edge cases when the SSP restoring for continuation after a
catch handler completes doesn't work correctly. The problem is caused by
the fact that we scan for the Rip of the frame handling the exception on
the shadow stack to find where to restore it, and in those edge cases,
the same address can be there multiple times and the first occurence is
not the right one. For example, when an exception is thrown from a catch
handler, it escapes the handler and the handler for the escaped
exception is in the same method as the one that invoked the handler.

This change fixes it by finding the SSP of the first managed frame where
we search for the handler and then updating the SSP with every unwind.
The SSP is stored in the REGDISPLAY. So when we reach the
CallCatchFunclet, the REGDISPLAY contains the SSP to restore.

There was also one more issue with restoring the SSP. I turned out that
the incsspq instruction uses only the lowest 8 bits of the argument to
increment the SSP, so the ClrRestoreNonVolatileContextWorker needs to
have a loop that repeats that instruction in case we need to move it by
more than 255 slots.

* Fix linux x64 build

* Fix release build REGDISPLAY size constant

* Add regression test

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

---------

Co-authored-by: Jan Vorlicek <jan.vorlicek@volny,cz>
Co-authored-by: Jan Kotas <jkotas@microsoft.com>

Author: 3 people

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>