From c86c55ace74345d8ad91763b803579bd034fd6eb Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Fri, 9 Feb 2018 20:59:33 +0100
Subject: [PATCH] [NTOS:KE:X64] Change the logic of KeSwitchKernelStack and
 friends to be standards conforming

The previous version (like the x86 one) used a combination of C and asm code, called from C code to switch the stack. This is problematic, since there is no guarantee what assumptions C code makes about the stack (i.e. it can place any kind of stack pointers into registers or on the stack itself.) The new algorithm returns back to the systemcall entry point in asm, which then calls KiConvertToGuiThread, which is also asm and calls KeSwitchKernelStack ...
---
 ntoskrnl/ke/amd64/stubs.c | 35 +++++++++++++++++++++++++----------
 ntoskrnl/ke/amd64/trap.S  | 34 ++++++++++++++++++++++++++++++++--
 2 files changed, 57 insertions(+), 12 deletions(-)

diff --git a/ntoskrnl/ke/amd64/stubs.c b/ntoskrnl/ke/amd64/stubs.c
index 9e3624d792f..c23c02b0849 100644
--- a/ntoskrnl/ke/amd64/stubs.c
+++ b/ntoskrnl/ke/amd64/stubs.c
@@ -13,6 +13,10 @@
 #define NDEBUG
 #include <debug.h>
 
+ULONG ProcessCount;
+BOOLEAN CcPfEnablePrefetcher;
+SIZE_T KeXStateLength = sizeof(XSAVE_FORMAT);
+
 VOID
 KiRetireDpcListInDpcStack(
     PKPRCB Prcb,
@@ -96,14 +100,28 @@ KiSwitchKernelStackHelper(
     LONG_PTR StackOffset,
     PVOID OldStackBase);
 
+/*
+ * Kernel stack layout (example pointers):
+ * 0xFFFFFC0F'2D008000 KTHREAD::StackBase
+ *    [XSAVE_AREA size == KeXStateLength = 0x440]
+ * 0xFFFFFC0F'2D007BC0 KTHREAD::StateSaveArea _XSAVE_FORMAT
+ * 0xFFFFFC0F'2D007B90 KTHREAD::InitialStack
+ *    [0x190 bytes KTRAP_FRAME]
+ * 0xFFFFFC0F'2D007A00 KTHREAD::TrapFrame
+ *    [KSTART_FRAME] or ...
+ *    [KSWITCH_FRAME]
+ * 0xFFFFFC0F'2D007230 KTHREAD::KernelStack
+ */
+
 PVOID
 NTAPI
-KeSwitchKernelStack(PVOID StackBase, PVOID StackLimit)
+KiSwitchKernelStack(PVOID StackBase, PVOID StackLimit)
 {
     PKTHREAD CurrentThread;
     PVOID OldStackBase;
     LONG_PTR StackOffset;
     SIZE_T StackSize;
+    PKIPCR Pcr;
 
     /* Get the current thread */
     CurrentThread = KeGetCurrentThread();
@@ -139,20 +157,17 @@ KeSwitchKernelStack(PVOID StackBase, PVOID StackLimit)
     CurrentThread->StackLimit = (ULONG_PTR)StackLimit;
     CurrentThread->LargeStack = TRUE;
 
-    /* Adjust the PCR fields */
-    __addgsqword(FIELD_OFFSET(KPCR, NtTib.StackBase), StackOffset);
-    __addgsqword(FIELD_OFFSET(KIPCR, Prcb.RspBase), StackOffset);
+    /* Adjust RspBase in the PCR */
+    Pcr = (PKIPCR)KeGetPcr();
+    Pcr->Prcb.RspBase += StackOffset;
 
-    /* Finally switch RSP to the new stack.
-       We pass OldStackBase to make sure it is not lost. */
-    OldStackBase = KiSwitchKernelStackHelper(StackOffset, OldStackBase);
-
-    /* Reenable interrupts */
-    _enable();
+    /* Adjust Rsp0 in the TSS */
+    Pcr->TssBase->Rsp0 += StackOffset;
 
     return OldStackBase;
 }
 
+
 NTSTATUS
 NTAPI
 KeUserModeCallback(IN ULONG RoutineIndex,
diff --git a/ntoskrnl/ke/amd64/trap.S b/ntoskrnl/ke/amd64/trap.S
index a5ba4998745..d361acb1fb1 100644
--- a/ntoskrnl/ke/amd64/trap.S
+++ b/ntoskrnl/ke/amd64/trap.S
@@ -22,7 +22,6 @@ EXTERN KiDeliverApc:PROC
 EXTERN KiDpcInterruptHandler:PROC
 EXTERN PsConvertToGuiThread:PROC
 EXTERN MmCreateKernelStack:PROC
-EXTERN KeSwitchKernelStack:PROC
 EXTERN MmDeleteKernelStack:PROC
 
 #ifdef _WINKD_
@@ -774,7 +773,7 @@ PUBLIC KiSystemCallEntry64
 
 .PROC KiSystemCall64Again
 
-   /* Old stack pointer is in rcx, lie and say we saved it in rbp */
+    /* Old stack pointer is in rcx, lie and say we saved it in rbp */
     .setframe rbp, 0
     .endprolog
 
@@ -1023,6 +1022,37 @@ KiSwitchKernelStackHelper:
     mov rax, rdx
     ret
 
+EXTERN KiSwitchKernelStack:PROC
+
+PUBLIC KeSwitchKernelStack
+FUNC KeSwitchKernelStack
+
+    sub rsp, 40
+    .allocstack 40
+
+    ; Save rcx
+    mov [rsp], rcx
+    .savereg rcx, 0
+    .endprolog
+
+    ; Call the C handler, which returns the old stack in rax
+   call KiSwitchKernelStack
+
+    ; Restore rcx (StackBase)
+    mov rcx, [rsp]
+
+    ; Switch to new stack: RSP += (StackBase - OldStackBase)
+    sub rcx, rax
+    add rsp, rcx
+
+    ; Deallocate the home frame
+    add rsp, 40
+    ret
+
+ENDFUNC
+
+
+
 
 #ifdef _MSC_VER
 #undef lgdt