From 34576c7015d74929070f0e765b6ce8c82108c58b Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Mon, 25 Mar 2024 20:57:41 +0200
Subject: [PATCH] [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 |  6 +++
 ntoskrnl/ke/amd64/usercall.c         |  2 +-
 ntoskrnl/ke/amd64/usercall_asm.S     | 76 ++++++++++++++++++++++++++++
 3 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/ntoskrnl/include/internal/amd64/ke.h b/ntoskrnl/include/internal/amd64/ke.h
index d55bd302510..7caf2b10d9a 100644
--- a/ntoskrnl/include/internal/amd64/ke.h
+++ b/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
diff --git a/ntoskrnl/ke/amd64/usercall.c b/ntoskrnl/ke/amd64/usercall.c
index 8eb059ccd54..9bcc81384b4 100644
--- a/ntoskrnl/ke/amd64/usercall.c
+++ b/ntoskrnl/ke/amd64/usercall.c
@@ -197,7 +197,7 @@ KiUserModeCallout(
     _enable();
 
     /* Exit to user-mode */
-    KiServiceExit(&CallbackTrapFrame, 0);
+    KiUserCallbackExit(&CallbackTrapFrame);
 }
 
 VOID
diff --git a/ntoskrnl/ke/amd64/usercall_asm.S b/ntoskrnl/ke/amd64/usercall_asm.S
index 8962f2470e0..7dc5bf1f6b8 100644
--- a/ntoskrnl/ke/amd64/usercall_asm.S
+++ b/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