Reconcile shadow stack with SP changes in RhpCallCatchFunclet

Author: VSadov

src/coreclr/nativeaot/Runtime/amd64/AsmOffsetsCpu.h

@@ -58,6 +58,7 @@ PLAT_ASM_OFFSET(0f0, PAL_LIMITED_CONTEXT, Xmm15)
 
 PLAT_ASM_SIZEOF(130, REGDISPLAY)
 PLAT_ASM_OFFSET(78, REGDISPLAY, SP)
+PLAT_ASM_OFFSET(80, REGDISPLAY, IP)
 
 PLAT_ASM_OFFSET(18, REGDISPLAY, pRbx)
 PLAT_ASM_OFFSET(20, REGDISPLAY, pRbp)

src/coreclr/nativeaot/Runtime/amd64/ExceptionHandling.asm

@@ -490,8 +490,9 @@ endif
         INLINE_THREAD_UNHIJACK rdx, rcx, r9                         ;; Thread in rdx, trashes rcx and r9
 
         mov     rcx, [rsp + rsp_offsetof_arguments + 18h]           ;; rcx <- current ExInfo *
+        mov     r10, [r8 + OFFSETOF__REGDISPLAY__IP]                ;; r10 <- original IP value
         mov     r8, [r8 + OFFSETOF__REGDISPLAY__SP]                 ;; r8 <- resume SP value
-        xor     r9d, r9d                                            ;; r9 <- 0
+        xor     r9, r9                                              ;; r9 <- 0
 
    @@:  mov     rcx, [rcx + OFFSETOF__ExInfo__m_pPrevExInfo]        ;; rcx <- next ExInfo
         cmp     rcx, r9
@@ -501,6 +502,20 @@ endif
 
    @@:  mov     [rdx + OFFSETOF__Thread__m_pExInfoStackHead], rcx   ;; store the new head on the Thread
 
+   ;; Sanity check: if we have shadow stack, it should agree with what we have in rsp
+   LOCAL_STACK_USE equ 118h
+   ifdef _DEBUG
+        rdsspq  r9
+        test    r9, r9
+        jz      @f
+        mov     r9, [r9]
+        cmp     [rsp + LOCAL_STACK_USE], r9
+        je      @f
+        int     3
+   @@:
+        xor     r9, r9                                              ;; r9 <- 0
+   endif
+
         test    [RhpTrapThreads], TrapThreadsFlags_AbortInProgress
         jz      @f
 
@@ -511,12 +526,28 @@ endif
         ;; It was the ThreadAbortException, so rethrow it
         mov     rcx, STATUS_REDHAWK_THREAD_ABORT
         mov     rdx, rax                                            ;; rdx <- continuation address as exception RIP
-        mov     rsp, r8                                             ;; reset the SP to resume SP value
-        jmp     RhpThrowHwEx ;; Throw the ThreadAbortException as a special kind of hardware exception
+        mov     rax, RhpThrowHwEx                                   ;; Throw the ThreadAbortException as a special kind of hardware exception
 
-        ;; reset RSP and jump to the continuation address
+        ;; reset RSP and jump to RAX
    @@:  mov     rsp, r8                                             ;; reset the SP to resume SP value
-        jmp     rax
+
+        ;; if have shadow stack, then we need to reconcile it with the rsp change we have just made
+        rdsspq  r9
+        test    r9, r9
+        jz      NoSSP
+
+        ;; Find the shadow stack pointer for the frame we are going to restore to.
+        ;; 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 
+        ;; the context and return one slot up from there.
+        ;; (Same logic as in GetSSPForFrameOnCurrentStack)
+        xor     r11, r11
+   @@:  inc     r11
+        cmp     [r9 + r11 * 8 - 8], r10
+        jne     @b
+
+        incsspq r11
+NoSSP:  jmp     rax
 
 
 NESTED_END RhpCallCatchFunclet, _TEXT