From 24b4026ce8aecb9d409b0ad7a378ad7a034bbc66 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Wed, 27 Mar 2024 23:17:29 +0200
Subject: [PATCH] [NTOS:KE/x64][SDK] Add HANDLE_USER_APCS asm macro

This also fixes delivering APCs from the system call handler, which previously would have clobbered rax. Also don't use the thread's TrapFrame member, which is not always set, when returning.
---
 ntoskrnl/ke/amd64/trap.S      | 26 +++++++++++++++++++-------
 sdk/include/asm/trapamd64.inc | 20 ++++++++++++++------
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/ntoskrnl/ke/amd64/trap.S b/ntoskrnl/ke/amd64/trap.S
index 88679949b17..79b58559129 100644
--- a/ntoskrnl/ke/amd64/trap.S
+++ b/ntoskrnl/ke/amd64/trap.S
@@ -850,12 +850,9 @@ GLOBAL_LABEL KiSystemServiceExit
 
     ASSERT_TRAP_FRAME_INTS_ENABLED rsp + MAX_SYSCALL_PARAM_SIZE
 
-    /* Check for pending user APC */
-    mov rcx, gs:qword ptr [PcCurrentThread]
-    cmp byte ptr [rcx + ThApcState + AsUserApcPending], 0
-    jz no_user_apc_pending
-    call KiInitiateUserApc
-no_user_apc_pending:
+    /* Check for pending user APCs */
+    mov rcx, gs:[PcCurrentThread]
+    HANDLE_USER_APCS rcx, rsp + MAX_SYSCALL_PARAM_SIZE
 
     /* Disable interrupts for return */
     cli
@@ -1137,12 +1134,23 @@ PUBLIC KiSetTrapContext
  */
 EXTERN KiDeliverApc:PROC
 
+/*
+ * VOID
+ * KiInitiateUserApc(
+ *     _In_ PKTRAP_FRAME TrapFrame@<rcx>);
+ *
+ * This function is called to deliver user mode APCs.
+ * It clobbers all non-volatile registers, except rax.
+ */
 PUBLIC KiInitiateUserApc
 .PROC KiInitiateUserApc
 
     /* Generate a KEXCEPTION_FRAME on the stack */
     GENERATE_EXCEPTION_FRAME
 
+    /* Save rax to not clobber the return for the system call handler */
+    mov [rsp + ExP1Home], rax
+
     /* Raise IRQL to APC_LEVEL */
     mov rax, APC_LEVEL
     mov cr8, rax
@@ -1150,6 +1158,9 @@ PUBLIC KiInitiateUserApc
     /* Get the current thread */
     mov rbp, gs:[PcCurrentThread]
 
+    /* Save the trap frame in rsi */
+    mov rsi, rcx
+
 deliver_apcs:
 
     /* Enable interrupts */
@@ -1158,7 +1169,7 @@ deliver_apcs:
     /* Call the C function */
     mov ecx, 1
     mov rdx, rsp
-    mov r8, [rbp + ThTrapFrame]
+    mov r8, rsi
     call KiDeliverApc
 
     /* Disable interrupts again */
@@ -1173,6 +1184,7 @@ deliver_apcs:
     mov cr8, rax
 
     /* Restore the registers from the KEXCEPTION_FRAME */
+    mov rax, [rsp + ExP1Home]
     RESTORE_EXCEPTION_STATE
 
     /* Return */
diff --git a/sdk/include/asm/trapamd64.inc b/sdk/include/asm/trapamd64.inc
index 5303adf2dbb..0c09063f058 100644
--- a/sdk/include/asm/trapamd64.inc
+++ b/sdk/include/asm/trapamd64.inc
@@ -37,6 +37,19 @@ IrqlIsPassive:
 #endif
 ENDM
 
+// Checks for user APCs and delivers them if necessary.
+// Clobbers all volatile registers except rax.
+MACRO(HANDLE_USER_APCS, ThreadReg, TrapFrame)
+    LOCAL NoUserApcPending
+
+    /* Check for pending user APC */
+    cmp byte ptr [ThreadReg + ThApcState + AsUserApcPending], 0
+    jz NoUserApcPending
+    lea rcx, [TrapFrame]
+    call KiInitiateUserApc
+NoUserApcPending:
+ENDM
+
 APIC_EOI = HEX(0FFFFFFFFFFFE00B0)
 
 TF_VOLATILES            = HEX(01)
@@ -195,7 +208,6 @@ ENDM
  */
 MACRO(ExitTrap, Flags)
     LOCAL kernel_mode_return
-    LOCAL NoUserApc
 
     ASSERT_TRAP_FRAME_IRQL_VALID rbp
 
@@ -217,12 +229,8 @@ MACRO(ExitTrap, Flags)
     jz kernel_mode_return
 
     if (Flags AND TF_CHECKUSERAPC)
-        /* Load current thread into r10 */
         mov r10, gs:[PcCurrentThread]
-        cmp byte ptr [r10 + KTHREAD_UserApcPending], 0
-        je NoUserApc
-        call KiInitiateUserApc
-    NoUserApc:
+        HANDLE_USER_APCS r10, rbp
     endif
 
     ASSERT_TRAP_FRAME_INTS_ENABLED rbp