[NTOS] Fix unwinding through KiThreadStartup

ntoskrnl/ke/amd64/ctxswitch.S

@@ -10,6 +10,7 @@
 /* INCLUDES ******************************************************************/
 
 #include <ksamd64.inc>
+#include <trapamd64.inc>
 
 /*
  * BOOLEAN
@@ -70,10 +71,7 @@ PUBLIC KiThreadStartup
      * mov [rsp + SfP3Home], r8
      * mov [rsp + SfP4Home], r9
      */
-
-    /* Terminate the unwind chain, by setting rbp as frame pointer,
-       which contains 0 */
-    .setframe rbp, 0
+    .allocstack (5 * 8)
     .endprolog
 
     /* Clear all the non-volatile registers, so the thread won't be tempted to
@@ -101,21 +99,50 @@ PUBLIC KiThreadStartup
     mov r8, [rsp + SfP3Home]  /* ? */
     call qword ptr [rsp + SfP4Home]     /* SystemRoutine */
 
-    /* The thread returned. If it was a user-thread, we have a return address
-       and all is well, otherwise this is very bad. */
-    mov rcx, [rsp + SfReturn]
-    or rcx, rcx
-    jnz .leave
+    /* Return to the exit code */
+    add rsp, 5 * 8
+    ret
+.ENDP
 
-    /* A system thread returned...this is very bad! */
-    int 3
+PUBLIC KiInvalidSystemThreadStartupExit
+.PROC KiInvalidSystemThreadStartupExit
+    .endprolog
 
-.leave:
-    /* It was a user thread, set our trapframe for the System Call Exit Dispatcher */
-    lea rcx, [rsp + 6 * 8 + KEXCEPTION_FRAME_LENGTH]
+    /* This is invalid! */
+    int HEX(2C)
+    nop
+.ENDP
+
+PUBLIC KiUserThreadStartupExit
+.PROC KiUserThreadStartupExit
+    .allocstack (KEXCEPTION_FRAME_LENGTH - 8)
+    .savereg rbp, ExRbp
+    .savereg rbx, ExRbx
+    .savereg rdi, ExRdi
+    .savereg rsi, ExRsi
+    .savereg r12, ExR12
+    .savereg r13, ExR13
+    .savereg r14, ExR14
+    .savereg r15, ExR15
+    .savexmm128 xmm6, ExXmm6
+    .savexmm128 xmm7, ExXmm7
+    .savexmm128 xmm8, ExXmm8
+    .savexmm128 xmm9, ExXmm9
+    .savexmm128 xmm10, ExXmm10
+    .savexmm128 xmm11, ExXmm11
+    .savexmm128 xmm12, ExXmm12
+    .savexmm128 xmm13, ExXmm13
+    .savexmm128 xmm14, ExXmm14
+    .savexmm128 xmm15, ExXmm15
+    .endprolog
+
+    /* Restore the exception frame */
+    RESTORE_EXCEPTION_STATE
+
+    /* Point rcx to the trap frame */
+    lea rcx, [rsp + 8]
 
     /* Return to the trap exit code */
-    add rsp, 5 * 8
     ret
 .ENDP
 

ntoskrnl/ke/amd64/thrdini.c

@@ -13,6 +13,10 @@
 #define NDEBUG
 #include <debug.h>
 
+extern void KiInvalidSystemThreadStartupExit(void);
+extern void KiUserThreadStartupExit(void);
+extern void KiServiceExit3(void);
+
 typedef struct _KUINIT_FRAME
 {
     KSWITCH_FRAME CtxSwitchFrame;
@@ -98,8 +102,11 @@ KiInitializeContextThread(IN PKTHREAD Thread,
         /* Terminate the Exception Handler List */
         TrapFrame->ExceptionFrame = 0;
 
-        /* We return to ... */
-        StartFrame->Return = (ULONG64)KiServiceExit2;
+        /* KiThreadStartup returns to KiUserThreadStartupExit */
+        StartFrame->Return = (ULONG64)KiUserThreadStartupExit;
+
+        /* KiUserThreadStartupExit returns to KiServiceExit3 */
+        InitFrame->ExceptionFrame.Return = (ULONG64)KiServiceExit3;
     }
     else
     {
@@ -121,8 +128,8 @@ KiInitializeContextThread(IN PKTHREAD Thread,
         /* No NPX State */
         Thread->NpxState = 0xA;
 
-        /* We have no return address! */
-        StartFrame->Return = 0;
+        /* This must never return! */
+        StartFrame->Return = (ULONG64)KiInvalidSystemThreadStartupExit;
     }
 
     /* Set up the Context Switch Frame */

ntoskrnl/ke/amd64/trap.S

@@ -930,16 +930,27 @@ PUBLIC KiServiceExit2
 .PROC KiServiceExit2
     .ENDPROLOG
 
+    // FIXME: this should probably also restore an exception frame
+
+    mov rsp, rcx
+.ENDP
+
+PUBLIC KiServiceExit3
+.PROC KiServiceExit3
+    .PUSHFRAME
+    .ALLOCSTACK (KTRAP_FRAME_LENGTH - MachineFrameLength)
+    .ENDPROLOG
+
 #if DBG
     /* Get the current IRQL and compare it to the trap frame */
     mov rax, cr8
-    cmp byte ptr [rcx + KTRAP_FRAME_PreviousIrql], al
+    cmp byte ptr [rsp + KTRAP_FRAME_PreviousIrql], al
     je KiServiceExit2_ok1
     int HEX(2C)
 
 KiServiceExit2_ok1:
     /* Check if this is a user mode exit */
-    mov ah, byte ptr [rcx + KTRAP_FRAME_SegCs]
+    mov ah, byte ptr [rsp + KTRAP_FRAME_SegCs]
     test ah, 1
     jz KiServiceExit2_kernel
 
@@ -951,10 +962,8 @@ KiServiceExit2_ok1:
 KiServiceExit2_kernel:
 #endif
 
-    mov rbp, rcx
-    mov rsp, rcx
-
     /* Return */
+    mov rbp, rsp
     ExitTrap TF_SAVE_ALL
 .ENDP