[NTOS:KE/x64] Implement KiUserCallbackExit

This is used in KiUserModeCallout instead of KiServiceExit2. The latter is broken, leaks non-volatile registers and will need to be modified to handle an exception frame, which we don't need/have here. It will also use sysret instead of iret and is generally simpler/faster.
Eventually it would be desirable to skip the entire trap frame setup and do everything in KiCallUserMode. This requires some cleanup and special handling for user APC delivery.

ntoskrnl/include/internal/amd64/ke.h

@@ -474,6 +474,12 @@ KiSetTrapContext(
     _In_ PCONTEXT Context,
     _In_ KPROCESSOR_MODE RequestorMode);
 
+// Exits to user mode, only restores the trap frame, zeroes the non-volatile registers
+DECLSPEC_NORETURN
+VOID
+KiUserCallbackExit(
+    _In_ PKTRAP_FRAME TrapFrame);
+
 #ifdef __cplusplus
 } // extern "C"
 #endif

ntoskrnl/ke/amd64/usercall.c

@@ -197,7 +197,7 @@ KiUserModeCallout(
     _enable();
 
     /* Exit to user-mode */
-    KiServiceExit(&CallbackTrapFrame, 0);
+    KiUserCallbackExit(&CallbackTrapFrame);
 }
 
 VOID

ntoskrnl/ke/amd64/usercall_asm.S

@@ -6,6 +6,10 @@
  */
 
 #include <ksamd64.inc>
+#include <trapamd64.inc>
+
+EXTERN KiInitiateUserApc:PROC
+EXTERN KeUserCallbackDispatcher:QWORD
 
 /*
  * NTSTATUS
@@ -45,6 +49,78 @@ PUBLIC KiCallUserMode
 
 .ENDP
 
+
+/*!
+ * \brief Exits to user mode, restores only rsp and rip, zeroes the rest.
+ *
+ * \todo Merge this with KiCallUserMode + KiUserModeCallout and skip the trap frame.
+ *
+ * DECLSPEC_NORETURN
+ * VOID
+ * KiUserCallbackExit(
+ *     _In_ PKTRAP_FRAME TrapFrame@<rcx>);
+ */
+ PUBLIC KiUserCallbackExit
+KiUserCallbackExit:
+
+    /* Point rsp to the trap frame */
+    mov rsp, rcx
+
+    /* Zero non-volatile registers for KiUserCallbackDispatcher.
+       This must be done before dispatching a pending user APC. */
+    xor rbx, rbx
+    xor rdi, rdi
+    xor rsi, rsi
+    xor rbp, rbp
+    xor r12, r12
+    xor r13, r13
+    xor r14, r14
+    xor r15, r15
+    pxor xmm6, xmm6
+    pxor xmm7, xmm7
+    pxor xmm8, xmm8
+    pxor xmm9, xmm9
+    pxor xmm10, xmm10
+    pxor xmm11, xmm11
+    pxor xmm12, xmm12
+    pxor xmm13, xmm13
+    pxor xmm14, xmm14
+    pxor xmm15, xmm15
+
+    /* Check for pending user APC */
+    mov rdx, gs:[PcCurrentThread]
+    HANDLE_USER_APCS rdx, rsp
+
+    /* Zero volatile registers */
+    xor rax, rax
+    xor rdx, rdx
+    xor r8, r8
+    xor r9, r9
+    xor r10, r10
+    pxor xmm0, xmm0
+    pxor xmm1, xmm1
+    pxor xmm2, xmm2
+    pxor xmm3, xmm3
+    pxor xmm4, xmm4
+    pxor xmm5, xmm5
+
+    /* Disable interrupts for return */
+    cli
+
+    /* Prepare user mode return address (rcx) and eflags (r11) for sysret */
+    mov rcx, qword ptr KeUserCallbackDispatcher[rip]
+    mov r11, EFLAGS_IF_MASK
+
+    /* Load user mode stack (It was copied to the trap frame) */
+    mov rsp, [rsp + KTRAP_FRAME_Rsp]
+
+    /* Swap gs back to user */
+    swapgs
+
+    /* return to user mode */
+    sysretq
+
+
 /*
  * DECLSPEC_NORETURN
  * VOID