From 1538712c0b61aadb9018fade9b6146a184afc73c Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Mon, 25 Mar 2024 15:57:34 +0200
Subject: [PATCH] [NTOS:KE/x64] Move setting the thread's trap frame to
 KiSystemCallEntry64

This is needed, because KiSystemCallHandler can be called multiple times for the same syscall entry, which would mess up the linkage. This replaces a previous hack and makes things cleaner.
---
 ntoskrnl/ke/amd64/trap.S        | 15 +++++++++------
 ntoskrnl/ke/amd64/traphandler.c |  4 ----
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/ntoskrnl/ke/amd64/trap.S b/ntoskrnl/ke/amd64/trap.S
index 043911cd102..388afe9c2e2 100644
--- a/ntoskrnl/ke/amd64/trap.S
+++ b/ntoskrnl/ke/amd64/trap.S
@@ -817,6 +817,15 @@ PUBLIC KiSystemCallEntry64
     stmxcsr [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_MxCsr]
     ldmxcsr gs:[PcMxCsr]
 
+    /* Get the current thread and the trap frame */
+    mov rax, gs:[PcCurrentThread]
+    mov rcx, [rax + ThTrapFrame]
+
+    /* Save the old trap frame */
+    lea rdx, [rsp + MAX_SYSCALL_PARAM_SIZE]
+    mov [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_TrapFrame], rcx
+    mov [rax + ThTrapFrame], rdx
+
 #if DBG
     /* Check IRQL */
     mov rax, cr8
@@ -1077,12 +1086,6 @@ AlreadyLargeStack:
     /* Disable interrupts for return */
     cli
 
-    // FIXME: should just do the trap frame switch in KiSystemCallHandler64
-    /* Restore old trap frame */
-    mov rcx, gs:[PcCurrentThread]
-    mov rdx, [rsp + 48 + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_TrapFrame]
-    mov [rcx + KTHREAD_TrapFrame], rdx
-
     // Restore register parameters
     mov rcx, [rsp + 48 + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_Rip]
     mov rdx, [rsp + 48 + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_Rdx]
diff --git a/ntoskrnl/ke/amd64/traphandler.c b/ntoskrnl/ke/amd64/traphandler.c
index c8e67ad583e..c4e973f4927 100644
--- a/ntoskrnl/ke/amd64/traphandler.c
+++ b/ntoskrnl/ke/amd64/traphandler.c
@@ -116,10 +116,6 @@ KiSystemCallHandler(
     /* Set previous mode */
     Thread->PreviousMode = TrapFrame->PreviousMode = UserMode;
 
-    /* Save the old trap frame and set the new */
-    TrapFrame->TrapFrame = (ULONG64)Thread->TrapFrame;
-    Thread->TrapFrame = TrapFrame;
-
     /* We don't have an exception frame yet */
     TrapFrame->ExceptionFrame = 0;