From 410cf69b0c1875ff6940132470d394bff1ffd60b Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Tue, 14 Nov 2023 16:45:18 +0200
Subject: [PATCH] [NTOS:KE/x64] Handle XState allocation on the stack

---
 ntoskrnl/ke/amd64/stubs.c   |  4 ++++
 ntoskrnl/ke/amd64/thrdini.c | 30 ++++++++++++++++++++++--------
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/ntoskrnl/ke/amd64/stubs.c b/ntoskrnl/ke/amd64/stubs.c
index 7326f7a2cd0..9b87b588df0 100644
--- a/ntoskrnl/ke/amd64/stubs.c
+++ b/ntoskrnl/ke/amd64/stubs.c
@@ -76,6 +76,10 @@ KiSwitchKernelStack(PVOID StackBase, PVOID StackLimit)
     CurrentThread->InitialStack = Add2Ptr(CurrentThread->InitialStack,
                                           StackOffset);
 
+    /* Switch StateSaveArea */
+    CurrentThread->StateSaveArea = Add2Ptr(CurrentThread->StateSaveArea,
+                                           StackOffset);
+
     /* Set the new stack limits */
     CurrentThread->StackBase = StackBase;
     CurrentThread->StackLimit = (ULONG_PTR)StackLimit;
diff --git a/ntoskrnl/ke/amd64/thrdini.c b/ntoskrnl/ke/amd64/thrdini.c
index 3185de21cd6..a1f9d614508 100644
--- a/ntoskrnl/ke/amd64/thrdini.c
+++ b/ntoskrnl/ke/amd64/thrdini.c
@@ -43,12 +43,30 @@ KiInitializeContextThread(IN PKTHREAD Thread,
                            IN PVOID StartContext,
                            IN PCONTEXT Context)
 {
-    //PFX_SAVE_AREA FxSaveArea;
-    //PFXSAVE_FORMAT FxSaveFormat;
     PKSTART_FRAME StartFrame;
     PKSWITCH_FRAME CtxSwitchFrame;
     PKTRAP_FRAME TrapFrame;
     ULONG ContextFlags;
+    PVOID InitialStack;
+
+    /* Allocate space on the stack for the XSAVE area */
+    InitialStack = (PUCHAR)Thread->InitialStack - KeXStateLength;
+    InitialStack = ALIGN_DOWN_POINTER_BY(InitialStack, 64);
+    Thread->InitialStack = InitialStack;
+
+    /* Initialize the state save area */
+    Thread->StateSaveArea = InitialStack;
+    RtlZeroMemory(Thread->StateSaveArea, KeXStateLength);
+    Thread->StateSaveArea->MxCsr = INITIAL_MXCSR;
+
+    /* Special initialization for XSAVES */
+    if (KeFeatureBits & KF_XSAVES)
+    {
+        /* Set bit 63 in XCOMP_BV to mark the area as compacted.
+           XRSTORS requires this and will #GP otherwise. */
+        PXSAVE_AREA XSaveArea = (PXSAVE_AREA)Thread->StateSaveArea;
+        XSaveArea->Header.Reserved[0] = 0x8000000000000000ULL;
+    }
 
     /* Check if this is a With-Context Thread */
     if (Context)
@@ -66,10 +84,8 @@ KiInitializeContextThread(IN PKTHREAD Thread,
         /* Tell the thread it will run in User Mode */
         Thread->PreviousMode = UserMode;
 
-        // FIXME Setup the Fx Area
-
         /* Set the Thread's NPX State */
-        Thread->NpxState = 0xA;
+        Thread->NpxState = SharedUserData->XState.EnabledFeatures;
         Thread->Header.NpxIrql = PASSIVE_LEVEL;
 
         /* Make sure, we have control registers, disable debug registers */
@@ -123,10 +139,8 @@ KiInitializeContextThread(IN PKTHREAD Thread,
         /* Tell the thread it will run in Kernel Mode */
         Thread->PreviousMode = KernelMode;
 
-        // FIXME Setup the Fx Area
-
         /* No NPX State */
-        Thread->NpxState = 0xA;
+        Thread->NpxState = 0;
 
         /* This must never return! */
         StartFrame->Return = (ULONG64)KiInvalidSystemThreadStartupExit;