[NTOS:KE] Fix NtRaiseException and NtContinue on x64

- Implement KiExceptionExit, which works like KiServiceExit2, but takes an exception frame as 2nd parameter
- Add a local exception frame to NtRaiseException and NtContinue, which is needed to exit with a full context on x64 and arm
- Use KeContextToTrapFrame again instead of KiSetTrapContext, since we have a proper exception frame now.

ntoskrnl/include/internal/amd64/ke.h

@@ -480,6 +480,12 @@ VOID
 KiUserCallbackExit(
     _In_ PKTRAP_FRAME TrapFrame);
 
+DECLSPEC_NORETURN
+VOID
+KiExceptionExit(
+    _In_ PKTRAP_FRAME TrapFrame,
+    _In_ PKEXCEPTION_FRAME ExceptionFrame);
+
 BOOLEAN
 KiProcessorFreezeHandler(
     _In_ PKTRAP_FRAME TrapFrame,

ntoskrnl/include/internal/i386/ke.h

@@ -668,6 +668,31 @@ KiConvertToGuiThread(
     VOID
 );
 
+DECLSPEC_NORETURN
+VOID
+FASTCALL
+KiServiceExit(
+    IN PKTRAP_FRAME TrapFrame,
+    IN NTSTATUS Status
+);
+
+DECLSPEC_NORETURN
+VOID
+FASTCALL
+KiServiceExit2(
+    IN PKTRAP_FRAME TrapFrame
+);
+
+FORCEINLINE
+DECLSPEC_NORETURN
+VOID
+KiExceptionExit(
+    _In_ PKTRAP_FRAME TrapFrame,
+    _In_ PKEXCEPTION_FRAME ExceptionFrame)
+{
+    KiServiceExit2(TrapFrame);
+}
+
 //
 // Global x86 only Kernel data
 //

ntoskrnl/include/internal/ke.h

@@ -883,21 +883,6 @@ KiContinue(
     IN PKTRAP_FRAME TrapFrame
 );
 
-DECLSPEC_NORETURN
-VOID
-FASTCALL
-KiServiceExit(
-    IN PKTRAP_FRAME TrapFrame,
-    IN NTSTATUS Status
-);
-
-DECLSPEC_NORETURN
-VOID
-FASTCALL
-KiServiceExit2(
-    IN PKTRAP_FRAME TrapFrame
-);
-
 #ifndef _M_AMD64
 VOID
 FASTCALL

ntoskrnl/ke/amd64/trap.S

@@ -944,62 +944,50 @@ NoUserApcPending:
 /*!
  * VOID
  * DECLSPEC_NORETURN
- * KiServiceExit(IN PKTRAP_FRAME TrapFrame, IN NTSTATUS Status));
+ * KiExceptionExit(
+ *     _In_ PKTRAP_FRAME TrapFrame@<rcx>,
+ *     _In_ PKEXCEPTION_FRAME ExceptionFrame@<rdx>);
  */
-PUBLIC KiServiceExit
-.PROC KiServiceExit
-    .endprolog
+PUBLIC KiExceptionExit
+KiExceptionExit:
 
-    lea rsp, [rcx - MAX_SYSCALL_PARAM_SIZE]
-    jmp KiSystemServiceExit
-
-.ENDP
-
-
-/*!
- * VOID
- * DECLSPEC_NORETURN
- * KiServiceExit2(IN PKTRAP_FRAME TrapFrame);
- */
-PUBLIC KiServiceExit2
-.PROC KiServiceExit2
-    .ENDPROLOG
-
-    // FIXME: this should probably also restore an exception frame
+    /* Restore registers from exception frame */
+    movaps xmm6, [rdx + ExXmm6]
+    movaps xmm7, [rdx + ExXmm7]
+    movaps xmm8, [rdx + ExXmm8]
+    movaps xmm9, [rdx + ExXmm9]
+    movaps xmm10, [rdx + ExXmm10]
+    movaps xmm11, [rdx + ExXmm11]
+    movaps xmm12, [rdx + ExXmm12]
+    movaps xmm13, [rdx + ExXmm13]
+    movaps xmm14, [rdx + ExXmm14]
+    movaps xmm15, [rdx + ExXmm15]
+    mov rbx, [rdx + ExRbx]
+    mov rdi, [rdx + ExRdi]
+    mov rsi, [rdx + ExRsi]
+    mov r12, [rdx + ExR12]
+    mov r13, [rdx + ExR13]
+    mov r14, [rdx + ExR14]
+    mov r15, [rdx + ExR15]
 
+    /* Point rsp at the trap frame */
     mov rsp, rcx
-.ENDP
+    /* Fall through */
 
+/*
+ * Internal function. Exits to user-mode with rsp pointing to the trap frame.
+ * All non-volatile register context must be set up already.
+ * Used by KiInitializeContextThread to set up the init path for a new thread.
+ */
 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 [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 [rsp + KTRAP_FRAME_SegCs]
-    test ah, 1
-    jz KiServiceExit2_kernel
-
-    /* Validate that we are at PASSIVE_LEVEL */
-    test al, al
-    jz KiServiceExit2_kernel
-    int HEX(2C)
-
-KiServiceExit2_kernel:
-#endif
-
     /* Return */
     mov rbp, rsp
-    ExitTrap TF_SAVE_ALL
+    ExitTrap (TF_SEGMENTS or TF_CHECKUSERAPC)
 .ENDP
 
 

ntoskrnl/ke/except.c

@@ -28,16 +28,12 @@ KiContinuePreviousModeUser(IN PCONTEXT Context,
     RtlCopyMemory(&LocalContext, Context, sizeof(CONTEXT));
     Context = &LocalContext;
 
-#ifdef _M_AMD64
-    KiSetTrapContext(TrapFrame, &LocalContext, UserMode);
-#else
     /* Convert the context into Exception/Trap Frames */
     KeContextToTrapFrame(&LocalContext,
                          ExceptionFrame,
                          TrapFrame,
                          LocalContext.ContextFlags,
                          UserMode);
-#endif
 }
 
 NTSTATUS
@@ -66,16 +62,12 @@ KiContinue(IN PCONTEXT Context,
         }
         else
         {
-#ifdef _M_AMD64
-            KiSetTrapContext(TrapFrame, Context, KernelMode);
-#else
             /* Convert the context into Exception/Trap Frames */
             KeContextToTrapFrame(Context,
                                  ExceptionFrame,
                                  TrapFrame,
                                  Context->ContextFlags,
                                  KernelMode);
-#endif
         }
     }
     _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
@@ -183,8 +175,14 @@ NtRaiseException(IN PEXCEPTION_RECORD ExceptionRecord,
     NTSTATUS Status;
     PKTHREAD Thread;
     PKTRAP_FRAME TrapFrame;
+#ifdef _M_IX86
+    PKEXCEPTION_FRAME ExceptionFrame = NULL;
+#else
+    KEXCEPTION_FRAME LocalExceptionFrame;
+    PKEXCEPTION_FRAME ExceptionFrame = &LocalExceptionFrame;
+#endif
 
-    /* Get trap frame and link previous one*/
+    /* Get trap frame and link previous one */
     Thread = KeGetCurrentThread();
     TrapFrame = Thread->TrapFrame;
     Thread->TrapFrame = KiGetLinkedTrapFrame(TrapFrame);
@@ -197,21 +195,23 @@ NtRaiseException(IN PEXCEPTION_RECORD ExceptionRecord,
     /* Raise the exception */
     Status = KiRaiseException(ExceptionRecord,
                               Context,
-                              NULL,
+                              ExceptionFrame,
                               TrapFrame,
                               FirstChance);
     if (NT_SUCCESS(Status))
     {
         /* It was handled, so exit restoring all state */
-        KiServiceExit2(TrapFrame);
+        KiExceptionExit(TrapFrame, ExceptionFrame);
     }
+#ifdef _M_IX86
     else
     {
         /* Exit with error */
         KiServiceExit(TrapFrame, Status);
     }
+#endif
 
-    /* We don't actually make it here */
+    /* Return to the caller */
     return Status;
 }
 
@@ -223,6 +223,12 @@ NtContinue(IN PCONTEXT Context,
     PKTHREAD Thread;
     NTSTATUS Status;
     PKTRAP_FRAME TrapFrame;
+#ifdef _M_IX86
+    PKEXCEPTION_FRAME ExceptionFrame = NULL;
+#else
+    KEXCEPTION_FRAME LocalExceptionFrame;
+    PKEXCEPTION_FRAME ExceptionFrame = &LocalExceptionFrame;
+#endif
 
     /* Get trap frame and link previous one*/
     Thread = KeGetCurrentThread();
@@ -230,22 +236,24 @@ NtContinue(IN PCONTEXT Context,
     Thread->TrapFrame = KiGetLinkedTrapFrame(TrapFrame);
 
     /* Continue from this point on */
-    Status = KiContinue(Context, NULL, TrapFrame);
+    Status = KiContinue(Context, ExceptionFrame, TrapFrame);
     if (NT_SUCCESS(Status))
     {
         /* Check if alert was requested */
         if (TestAlert) KeTestAlertThread(Thread->PreviousMode);
 
         /* Exit to new trap frame */
-        KiServiceExit2(TrapFrame);
+        KiExceptionExit(TrapFrame, ExceptionFrame);
     }
+#ifdef _M_IX86
     else
     {
         /* Exit with an error */
         KiServiceExit(TrapFrame, Status);
     }
+#endif
 
-    /* We don't actually make it here */
+    /* Return to the caller */
     return Status;
 }